refactored body limit middleware to prevent unnecessery reads if already beyound the limit

CHANGELOG.md

@@ -3,12 +3,14 @@
 - Minor UI improvements (updated dark primary btn contrast).
 
 - Write the status header for JSON responses only if the fields picker succeed or has acceptable fallback.
-    _This is to allow custom response status code in reaction to failed json writes._
+    _This is to allow custom response status code for failed json writes._
 
 - Fixed collection index validator to allow expressions with parenthesis in the optional `WHERE` clause.
 
 - Clamped arccosine to [-1,1] in the Harvesine formula for the `geoDistance()` filter function to workaround edge case related to float rounding errors for some coordinates.
 
+- Prevent unnecessery body chunk read if we already known that we are beyond the allowed limit.
+
 - Bumped `golang.org/x/*` dependencies to silence security scanners ([#7829](https://github.com/pocketbase/pocketbase/discussions/7829)).
 
 

apis/middlewares_body_limit.go

@@ -84,32 +84,39 @@ func applyBodyLimit(e *core.RequestEvent, limitBytes int64) error {
 	}
 
 	// replace the request body
-	//
-	// note: we don't use sync.Pool since the size of the elements could vary too much
-	// and it might not be efficient (see https://github.com/golang/go/issues/23199)
-	e.Request.Body = &limitedReader{ReadCloser: e.Request.Body, limit: limitBytes}
+	e.Request.Body = newLimitedReader(e.Request.Body, limitBytes)
 
 	return nil
 }
 
+func newLimitedReader(body io.ReadCloser, limitBytes int64) *limitedReader {
+	return &limitedReader{
+		ReadCloser: body,
+		limit:      limitBytes,
+		remaining:  limitBytes,
+	}
+}
+
 type limitedReader struct {
 	io.ReadCloser
 	limit     int64
-	totalRead int64
+	remaining int64
 }
 
 func (r *limitedReader) Read(b []byte) (int, error) {
+	if r.remaining <= 0 {
+		return 0, ErrRequestEntityTooLarge
+	}
+
+	if int64(len(b)) > r.remaining {
+		b = b[0:r.remaining]
+	}
+
 	n, err := r.ReadCloser.Read(b)
-	if err != nil {
-		return n, err
-	}
 
-	r.totalRead += int64(n)
-	if r.totalRead > r.limit {
-		return n, ErrRequestEntityTooLarge
-	}
+	r.remaining -= int64(n)
 
-	return n, nil
+	return n, err
 }
 
 // explicit casts to ensure that the main struct methods will be invoked
@@ -120,6 +127,7 @@ func (r *limitedReader) Reread() {
 	rereader, ok := r.ReadCloser.(router.Rereader)
 	if ok {
 		rereader.Reread()
+		r.remaining = r.limit
 	}
 }
 

apis/middlewares_body_limit_test.go

@@ -1,9 +1,9 @@
 package apis_test
 
 import (
-	"bytes"
-	"fmt"
+	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/pocketbase/pocketbase/apis"
@@ -19,13 +19,26 @@ func TestBodyLimitMiddleware(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	pbRouter.POST("/a", func(e *core.RequestEvent) error {
-		return e.String(200, "a")
-	}) // default global BodyLimit check
 
-	pbRouter.POST("/b", func(e *core.RequestEvent) error {
-		return e.String(200, "b")
-	}).Bind(apis.BodyLimit(20))
+	testHandler := func(e *core.RequestEvent) error {
+		// read the body multiple times to ensure that the limited
+		// reader guards and rereads are invoked
+		var result any
+		if err := e.BindBody(&result); err != nil {
+			return err
+		}
+
+		if err := e.BindBody(&result); err != nil {
+			return err
+		}
+
+		return e.JSON(200, result)
+	}
+
+	const customLimit = 20
+
+	pbRouter.POST("/a", testHandler) // default global BodyLimit check
+	pbRouter.POST("/b", testHandler).Bind(apis.BodyLimit(customLimit))
 
 	mux, err := pbRouter.BuildMux()
 	if err != nil {
@@ -33,20 +46,82 @@ func TestBodyLimitMiddleware(t *testing.T) {
 	}
 
 	scenarios := []struct {
-		url            string
-		size           int64
-		expectedStatus int
+		name              string
+		url               string
+		body              string
+		lazyContentLength bool
+		expectedStatus    int
 	}{
-		{"/a", 21, 200},
-		{"/a", apis.DefaultMaxBodySize + 1, 413},
-		{"/b", 20, 200},
-		{"/b", 21, 413},
+		{
+			"(eager content-length check) with body = default limit",
+			"/a",
+			`"` + strings.Repeat("a", int(apis.DefaultMaxBodySize-2)) + `"`,
+			false,
+			http.StatusOK,
+		},
+		{
+			"(eager content-length check) with body > default limit",
+			"/a",
+			`"` + strings.Repeat("a", int(apis.DefaultMaxBodySize)) + `"`,
+			false,
+			http.StatusRequestEntityTooLarge,
+		},
+		{
+			"(lazy content-length check) with body = default limit",
+			"/a",
+			`"` + strings.Repeat("a", int(apis.DefaultMaxBodySize-2)) + `"`,
+			true,
+			http.StatusOK,
+		},
+		{
+			"(lazy content-length check) with body > default limit",
+			"/a",
+			`"` + strings.Repeat("a", int(apis.DefaultMaxBodySize)) + `"`,
+			true,
+			http.StatusRequestEntityTooLarge,
+		},
+		// ---
+		{
+			"(eager content-length check) with body = custom limit",
+			"/b",
+			`"` + strings.Repeat("a", customLimit-2) + `"`,
+			false,
+			http.StatusOK,
+		},
+		{
+			"(eager content-length check) with body > custom limit",
+			"/b",
+			`"` + strings.Repeat("a", customLimit) + `"`,
+			false,
+			http.StatusRequestEntityTooLarge,
+		},
+		{
+			"(lazy content-length check) with body = custom limit",
+			"/b",
+			`"` + strings.Repeat("a", customLimit-2) + `"`,
+			true,
+			http.StatusOK,
+		},
+		{
+			"(lazy content-length check) with body > custom limit",
+			"/b",
+			`"` + strings.Repeat("a", customLimit) + `"`,
+			true,
+			http.StatusRequestEntityTooLarge,
+		},
 	}
 
 	for _, s := range scenarios {
-		t.Run(fmt.Sprintf("%s_%d", s.url, s.size), func(t *testing.T) {
+		t.Run(s.name, func(t *testing.T) {
 			rec := httptest.NewRecorder()
-			req := httptest.NewRequest("POST", s.url, bytes.NewReader(make([]byte, s.size)))
+
+			req := httptest.NewRequest("POST", s.url, strings.NewReader(s.body))
+			req.Header.Set("Content-Type", "application/json")
+
+			if s.lazyContentLength {
+				req.ContentLength = -1
+			}
+
 			mux.ServeHTTP(rec, req)
 
 			result := rec.Result()