fix: restore original writer so unmatched routes still return 404

timeout.go

@@ -122,6 +122,11 @@ func New(opts ...Option) gin.HandlerFunc {
 			}
 			tw.FreeBuffer()
 			bufPool.Put(buffer)
+			// Restore the original writer so anything gin writes after this
+			// middleware returns - such as the default 404 body from
+			// serveError() when no route matched - reaches the real
+			// ResponseWriter instead of the freed buffer.
+			c.Writer = w
 
 		case <-timer.C:
 			tw.mu.Lock()
@@ -145,7 +150,12 @@ func New(opts ...Option) gin.HandlerFunc {
 			case <-panicChan:
 			}
 
-			// Goroutine is done. Safe to modify c.index now.
+			// Goroutine is done. Safe to modify c and c.index now. Restoring the
+			// writer matters here too: FreeBuffer left tw reporting Size() == -1,
+			// so middleware that inspects c.Writer after c.Next() - gin's own
+			// logger reading BodySize, for one - would read that instead of the
+			// timeout response actually written to w.
+			c.Writer = w
 			c.Abort()
 		}
 	}

timeout_test.go

@@ -296,3 +296,61 @@ func TestContextDeadlineSet(t *testing.T) {
 	assert.True(t, hasDeadline, "request context should have a deadline set by the middleware")
 	assert.Equal(t, http.StatusOK, w.Code)
 }
+
+func TestNoRouteWithUse(t *testing.T) {
+	r := gin.New()
+	r.Use(New(
+		WithTimeout(1 * time.Second),
+	))
+	r.GET("/hello", func(c *gin.Context) {
+		c.String(http.StatusOK, "world")
+	})
+
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequestWithContext(context.Background(), "GET", "/no/such/route", nil)
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusNotFound, w.Code)
+	assert.Equal(t, "404 page not found", w.Body.String())
+}
+
+func TestMethodNotAllowedWithUse(t *testing.T) {
+	r := gin.New()
+	r.HandleMethodNotAllowed = true
+	r.Use(New(
+		WithTimeout(1 * time.Second),
+	))
+	r.GET("/hello", func(c *gin.Context) {
+		c.String(http.StatusOK, "world")
+	})
+
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequestWithContext(context.Background(), "POST", "/hello", nil)
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusMethodNotAllowed, w.Code)
+	assert.Equal(t, "405 method not allowed", w.Body.String())
+}
+
+func TestOuterMiddlewareSeesResponseSize(t *testing.T) {
+	var size int
+	r := gin.New()
+	// Runs before the timeout middleware, so after c.Next() it reads whatever
+	// writer was left on the context - the same position as gin.Logger().
+	r.Use(func(c *gin.Context) {
+		c.Next()
+		size = c.Writer.Size()
+	})
+	r.Use(New(
+		WithTimeout(5 * time.Millisecond),
+	))
+	r.GET("/", emptySuccessResponse)
+
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequestWithContext(context.Background(), "GET", "/", nil)
+	r.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusRequestTimeout, w.Code)
+	assert.Equal(t, len(http.StatusText(http.StatusRequestTimeout)), size,
+		"outer middleware should see the size of the timeout response, not the freed buffer")
+}