From a720df594a48abf63b96828462a512aec85cca9e Mon Sep 17 00:00:00 2001 From: Kero Date: Wed, 25 Dec 2024 16:16:02 +0800 Subject: [PATCH] fix(host): limit body size using DefaultBodyLimit (#437) The original implementation of check_max_body_size relied on the size_hint provided by Axum. However, we've encountered instances where check_max_body_size incorrectly returns a BAD_REQUEST response due to inaccurate size_hint values. This commit changes to use `DefaultBodyLimit`, which is the suggestion by axum. See also - https://github.com/tokio-rs/axum/pull/1400 - https://github.com/tokio-rs/axum/pull/1397 --- host/src/server/api/mod.rs | 25 ++++--------------------- 1 file changed, 4 insertions(+), 21 deletions(-) diff --git a/host/src/server/api/mod.rs b/host/src/server/api/mod.rs index 45be92f15..cf6a71412 100644 --- a/host/src/server/api/mod.rs +++ b/host/src/server/api/mod.rs @@ -1,9 +1,6 @@ use axum::{ - body::HttpBody, - extract::Request, + extract::DefaultBodyLimit, http::{header, HeaderName, Method, StatusCode, Uri}, - middleware::{self, Next}, - response::Response, Router, }; use tower::ServiceBuilder; @@ -20,6 +17,8 @@ pub mod v1; pub mod v2; pub mod v3; +pub const MAX_BODY_SIZE: usize = 1 << 20; + pub fn create_router(concurrency_limit: usize, jwt_secret: Option<&str>) -> Router { let cors = CorsLayer::new() .allow_methods([Method::GET, Method::POST, Method::OPTIONS]) @@ -46,7 +45,7 @@ pub fn create_router(concurrency_limit: usize, jwt_secret: Option<&str>) -> Rout .nest("/v3", v3_api.clone()) .merge(v3_api) .layer(middleware) - .layer(middleware::from_fn(check_max_body_size)) + .layer(DefaultBodyLimit::max(MAX_BODY_SIZE)) .layer(trace) .fallback(|uri: Uri| async move { (StatusCode::NOT_FOUND, format!("No handler found for {uri}")) @@ -63,19 +62,3 @@ pub fn create_router(concurrency_limit: usize, jwt_secret: Option<&str>) -> Rout pub fn create_docs() -> utoipa::openapi::OpenApi { v3::create_docs() } - -async fn check_max_body_size(req: Request, next: Next) -> Response { - const MAX_BODY_SIZE: u64 = 1 << 20; - let response_content_length = match req.body().size_hint().upper() { - Some(v) => v, - None => MAX_BODY_SIZE + 1, - }; - - if response_content_length > MAX_BODY_SIZE { - let mut resp = Response::new(axum::body::Body::from("request too large")); - *resp.status_mut() = StatusCode::BAD_REQUEST; - return resp; - } - - next.run(req).await -}