refactor: simplify the code (#2763)

* refactor: simplify the code

* fix: fix data race

* refactor: simplify the code

* refactor: simplify the code

rest/handler/loghandler.go

@@ -22,6 +22,7 @@ import (
 	"github.com/zeromicro/go-zero/core/utils"
 	"github.com/zeromicro/go-zero/rest/httpx"
 	"github.com/zeromicro/go-zero/rest/internal"
+	"github.com/zeromicro/go-zero/rest/internal/response"
 )
 
 const (
@@ -31,66 +32,30 @@ const (
 
 var slowThreshold = syncx.ForAtomicDuration(defaultSlowThreshold)
 
-type loggedResponseWriter struct {
-	w    http.ResponseWriter
-	r    *http.Request
-	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) {
-	if hijacked, ok := w.w.(http.Hijacker); ok {
-		return hijacked.Hijack()
-	}
-
-	return nil, nil, errors.New("server doesn't support hijacking")
-}
-
-func (w *loggedResponseWriter) Write(bytes []byte) (int, error) {
-	return w.w.Write(bytes)
-}
-
-func (w *loggedResponseWriter) WriteHeader(code int) {
-	w.w.WriteHeader(code)
-	w.code = code
-}
-
 // 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) {
 		timer := utils.NewElapsedTimer()
 		logs := new(internal.LogCollector)
-		lrw := loggedResponseWriter{
-			w:    w,
-			r:    r,
-			code: http.StatusOK,
+		lrw := response.WithCodeResponseWriter{
+			Writer: w,
+			Code:   http.StatusOK,
 		}
 
 		var dup io.ReadCloser
 		r.Body, dup = iox.DupReadCloser(r.Body)
 		next.ServeHTTP(&lrw, r.WithContext(context.WithValue(r.Context(), internal.LogContext, logs)))
 		r.Body = dup
-		logBrief(r, lrw.code, timer, logs)
+		logBrief(r, lrw.Code, timer, logs)
 	})
 }
 
 type detailLoggedResponseWriter struct {
-	writer *loggedResponseWriter
+	writer *response.WithCodeResponseWriter
 	buf    *bytes.Buffer
 }
 
-func newDetailLoggedResponseWriter(writer *loggedResponseWriter, buf *bytes.Buffer) *detailLoggedResponseWriter {
+func newDetailLoggedResponseWriter(writer *response.WithCodeResponseWriter, buf *bytes.Buffer) *detailLoggedResponseWriter {
 	return &detailLoggedResponseWriter{
 		writer: writer,
 		buf:    buf,
@@ -108,7 +73,7 @@ func (w *detailLoggedResponseWriter) Header() http.Header {
 // Hijack implements the http.Hijacker interface.
 // This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
 func (w *detailLoggedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-	if hijacked, ok := w.writer.w.(http.Hijacker); ok {
+	if hijacked, ok := w.writer.Writer.(http.Hijacker); ok {
 		return hijacked.Hijack()
 	}
 
@@ -129,10 +94,9 @@ func DetailedLogHandler(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		timer := utils.NewElapsedTimer()
 		var buf bytes.Buffer
-		lrw := newDetailLoggedResponseWriter(&loggedResponseWriter{
-			w:    w,
-			r:    r,
-			code: http.StatusOK,
+		lrw := newDetailLoggedResponseWriter(&response.WithCodeResponseWriter{
+			Writer: w,
+			Code:   http.StatusOK,
 		}, &buf)
 
 		var dup io.ReadCloser
@@ -203,7 +167,7 @@ func logDetails(r *http.Request, response *detailLoggedResponseWriter, timer *ut
 	logs *internal.LogCollector) {
 	var buf bytes.Buffer
 	duration := timer.Duration()
-	code := response.writer.code
+	code := response.writer.Code
 	logger := logx.WithContext(r.Context())
 	buf.WriteString(fmt.Sprintf("[HTTP] %s - %d - %s - %s\n=> %s\n",
 		r.Method, code, r.RemoteAddr, timex.ReprOfDuration(duration), dumpRequest(r)))

rest/handler/loghandler_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/rest/internal"
+	"github.com/zeromicro/go-zero/rest/internal/response"
 )
 
 func init() {
@@ -54,7 +55,7 @@ func TestLogHandlerVeryLong(t *testing.T) {
 	req := httptest.NewRequest(http.MethodPost, "http://localhost", &buf)
 	handler := LogHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		r.Context().Value(internal.LogContext).(*internal.LogCollector).Append("anything")
-		io.Copy(io.Discard, r.Body)
+		_, _ = io.Copy(io.Discard, r.Body)
 		w.Header().Set("X-Test", "test")
 		w.WriteHeader(http.StatusServiceUnavailable)
 		_, err := w.Write([]byte("content"))
@@ -89,45 +90,26 @@ func TestLogHandlerSlow(t *testing.T) {
 		assert.Equal(t, http.StatusOK, resp.Code)
 	}
 }
-
-func TestLogHandler_Hijack(t *testing.T) {
-	resp := httptest.NewRecorder()
-	writer := &loggedResponseWriter{
-		w: resp,
-	}
-	assert.NotPanics(t, func() {
-		writer.Hijack()
-	})
-
-	writer = &loggedResponseWriter{
-		w: mockedHijackable{resp},
-	}
-	assert.NotPanics(t, func() {
-		writer.Hijack()
-	})
-}
-
 func TestDetailedLogHandler_Hijack(t *testing.T) {
 	resp := httptest.NewRecorder()
 	writer := &detailLoggedResponseWriter{
-		writer: &loggedResponseWriter{
-			w: resp,
+		writer: &response.WithCodeResponseWriter{
+			Writer: resp,
 		},
 	}
 	assert.NotPanics(t, func() {
-		writer.Hijack()
+		_, _, _ = writer.Hijack()
 	})
 
 	writer = &detailLoggedResponseWriter{
-		writer: &loggedResponseWriter{
-			w: mockedHijackable{resp},
+		writer: &response.WithCodeResponseWriter{
+			Writer: mockedHijackable{resp},
 		},
 	}
 	assert.NotPanics(t, func() {
-		writer.Hijack()
+		_, _, _ = writer.Hijack()
 	})
 }
-
 func TestSetSlowThreshold(t *testing.T) {
 	assert.Equal(t, defaultSlowThreshold, slowThreshold.Load())
 	SetSlowThreshold(time.Second)