Fix bug: the result will be incorrect when call gin.Context.Writer.Status() in our custom gin middleware used after this middleware

timeout_test.go

@@ -4,6 +4,7 @@
  "context"
  "net/http"
  "net/http/httptest"
+ "strconv"
  "sync"
  "testing"
  "time"
@@ -102,6 +103,37 @@
  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