Add WebSocket support methods to statusWriter in gateway.go
- Implemented Hijack method to preserve WebSocket upgrade support through middleware. - Added ReadFrom method to maintain fast-path behavior for io.Copy when the underlying writer supports it. - Introduced Unwrap method to allow access to the original ResponseWriter, enhancing flexibility in response handling.
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user