fix: the result will be incorrect when call gin.Context.Writer.Status…

…() in our custom gin middleware used after this middleware (#46)

Co-authored-by: zhangyi <zhangyi905@guance.com>

timeout_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"net/http"
 	"net/http/httptest"
+	"strconv"
 	"sync"
 	"testing"
 	"time"
@@ -102,6 +103,37 @@ func TestPanic(t *testing.T) {
 	assert.Equal(t, "", w.Body.String())
 }
 
+func TestWriter_Status(t *testing.T) {
+	r := gin.New()
+
+	r.Use(New(
+		WithTimeout(1*time.Second),
+		WithHandler(func(c *gin.Context) {
+			c.Next()
+		}),
+		WithResponse(testResponse),
+	))
+
+	r.Use(func(c *gin.Context) {
+		c.Next()
+		statusInMW := c.Writer.Status()
+		c.Request.Header.Set("X-Status-Code-MW-Set", strconv.Itoa(statusInMW))
+		t.Logf("[%s] %s %s %d\n", time.Now().Format(time.RFC3339), c.Request.Method, c.Request.URL, statusInMW)
+	})
+
+	r.GET("/test", func(c *gin.Context) {
+		c.Writer.WriteHeader(http.StatusInternalServerError)
+	})
+
+	w := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodGet, "/test", nil)
+
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, w.Result().StatusCode, http.StatusInternalServerError)
+	assert.Equal(t, strconv.Itoa(http.StatusInternalServerError), req.Header.Get("X-Status-Code-MW-Set"))
+}
+
 func TestDeadlineExceeded(t *testing.T) {
 	var wg sync.WaitGroup
 	wg.Add(1)

writer.go

@@ -55,6 +55,16 @@ func (w *Writer) writeHeader(code int) {
 	w.code = code
 }
 
+// Status we must override Status func here,
+// or the http status code returned by gin.Context.Writer.Status()
+// will always be 200 in other custom gin middlewares.
+func (w *Writer) Status() int {
+	if w.code == 0 || w.timeout {
+		return w.ResponseWriter.Status()
+	}
+	return w.code
+}
+
 // Header will get response headers
 func (w *Writer) Header() http.Header {
 	return w.headers