use axum::{
body::Body,
extract::Request,
http::StatusCode,
middleware::Next,
response::{IntoResponse, Response},
};
use tracing::warn;
pub async fn header_limits_middleware(
request: Request<Body>,
next: Next,
max_header_count: usize,
max_header_bytes: usize,
) -> Response {
let headers = request.headers();
let header_count = headers.len();
if header_count > max_header_count {
warn!(header_count, max_header_count, "Request rejected: too many headers");
return (StatusCode::REQUEST_HEADER_FIELDS_TOO_LARGE, "Too many request headers")
.into_response();
}
let total_bytes: usize =
headers.iter().map(|(name, value)| name.as_str().len() + value.len()).sum();
if total_bytes > max_header_bytes {
warn!(total_bytes, max_header_bytes, "Request rejected: headers too large");
return (StatusCode::REQUEST_HEADER_FIELDS_TOO_LARGE, "Request headers too large")
.into_response();
}
next.run(request).await
}