add http hijack methods (#555)

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)
 }

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)
 }

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) {

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)