tower_async_http/limit/
service.rs1use 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#[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 pub fn new(inner: S, limit: usize) -> Self {
22 Self { inner, limit }
23 }
24
25 define_inner_service_accessors!();
26
27 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}