add MaxBodyBytes to EndpointMeta for enforcing req payload size

Author: bgentry

apiendpoint/api_endpoint.go

@@ -11,6 +11,7 @@ import (
 	"io"
 	"log/slog"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/jackc/pgerrcode"
@@ -68,6 +69,11 @@ type EndpointExecuteInterface[TReq any, TResp any] interface {
 
 // EndpointMeta is metadata about an API endpoint.
 type EndpointMeta struct {
+	// MaxBodyBytes is the maximum number of bytes that can be read from the
+	// request body. If the request body exceeds this number, a 413 error will
+	// be returned.
+	MaxBodyBytes int64
+
 	// Pattern is the API endpoint's HTTP method and path where it should be
 	// mounted, which is passed to http.ServeMux by Mount. It should start with
 	// a verb like `GET` or `POST`, and may contain Go 1.22 path variables like
@@ -136,8 +142,15 @@ func executeAPIEndpoint[TReq any, TResp any](w http.ResponseWriter, r *http.Requ
 	err := func() error {
 		var req TReq
 		if r.Method != http.MethodGet {
+			if meta.MaxBodyBytes > 0 {
+				r.Body = http.MaxBytesReader(w, r.Body, meta.MaxBodyBytes)
+			}
+
 			reqData, err := io.ReadAll(r.Body)
 			if err != nil {
+				if strings.Contains(err.Error(), "request body too large") {
+					return apierror.NewRequestEntityTooLarge("Request entity too large")
+				}
 				return fmt.Errorf("error reading request body: %w", err)
 			}
 

apiendpoint/api_endpoint_test.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"log/slog"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -25,6 +26,7 @@ func TestMountAndServe(t *testing.T) {
 	ctx := context.Background()
 
 	type testBundle struct {
+		logger   *slog.Logger
 		recorder *httptest.ResponseRecorder
 	}
 
@@ -41,6 +43,7 @@ func TestMountAndServe(t *testing.T) {
 		Mount(mux, &postEndpoint{}, opts)
 
 		return mux, &testBundle{
+			logger:   logger,
 			recorder: httptest.NewRecorder(),
 		}
 	}
@@ -68,6 +71,28 @@ func TestMountAndServe(t *testing.T) {
 		requireStatusAndJSONResponse(t, http.StatusOK, &postResponse{Message: "Hello."}, bundle.recorder)
 	})
 
+	t.Run("MaxBodyBytes", func(t *testing.T) {
+		t.Parallel()
+
+		_, bundle := setup(t)
+
+		payload := mustMarshalJSON(t, &postRequest{Message: "Hello."})
+
+		mux := http.NewServeMux()
+		endpoint := &postEndpoint{MaxBodyBytes: int64(len(payload))}
+		Mount(mux, endpoint, &MountOpts{Logger: bundle.logger})
+
+		req := httptest.NewRequest(http.MethodPost, "/api/post-endpoint", bytes.NewBuffer(payload))
+		mux.ServeHTTP(bundle.recorder, req)
+		requireStatusAndResponse(t, http.StatusCreated, `{"message":"Hello."}`, bundle.recorder)
+
+		bundle.recorder = httptest.NewRecorder()
+		payload = mustMarshalJSON(t, &postRequest{Message: "Hello!!"}) // one longer than previous payload
+		req = httptest.NewRequest(http.MethodPost, "/api/post-endpoint", bytes.NewBuffer(payload))
+		mux.ServeHTTP(bundle.recorder, req)
+		requireStatusAndJSONResponse(t, http.StatusRequestEntityTooLarge, &apierror.APIError{Message: "Request entity too large"}, bundle.recorder)
+	})
+
 	t.Run("MethodNotAllowed", func(t *testing.T) {
 		t.Parallel()
 
@@ -275,12 +300,14 @@ func (a *getEndpoint) Execute(_ context.Context, req *getRequest) (*getResponse,
 
 type postEndpoint struct {
 	Endpoint[postRequest, postResponse]
+	MaxBodyBytes int64
 }
 
-func (*postEndpoint) Meta() *EndpointMeta {
+func (a *postEndpoint) Meta() *EndpointMeta {
 	return &EndpointMeta{
-		Pattern:    "POST /api/post-endpoint",
-		StatusCode: http.StatusCreated,
+		MaxBodyBytes: a.MaxBodyBytes,
+		Pattern:      "POST /api/post-endpoint",
+		StatusCode:   http.StatusCreated,
 	}
 }
 

apierror/api_error.go

@@ -131,6 +131,23 @@ func NewNotFoundf(format string, a ...any) *NotFound {
 	return NewNotFound(fmt.Sprintf(format, a...))
 }
 
+//
+// RequestEntityTooLarge
+//
+
+type RequestEntityTooLarge struct { //nolint:errname
+	APIError
+}
+
+func NewRequestEntityTooLarge(message string) *RequestEntityTooLarge {
+	return &RequestEntityTooLarge{
+		APIError: APIError{
+			Message:    message,
+			StatusCode: http.StatusRequestEntityTooLarge,
+		},
+	}
+}
+
 //
 // ServiceUnavailable
 //