diff --git a/huma.go b/huma.go index a80bf2c2..8c9e5a47 100644 --- a/huma.go +++ b/huma.go @@ -14,6 +14,7 @@ import ( "errors" "fmt" "io" + "mime" "mime/multipart" "net" "net/http" @@ -1020,7 +1021,7 @@ func Register[I, O any](api API, op Operation, handler func(context.Context, *I) if rbt.isMultipart() { // Read form - form, err := readForm(ctx) + form, err := readForm(ctx, op.MaxBodyBytes) jsonUnmarshaler := func(data []byte, v any) error { return api.Unmarshal("application/json", data, v) } if err != nil { @@ -2043,13 +2044,36 @@ func processMultipartMsgBody(form *multipart.Form, op Operation, v reflect.Value return nil } -func readForm(ctx Context) (*multipart.Form, *ErrorDetail) { - form, err := ctx.GetMultipartForm() +func readForm(ctx Context, maxBytes int64) (*multipart.Form, error) { + if maxBytes <= 0 { + form, err := ctx.GetMultipartForm() + if err != nil { + return form, &ErrorDetail{ + Location: "body", + Message: "cannot read multipart form: " + err.Error(), + } + } + return form, nil + } + + _, params, err := mime.ParseMediaType(ctx.Header("Content-Type")) if err != nil { - return form, &ErrorDetail{ - Location: "body", - Message: "cannot read multipart form: " + err.Error(), + return nil, &ErrorDetail{Location: "body", Message: "cannot read multipart form: " + err.Error()} + } + + r := &io.LimitedReader{R: ctx.BodyReader(), N: maxBytes} + form, err := multipart.NewReader(r, params["boundary"]).ReadForm(8 << 10) + if err == nil { + _, err = io.Copy(io.Discard, r) + } + if r.N == 0 { + if form != nil { + _ = form.RemoveAll() } + return nil, Error413RequestEntityTooLarge(fmt.Sprintf("request body is too large limit=%d bytes", maxBytes)) + } + if err != nil { + return form, &ErrorDetail{Location: "body", Message: "cannot read multipart form: " + err.Error()} } return form, nil } diff --git a/huma_test.go b/huma_test.go index 5d35f4c1..134efaf4 100644 --- a/huma_test.go +++ b/huma_test.go @@ -2919,10 +2919,10 @@ Content-Type: text/plain }, Method: http.MethodGet, URL: "/transform", -Assert: func(t *testing.T, resp *httptest.ResponseRecorder) { - assert.Equal(t, http.StatusOK, resp.Code) - assert.JSONEq(t, `null`, resp.Body.String()) -}, + Assert: func(t *testing.T, resp *httptest.ResponseRecorder) { + assert.Equal(t, http.StatusOK, resp.Code) + assert.JSONEq(t, `null`, resp.Body.String()) + }, }, { Name: "schema-url-from-x-forwarded-host", @@ -3711,6 +3711,52 @@ func TestMultipartStructFieldRequiresJSONTag(t *testing.T) { }) } +func TestMultipartMaxBodyBytesEnforced(t *testing.T) { + const maxBodyBytes = 1024 + + mux, api := humatest.New(t, huma.DefaultConfig("Test", "1.0.0")) + + var handlerCalled bool + huma.Register(api, huma.Operation{ + Method: http.MethodPost, + Path: "/upload", + MaxBodyBytes: maxBodyBytes, + }, func(_ context.Context, _ *struct { + RawBody multipart.Form + }) (*struct{}, error) { + handlerCalled = true + return &struct{}{}, nil + }) + + var body bytes.Buffer + mw := multipart.NewWriter(&body) + part, err := mw.CreateFormField("large") + if err != nil { + t.Fatal(err) + } + // Much larger than MaxBodyBytes + if _, err := part.Write([]byte(strings.Repeat("a", maxBodyBytes+1))); err != nil { + t.Fatal(err) + } + if err := mw.Close(); err != nil { + t.Fatal(err) + } + mpbodylen := body.Len() + + req := httptest.NewRequest(http.MethodPost, "/upload", &body) + req.Header.Set("Content-Type", mw.FormDataContentType()) + res := httptest.NewRecorder() + mux.ServeHTTP(res, req) + + if handlerCalled { + t.Fatalf("handler was called even though multipart body (%d bytes) exceeds MaxBodyBytes (%d bytes)", mpbodylen, maxBodyBytes) + } + + if got, want := res.Code, http.StatusRequestEntityTooLarge; got != want { + t.Fatalf("invalid status code: got %d, want %d", got, want) + } +} + func TestOpenAPI(t *testing.T) { r, api := humatest.New(t, huma.DefaultConfig("Features Test API", "1.0.0"))