oxidecomputer · sunshowers · Mar 21, 2023 · Mar 21, 2023 · Mar 21, 2023 · davepacheco
diff --git a/dropshot/src/api_description.rs b/dropshot/src/api_description.rs
@@ -31,13 +31,15 @@ use std::collections::HashSet;
 /// provided explicitly--as well as parameters and a description which can be
 /// inferred from function parameter types and doc comments (respectively).
 #[derive(Debug)]
+#[non_exhaustive]
 pub struct ApiEndpoint<Context: ServerContext> {
     pub operation_id: String,
     pub handler: Box<dyn RouteHandler<Context>>,
     pub method: Method,
     pub path: String,
     pub parameters: Vec<ApiEndpointParameter>,
     pub body_content_type: ApiEndpointBodyContentType,
+    pub request_body_max_bytes: Option<usize>,
     pub response: ApiEndpointResponse,
     pub summary: Option<String>,
     pub description: Option<String>,
@@ -72,6 +74,7 @@ impl<'a, Context: ServerContext> ApiEndpoint<Context> {
             path: path.to_string(),
             parameters: func_parameters.parameters,
             body_content_type,
+            request_body_max_bytes: None,
             response,
             summary: None,
             description: None,
@@ -92,6 +95,11 @@ impl<'a, Context: ServerContext> ApiEndpoint<Context> {
         self
     }
 
