diff --git a/middleware/body_limit.go b/middleware/body_limit.go index d13ad2c4e..e48aff61c 100644 --- a/middleware/body_limit.go +++ b/middleware/body_limit.go @@ -92,12 +92,18 @@ func BodyLimitWithConfig(config BodyLimitConfig) echo.MiddlewareFunc { } func (r *limitedReader) Read(b []byte) (n int, err error) { + if r.limit > 0 && r.read > r.limit { + return 0, echo.ErrStatusRequestEntityTooLarge + } + n, err = r.reader.Read(b) r.read += int64(n) - if r.read > r.limit { + + if r.limit > 0 && r.read > r.limit { return n, echo.ErrStatusRequestEntityTooLarge } - return + + return n, err } func (r *limitedReader) Close() error { diff --git a/middleware/body_limit_test.go b/middleware/body_limit_test.go index d14c2b649..8ad9004c1 100644 --- a/middleware/body_limit_test.go +++ b/middleware/body_limit_test.go @@ -171,3 +171,57 @@ func TestBodyLimit_panicOnInvalidLimit(t *testing.T) { func() { BodyLimit("") }, ) } + +func TestBodyLimit_Middleware_BodyRestoration(t *testing.T) { + e := echo.New() + + e.Use(BodyLimit("1KB")) + + e.Use(func(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + bodyBytes, err := io.ReadAll(c.Request().Body) + if err != nil { + return err + } + + c.Request().Body.Close() + + c.Request().Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) + + return next(c) + } + }) + + e.POST("/", func(c echo.Context) error { + type Payload struct { + Message string `json:"message"` + } + p := new(Payload) + if err := c.Bind(p); err != nil { + return err + } + return c.String(http.StatusOK, p.Message) + }) + + t.Run("valid request under 1KB binds successfully", func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader([]byte(`{"message": "hello"}`))) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + + e.ServeHTTP(rec, req) + + assert.Equal(t, http.StatusOK, rec.Code) + assert.Equal(t, "hello", rec.Body.String()) + }) + + t.Run("request exceeding 1KB returns 413 at middleware read phase", func(t *testing.T) { + largePayload := `{"message": "` + string(bytes.Repeat([]byte("A"), 2000)) + `"}` + req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader([]byte(largePayload))) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + + e.ServeHTTP(rec, req) + + assert.Equal(t, http.StatusRequestEntityTooLarge, rec.Code) + }) +}