Implement body size limit middleware with dynamic config reload
Add body_limit middleware that reads limit from ArcSwap<DynamicConfig> on each request, enabling runtime config changes without restart. Uses Content-Length header check for fast rejection and http_body_util::Limited for streaming body enforcement. Default limit: 100 MB (104,857,600 bytes). Returns 413 Payload Too Large when exceeded.
This commit is contained in:
51
src/proxy/body_limit.rs
Normal file
51
src/proxy/body_limit.rs
Normal file
@@ -0,0 +1,51 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use axum::body::Body;
|
||||
use axum::extract::State;
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::IntoResponse;
|
||||
use http_body_util::Limited;
|
||||
|
||||
use crate::config::DynamicConfig;
|
||||
|
||||
pub const DEFAULT_BODY_LIMIT_BYTES: u64 = 104_857_600;
|
||||
|
||||
pub async fn body_limit_middleware(
|
||||
State(config): State<Arc<ArcSwap<DynamicConfig>>>,
|
||||
request: axum::extract::Request,
|
||||
next: axum::middleware::Next,
|
||||
) -> axum::response::Response {
|
||||
let limit = config.load().body.limit_bytes;
|
||||
let limit = if limit == 0 {
|
||||
DEFAULT_BODY_LIMIT_BYTES
|
||||
} else {
|
||||
limit
|
||||
};
|
||||
|
||||
if let Some(content_length) = request.headers().get("content-length") {
|
||||
if let Ok(length_str) = content_length.to_str() {
|
||||
if let Ok(length) = length_str.parse::<u64>() {
|
||||
if length > limit {
|
||||
return (StatusCode::PAYLOAD_TOO_LARGE, "Payload Too Large").into_response();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let (parts, body) = request.into_parts();
|
||||
let limited_body = Limited::new(body, limit as usize);
|
||||
let request = axum::extract::Request::from_parts(parts, Body::new(limited_body));
|
||||
|
||||
next.run(request).await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn default_body_limit_is_100mb() {
|
||||
assert_eq!(DEFAULT_BODY_LIMIT_BYTES, 104_857_600);
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,20 @@
|
||||
pub mod body_limit;
|
||||
pub mod error;
|
||||
pub mod handler;
|
||||
pub mod headers;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
|
||||
use crate::config::DynamicConfig;
|
||||
|
||||
pub fn router_with_body_limit(
|
||||
router: axum::Router,
|
||||
config: Arc<ArcSwap<DynamicConfig>>,
|
||||
) -> axum::Router {
|
||||
router.layer(axum::middleware::from_fn_with_state(
|
||||
config,
|
||||
body_limit::body_limit_middleware,
|
||||
))
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user