diff --git a/internal/server/gateway.go b/internal/server/gateway.go index 280b10f..8a00c44 100644 --- a/internal/server/gateway.go +++ b/internal/server/gateway.go @@ -1,9 +1,11 @@ package server import ( + "bufio" "context" "encoding/json" "fmt" + "io" "log/slog" "net" "net/http" @@ -403,6 +405,27 @@ func (s *statusWriter) Flush() { } } +// Hijack preserves WebSocket upgrade support through middleware wrappers. +func (s *statusWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + if h, ok := s.ResponseWriter.(http.Hijacker); ok { + return h.Hijack() + } + return nil, nil, http.ErrNotSupported +} + +// ReadFrom keeps io.Copy fast-path behavior when the underlying writer supports it. +func (s *statusWriter) ReadFrom(r io.Reader) (int64, error) { + if rf, ok := s.ResponseWriter.(io.ReaderFrom); ok { + return rf.ReadFrom(r) + } + return io.Copy(s.ResponseWriter, r) +} + +// Unwrap lets net/http.ResponseController reach the original writer. +func (s *statusWriter) Unwrap() http.ResponseWriter { + return s.ResponseWriter +} + // Shutdown idle connections on the shared transport. func (g *Gateway) Shutdown(ctx context.Context) error { g.transport.CloseIdleConnections()