diff --git a/rest/handler/authhandler.go b/rest/handler/authhandler.go index 9cb691aa..e6ccf695 100644 --- a/rest/handler/authhandler.go +++ b/rest/handler/authhandler.go @@ -1,8 +1,10 @@ package handler import ( + "bufio" "context" "errors" + "net" "net/http" "net/http/httputil" @@ -138,6 +140,12 @@ func (grw *guardedResponseWriter) Header() http.Header { return grw.writer.Header() } +// Hijack implements the http.Hijacker interface. +// This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it. +func (grw *guardedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return grw.writer.(http.Hijacker).Hijack() +} + func (grw *guardedResponseWriter) Write(body []byte) (int, error) { return grw.writer.Write(body) } diff --git a/rest/handler/cryptionhandler.go b/rest/handler/cryptionhandler.go index beb8700a..caae3f9d 100644 --- a/rest/handler/cryptionhandler.go +++ b/rest/handler/cryptionhandler.go @@ -1,11 +1,13 @@ package handler import ( + "bufio" "bytes" "encoding/base64" "errors" "io" "io/ioutil" + "net" "net/http" "github.com/tal-tech/go-zero/core/codec" @@ -94,6 +96,12 @@ func (w *cryptionResponseWriter) Header() http.Header { return w.ResponseWriter.Header() } +// Hijack implements the http.Hijacker interface. +// This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it. +func (w *cryptionResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return w.ResponseWriter.(http.Hijacker).Hijack() +} + func (w *cryptionResponseWriter) Write(p []byte) (int, error) { return w.buf.Write(p) } diff --git a/rest/handler/loghandler.go b/rest/handler/loghandler.go index 3663e274..4a0800d5 100644 --- a/rest/handler/loghandler.go +++ b/rest/handler/loghandler.go @@ -1,10 +1,12 @@ package handler import ( + "bufio" "bytes" "context" "fmt" "io" + "net" "net/http" "net/http/httputil" "time" @@ -25,10 +27,22 @@ type loggedResponseWriter struct { code int } +func (w *loggedResponseWriter) Flush() { + if flusher, ok := w.w.(http.Flusher); ok { + flusher.Flush() + } +} + func (w *loggedResponseWriter) Header() http.Header { return w.w.Header() } +// Hijack implements the http.Hijacker interface. +// This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it. +func (w *loggedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return w.w.(http.Hijacker).Hijack() +} + func (w *loggedResponseWriter) Write(bytes []byte) (int, error) { return w.w.Write(bytes) } @@ -38,12 +52,6 @@ func (w *loggedResponseWriter) WriteHeader(code int) { w.code = code } -func (w *loggedResponseWriter) Flush() { - if flusher, ok := w.w.(http.Flusher); ok { - flusher.Flush() - } -} - // LogHandler returns a middleware that logs http request and response. func LogHandler(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/rest/internal/security/withcoderesponsewriter.go b/rest/internal/security/withcoderesponsewriter.go index 20e77d1e..2496e301 100644 --- a/rest/internal/security/withcoderesponsewriter.go +++ b/rest/internal/security/withcoderesponsewriter.go @@ -1,6 +1,10 @@ package security -import "net/http" +import ( + "bufio" + "net" + "net/http" +) // A WithCodeResponseWriter is a helper to delay sealing a http.ResponseWriter on writing code. type WithCodeResponseWriter struct { @@ -20,6 +24,12 @@ func (w *WithCodeResponseWriter) Header() http.Header { return w.Writer.Header() } +// Hijack implements the http.Hijacker interface. +// This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it. +func (w *WithCodeResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return w.Writer.(http.Hijacker).Hijack() +} + // Write writes bytes into w. func (w *WithCodeResponseWriter) Write(bytes []byte) (int, error) { return w.Writer.Write(bytes)