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
|
package server
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"bufio"
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"io"
|
||||||
"log/slog"
|
"log/slog"
|
||||||
"net"
|
"net"
|
||||||
"net/http"
|
"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.
|
// Shutdown idle connections on the shared transport.
|
||||||
func (g *Gateway) Shutdown(ctx context.Context) error {
|
func (g *Gateway) Shutdown(ctx context.Context) error {
|
||||||
g.transport.CloseIdleConnections()
|
g.transport.CloseIdleConnections()
|
||||||
|
|||||||
Reference in New Issue
Block a user