api_tools/server/axum/layers/
body_limiter.rs

1//! Request body size limiter layer
2//!
3//! Rejects incoming requests whose advertised `Content-Length` exceeds a
4//! configured maximum, **before** the body is read. The check is O(1) and does
5//! not buffer anything, so it preserves streaming for legitimate requests.
6//!
7//! # Scope
8//!
9//! The limit is enforced from the `Content-Length` header. A chunked request
10//! that omits `Content-Length` is **not** capped by this layer.
11//!
12//! # Hard guarantee on chunked bodies
13//!
14//! For a hard cap that also covers requests without a `Content-Length` (chunked
15//! / streamed uploads), stack this layer on top of the re-exported
16//! [`RequestBodyLimitLayer`]. `BodyLimiterLayer` rejects the common
17//! honest-client case early with a JSON `413`, while `RequestBodyLimitLayer`
18//! wraps the body and enforces the same limit as it is read, backstopping the
19//! chunked case (with tower-http's plain `413`):
20//!
21//! ```
22//! use api_tools::server::axum::layers::body_limiter::{
23//!     BodyLimiterConfig, BodyLimiterLayer, RequestBodyLimitLayer,
24//! };
25//! use tower::ServiceBuilder;
26//!
27//! const MAX: usize = 2 * 1024 * 1024; // 2 MiB
28//!
29//! let _layers = ServiceBuilder::new()
30//!     .layer(BodyLimiterLayer::new(&BodyLimiterConfig { body_max_size: MAX }))
31//!     .layer(RequestBodyLimitLayer::new(MAX));
32//! ```
33
34use crate::server::axum::response::ApiError;
35use axum::body::Body;
36use axum::http::Request;
37use axum::http::header::CONTENT_LENGTH;
38use axum::response::{IntoResponse, Response};
39use futures::future::BoxFuture;
40use std::task::{Context, Poll};
41use tower::{Layer, Service};
42#[doc(inline)]
43pub use tower_http::limit::RequestBodyLimitLayer;
44
45/// Configuration for the `BodyLimiterLayer`
46#[derive(Clone, Debug)]
47pub struct BodyLimiterConfig {
48    /// Maximum size of the request body in bytes
49    pub body_max_size: usize,
50}
51
52#[derive(Clone)]
53pub struct BodyLimiterLayer {
54    pub config: BodyLimiterConfig,
55}
56
57impl BodyLimiterLayer {
58    /// Create a new `BodyLimiterLayer`
59    pub fn new(config: &BodyLimiterConfig) -> Self {
60        Self { config: config.clone() }
61    }
62}
63
64impl<S> Layer<S> for BodyLimiterLayer {
65    type Service = BodyLimiterMiddleware<S>;
66
67    fn layer(&self, inner: S) -> Self::Service {
68        BodyLimiterMiddleware {
69            inner,
70            config: self.config.clone(),
71        }
72    }
73}
74
75#[derive(Clone)]
76pub struct BodyLimiterMiddleware<S> {
77    inner: S,
78    config: BodyLimiterConfig,
79}
80
81impl<S> Service<Request<Body>> for BodyLimiterMiddleware<S>
82where
83    S: Service<Request<Body>, Response = Response> + Send + 'static,
84    S::Future: Send + 'static,
85{
86    type Response = S::Response;
87    type Error = S::Error;
88    // `BoxFuture` is a type alias for `Pin<Box<dyn Future + Send + 'a>>`
89    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
90
91    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
92        self.inner.poll_ready(cx)
93    }
94
95    fn call(&mut self, request: Request<Body>) -> Self::Future {
96        // Reject early, without reading the body, when the advertised
97        // `Content-Length` exceeds the limit. This is O(1) and preserves
98        // streaming for legitimate requests.
99        let content_length = request
100            .headers()
101            .get(CONTENT_LENGTH)
102            .and_then(|value| value.to_str().ok())
103            .and_then(|value| value.parse::<usize>().ok());
104
105        if let Some(length) = content_length
106            && length > self.config.body_max_size
107        {
108            return Box::pin(async move { Ok(ApiError::PayloadTooLarge.into_response()) });
109        }
110
111        Box::pin(self.inner.call(request))
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use axum::http::StatusCode;
119    use std::convert::Infallible;
120    use tower::{ServiceBuilder, ServiceExt};
121
122    fn layer() -> BodyLimiterLayer {
123        BodyLimiterLayer::new(&BodyLimiterConfig { body_max_size: 8 })
124    }
125
126    fn ok_service() -> impl Service<Request<Body>, Response = Response, Error = Infallible> + Clone {
127        ServiceBuilder::new()
128            .layer(layer())
129            .service(tower::service_fn(|_req: Request<Body>| async {
130                Ok::<_, Infallible>(
131                    Response::builder()
132                        .status(StatusCode::OK)
133                        .body(Body::from("ok"))
134                        .unwrap(),
135                )
136            }))
137    }
138
139    async fn read_body(response: Response) -> String {
140        let body = axum::body::to_bytes(response.into_body(), 4096).await.unwrap();
141        String::from_utf8(body.to_vec()).unwrap()
142    }
143
144    /// A `Content-Length` above the limit is rejected with a JSON `413` and the
145    /// inner handler is never invoked.
146    #[tokio::test]
147    async fn content_length_over_limit_is_rejected_with_413() {
148        let svc = ServiceBuilder::new()
149            .layer(layer())
150            .service(tower::service_fn(|_req: Request<Body>| async {
151                // Should never be reached.
152                Ok::<_, Infallible>(
153                    Response::builder()
154                        .status(StatusCode::OK)
155                        .body(Body::from("ok"))
156                        .unwrap(),
157                )
158            }));
159
160        let response = svc
161            .oneshot(
162                Request::builder()
163                    .uri("/")
164                    .header(CONTENT_LENGTH, "9")
165                    .body(Body::from("123456789"))
166                    .unwrap(),
167            )
168            .await
169            .unwrap();
170
171        assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
172        let body = read_body(response).await;
173        assert!(body.contains("\"code\":413"), "body was: {body}");
174    }
175
176    /// A `Content-Length` at the limit passes through.
177    #[tokio::test]
178    async fn content_length_at_limit_passes_through() {
179        let response = ok_service()
180            .oneshot(
181                Request::builder()
182                    .uri("/")
183                    .header(CONTENT_LENGTH, "8")
184                    .body(Body::from("12345678"))
185                    .unwrap(),
186            )
187            .await
188            .unwrap();
189
190        assert_eq!(response.status(), StatusCode::OK);
191        assert_eq!(read_body(response).await, "ok");
192    }
193
194    /// A `Content-Length` below the limit passes through.
195    #[tokio::test]
196    async fn content_length_under_limit_passes_through() {
197        let response = ok_service()
198            .oneshot(
199                Request::builder()
200                    .uri("/")
201                    .header(CONTENT_LENGTH, "3")
202                    .body(Body::from("abc"))
203                    .unwrap(),
204            )
205            .await
206            .unwrap();
207
208        assert_eq!(response.status(), StatusCode::OK);
209        assert_eq!(read_body(response).await, "ok");
210    }
211
212    /// Without `Content-Length` the request is not capped by this layer
213    /// (documented limitation) and passes through.
214    #[tokio::test]
215    async fn missing_content_length_passes_through() {
216        let response = ok_service()
217            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
218            .await
219            .unwrap();
220
221        assert_eq!(response.status(), StatusCode::OK);
222        assert_eq!(read_body(response).await, "ok");
223    }
224}
api_tools/server/axum/layers/body_limiter.rs

api_tools/server/axum/layers/
body_limiter.rs
