diff --git a/timeout.go b/timeout.go index 6cdfc10..383c2fe 100644 --- a/timeout.go +++ b/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() } } diff --git a/timeout_test.go b/timeout_test.go index c6f8cfc..aec379a 100644 --- a/timeout_test.go +++ b/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") +}