tower_async_http/limit/
service.rs

1use super::body::create_error_response;
2use super::{RequestBodyLimitLayer, ResponseBody};
3
4use http::{Request, Response};
5use http_body::Body;
6use http_body_util::Limited;
7use tower_async_service::Service;
8
9/// Middleware that intercepts requests with body lengths greater than the
10/// configured limit and converts them into `413 Payload Too Large` responses.
11///
12/// See the [module docs](crate::limit) for an example.
13#[derive(Clone, Copy, Debug)]
14pub struct RequestBodyLimit<S> {
15    pub(crate) inner: S,
16    pub(crate) limit: usize,
17}
18
19impl<S> RequestBodyLimit<S> {
20    /// Create a new `RequestBodyLimit` with the given body length limit.
21    pub fn new(inner: S, limit: usize) -> Self {
22        Self { inner, limit }
23    }
24
25    define_inner_service_accessors!();
26
27    /// Returns a new [`Layer`] that wraps services with a `RequestBodyLimit` middleware.
28    ///
29    /// [`Layer`]: tower_async_layer::Layer
30    pub fn layer(limit: usize) -> RequestBodyLimitLayer {
31        RequestBodyLimitLayer::new(limit)
32    }
33}
34
35impl<ReqBody, ResBody, S> Service<Request<ReqBody>> for RequestBodyLimit<S>
36where
37    ResBody: Body,
38    S: Service<Request<Limited<ReqBody>>, Response = Response<ResBody>>,
39{
40    type Response = Response<ResponseBody<ResBody>>;
41    type Error = S::Error;
42
43    async fn call(&self, req: Request<ReqBody>) -> Result<Self::Response, Self::Error> {
44        let content_length = req
45            .headers()
46            .get(http::header::CONTENT_LENGTH)
47            .and_then(|value| value.to_str().ok()?.parse::<usize>().ok());
48
49        let body_limit = match content_length {
50            Some(len) if len > self.limit => return Ok(create_error_response()),
51            Some(len) => self.limit.min(len),
52            None => self.limit,
53        };
54
55        let req = req.map(|body| Limited::new(body, body_limit));
56        Ok(self.inner.call(req).await?.map(ResponseBody::new))
57    }
58}