+    pub fn request_body_max_bytes(mut self, max_bytes: usize) -> Self {
+        self.request_body_max_bytes = Some(max_bytes);
+        self
+    }
+
     pub fn tag<T: ToString>(mut self, tag: T) -> Self {
         self.tags.push(tag.to_string());
         self

diff --git a/dropshot/src/extractor/body.rs b/dropshot/src/extractor/body.rs
@@ -57,9 +57,8 @@ async fn http_request_load_body<Context: ServerContext, BodyType>(
 where
     BodyType: JsonSchema + DeserializeOwned + Send + Sync,
 {
-    let server = &rqctx.server;
     let (parts, body) = request.into_parts();
-    let body = StreamingBody::new(body, server.config.request_body_max_bytes)
+    let body = StreamingBody::new(body, rqctx.request_body_max_bytes)
         .into_bytes_mut()
         .await?;
 
@@ -191,12 +190,10 @@ impl ExclusiveExtractor for UntypedBody {
         rqctx: &RequestContext<Context>,
         request: hyper::Request<hyper::Body>,
     ) -> Result<UntypedBody, HttpError> {
-        let server = &rqctx.server;
         let body = request.into_body();
-        let body_bytes =
-            StreamingBody::new(body, server.config.request_body_max_bytes)
-                .into_bytes_mut()
-                .await?;
+        let body_bytes = StreamingBody::new(body, rqctx.request_body_max_bytes)
+            .into_bytes_mut()
+            .await?;
         Ok(UntypedBody { content: body_bytes.freeze() })
     }
 
@@ -396,11 +393,9 @@ impl ExclusiveExtractor for StreamingBody {
         rqctx: &RequestContext<Context>,
         request: hyper::Request<hyper::Body>,
     ) -> Result<Self, HttpError> {
-        let server = &rqctx.server;
-
         Ok(Self {
             body: request.into_body(),
-            cap: server.config.request_body_max_bytes,
+            cap: rqctx.request_body_max_bytes,
         })
     }
 

diff --git a/dropshot/src/handler.rs b/dropshot/src/handler.rs
@@ -72,13 +72,18 @@ pub type HttpHandlerResult = Result<Response<Body>, HttpError>;
 
 /// Handle for various interfaces useful during request processing.
 #[derive(Debug)]
+#[non_exhaustive]
 pub struct RequestContext<Context: ServerContext> {
     /// shared server state
     pub server: Arc<DropshotState<Context>>,
     /// HTTP request routing variables
     pub path_variables: VariableSet,
     /// expected request body mime type
     pub body_content_type: ApiEndpointBodyContentType,
+    /// Maximum request body size: typically the same as
+    /// [`server.config.request_body_max_bytes`], but can be overridden for an
+    /// individual request
+    pub request_body_max_bytes: usize,
     /// unique id assigned to this request
     pub request_id: String,
     /// logger for this specific request

diff --git a/dropshot/src/lib.rs b/dropshot/src/lib.rs
@@ -191,6 +191,9 @@
 //!
 //!     // Optional fields
 //!     tags = [ "all", "your", "OpenAPI", "tags" ],
+//!     // An optional limit for the request body size that overrides the
+//!     // default server config.
+//!     request_body_max_bytes = 1048576,
 //! }]
 //! ```
 //!

diff --git a/dropshot/src/router.rs b/dropshot/src/router.rs
@@ -207,10 +207,12 @@ impl MapValue for VariableValue {
 /// corresponding values in the actual path, and the expected body
 /// content type.
 #[derive(Debug)]
+#[non_exhaustive]
 pub struct RouterLookupResult<'a, Context: ServerContext> {
     pub handler: &'a dyn RouteHandler<Context>,
     pub variables: VariableSet,
     pub body_content_type: ApiEndpointBodyContentType,
+    pub request_body_max_bytes: Option<usize>,
 }
 
 impl<Context: ServerContext> HttpRouterNode<Context> {
@@ -483,6 +485,7 @@ impl<Context: ServerContext> HttpRouter<Context> {
                 handler: &*handler.handler,
                 variables,
                 body_content_type: handler.body_content_type.clone(),
+                request_body_max_bytes: handler.request_body_max_bytes,
             })
             .ok_or_else(|| {
                 HttpError::for_status(None, StatusCode::METHOD_NOT_ALLOWED)
@@ -766,6 +769,7 @@ mod test {
             parameters: vec![],
             body_content_type: ApiEndpointBodyContentType::default(),
             response: ApiEndpointResponse::default(),
+            request_body_max_bytes: None,
             summary: None,
             description: None,
             tags: vec![],

diff --git a/dropshot/src/server.rs b/dropshot/src/server.rs
@@ -789,6 +789,9 @@ async fn http_request_handle<C: ServerContext>(
         request: RequestInfo::from(&request),
         path_variables: lookup_result.variables,
         body_content_type: lookup_result.body_content_type,
+        request_body_max_bytes: lookup_result
+            .request_body_max_bytes
+            .unwrap_or(server.config.request_body_max_bytes),
         request_id: request_id.to_string(),
         log: request_log,
     };

diff --git a/dropshot/src/websocket.rs b/dropshot/src/websocket.rs
@@ -335,6 +335,7 @@ mod tests {
             request: RequestInfo::from(&request),
             path_variables: Default::default(),
             body_content_type: Default::default(),
+            request_body_max_bytes: 0,
             request_id: "".to_string(),
             log: log.clone(),
         };

diff --git a/dropshot/tests/fail/bad_endpoint20.rs b/dropshot/tests/fail/bad_endpoint20.rs
@@ -0,0 +1,22 @@
+// Copyright 2023 Oxide Computer Company
+
+#![allow(unused_imports)]
+
+use dropshot::endpoint;
+use dropshot::HttpError;
+use dropshot::HttpResponseOk;
+use dropshot::UntypedBody;
+
+#[endpoint {
+    method = GET,
+    path = "/test",
+    request_body_max_bytes = true,
+}]
+async fn bad_request_body_max_bytes(
+    _rqctx: RequestContext<()>,
+    param: UntypedBody,
+) -> Result<HttpResponseOk<()>, HttpError> {
+    Ok(HttpResponseOk(()))
+}
+
+fn main() {}
diff --git a/dropshot/tests/fail/bad_endpoint20.stderr b/dropshot/tests/fail/bad_endpoint20.stderr
@@ -0,0 +1,5 @@
+error: expected u64, but found `true`
+  --> tests/fail/bad_endpoint20.rs:13:30
+   |
+13 |     request_body_max_bytes = true,
+   |                              ^^^^
diff --git a/dropshot/tests/fail/bad_endpoint21.rs b/dropshot/tests/fail/bad_endpoint21.rs
@@ -0,0 +1,26 @@
+// Copyright 2023 Oxide Computer Company
+
+// This does not currently work, but we may want to support it in the future.
+
+#![allow(unused_imports)]
+
+use dropshot::endpoint;
+use dropshot::HttpError;
+use dropshot::HttpResponseOk;
+use dropshot::UntypedBody;
+
+const MAX_REQUEST_BYTES: u64 = 400_000;
+
+#[endpoint {
+    method = GET,
+    path = "/test",
+    request_body_max_bytes = MAX_REQUEST_BYTES,
+}]
+async fn bad_request_body_max_bytes(
+    _rqctx: RequestContext<()>,
+    param: UntypedBody,
+) -> Result<HttpResponseOk<()>, HttpError> {
+    Ok(HttpResponseOk(()))
+}
+
+fn main() {}
diff --git a/dropshot/tests/fail/bad_endpoint21.stderr b/dropshot/tests/fail/bad_endpoint21.stderr
@@ -0,0 +1,5 @@
+error: expected u64, but found `MAX_REQUEST_BYTES`
+  --> tests/fail/bad_endpoint21.rs:17:30
+   |
+17 |     request_body_max_bytes = MAX_REQUEST_BYTES,
+   |                              ^^^^^^^^^^^^^^^^^