From e493be8d3cf392ecedcec3a3902ed32f9fee54bb Mon Sep 17 00:00:00 2001
From: Sergio Benitez <sb@sergio.bz>
Date: Wed, 17 Jan 2024 12:15:01 -0800
Subject: [PATCH] Forward whole-form errors in 'FromForm' derive.

Resolves #2672.
---
 core/codegen/src/derive/from_form.rs |  6 ++++
 core/lib/src/form/parser.rs          |  9 ++++-
 core/lib/tests/multipart-limit.rs    | 54 ++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 1 deletion(-)
 create mode 100644 core/lib/tests/multipart-limit.rs

diff --git a/core/codegen/src/derive/from_form.rs b/core/codegen/src/derive/from_form.rs
index bf55d3bb..93ca2ec8 100644
--- a/core/codegen/src/derive/from_form.rs
+++ b/core/codegen/src/derive/from_form.rs
@@ -211,6 +211,12 @@ pub fn derive_from_form(input: proc_macro::TokenStream) -> TokenStream {
                 __fut.await;
             })))
         )
+        .inner_mapper(MapperBuild::new()
+            .with_output(|_, _| quote! {
+                fn push_error(__c: &mut Self::Context, __e: #_form::Error<'r>) {
+                    __c.__errors.push(__e);
+                }
+            }))
         .inner_mapper(MapperBuild::new()
             .with_output(|_, output| quote! {
                 fn finalize(mut __c: Self::Context) -> #_Result<Self, #_form::Errors<'r>> {
diff --git a/core/lib/src/form/parser.rs b/core/lib/src/form/parser.rs
index b53e2a50..4b5fe0a7 100644
--- a/core/lib/src/form/parser.rs
+++ b/core/lib/src/form/parser.rs
@@ -67,10 +67,17 @@ impl<'r, 'i> Parser<'r, 'i> {
             .get("data-form")
             .unwrap_or(Limits::DATA_FORM);
 
+        // Increase internal limit by 1 so multer can limit to `form_limit`.
+        let stream = data.open(form_limit + 1);
+        let constraints = multer::Constraints::new()
+            .size_limit(multer::SizeLimit::new()
+                .whole_stream(form_limit.into())
+                .per_field(form_limit.into()));
+
         Ok(Parser::Multipart(MultipartParser {
             request: req,
             buffer: local_cache_once!(req, SharedStack::new()),
-            source: Multipart::with_reader(data.open(form_limit), boundary),
+            source: Multipart::with_reader_with_constraints(stream, boundary, constraints),
             done: false,
         }))
     }
diff --git a/core/lib/tests/multipart-limit.rs b/core/lib/tests/multipart-limit.rs
new file mode 100644
index 00000000..31a13dc1
--- /dev/null
+++ b/core/lib/tests/multipart-limit.rs
@@ -0,0 +1,54 @@
+#[macro_use] extern crate rocket;
+
+use rocket::{Config, Build, Rocket};
+use rocket::{data::Limits, form::Form};
+use rocket::http::{ContentType, Status};
+use ubyte::{ToByteUnit, ByteUnit};
+
+#[derive(FromForm)]
+struct Data<'r> {
+    foo: Option<&'r str>,
+}
+
+#[rocket::post("/", data = "<form>")]
+fn form<'r>(form: Form<Data<'r>>) -> &'r str {
+    form.foo.unwrap_or("missing")
+}
+
+fn rocket_with_form_data_limit(limit: ByteUnit) -> Rocket<Build> {
+    rocket::custom(Config {
+        limits: Limits::default().limit("data-form", limit),
+        ..Config::debug_default()
+    }).mount("/", routes![form])
+}
+
+#[test]
+fn test_multipart_limit() {
+    use rocket::local::blocking::Client;
+
+    let body = &[
+        "--X-BOUNDARY",
+        r#"Content-Disposition: form-data; name="foo"; filename="foo.txt""#,
+        "Content-Type: text/plain",
+        "",
+        "hi",
+        "--X-BOUNDARY--",
+        "",
+    ].join("\r\n");
+
+    let client = Client::debug(rocket_with_form_data_limit(body.len().bytes())).unwrap();
+    let response = client.post("/")
+        .header("multipart/form-data; boundary=X-BOUNDARY".parse::<ContentType>().unwrap())
+        .body(body)
+        .dispatch();
+
+    assert_eq!(response.into_string().unwrap(), "hi");
+
+    let client = Client::debug(rocket_with_form_data_limit(body.len().bytes() - 1)).unwrap();
+    let response = client.post("/")
+        .header("multipart/form-data; boundary=X-BOUNDARY".parse::<ContentType>().unwrap())
+        .body(body)
+        .dispatch();
+
+    assert_eq!(response.status(), Status::PayloadTooLarge);
+}