rama_http/layer/
body_limit.rs

1//! Apply a limit to the request body.
2//!
3//! # Example
4//!
5//! ```
6//! use rama_http::{Body, Request, Response};
7//! use std::convert::Infallible;
8//! use rama_core::service::service_fn;
9//! use rama_core::{Layer, Service};
10//! use rama_http::layer::body_limit::BodyLimitLayer;
11//!
12//! async fn handle<B>(_: Request<B>) -> Result<Response, Infallible> {
13//!     // ...
14//!     # Ok(Response::new(Body::default()))
15//! }
16//!
17//! # #[tokio::main]
18//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
19//! let mut svc = (
20//!      // Limit the request body to 2MB
21//!     BodyLimitLayer::new(2*1024*1024),
22//! ).into_layer(service_fn(handle));
23//!
24//! // Call the service
25//! let request = Request::new(Body::default());
26//!
27//! svc.serve(request).await?;
28//! # Ok(())
29//! # }
30//! ```
31
32use crate::{Body, Request, StreamingBody, body::util::Limited};
33use rama_core::{Layer, Service, bytes::Bytes, error::BoxError};
34use rama_utils::macros::define_inner_service_accessors;
35use std::fmt;
36
37/// Apply a limit to the request body's size.
38///
39/// See the [module docs](crate::layer::body_limit) for an example.
40#[derive(Debug, Clone)]
41pub struct BodyLimitLayer {
42    size: usize,
43}
44
45impl BodyLimitLayer {
46    /// Create a new [`BodyLimitLayer`].
47    #[must_use]
48    pub const fn new(size: usize) -> Self {
49        Self { size }
50    }
51}
52
53impl<S> Layer<S> for BodyLimitLayer {
54    type Service = BodyLimitService<S>;
55
56    fn layer(&self, inner: S) -> Self::Service {
57        BodyLimitService::new(inner, self.size)
58    }
59}
60
61/// Apply a transformation to the request body.
62///
63/// See the [module docs](crate::layer::body_limit) for an example.
64#[derive(Clone)]
65pub struct BodyLimitService<S> {
66    inner: S,
67    size: usize,
68}
69
70impl<S> BodyLimitService<S> {
71    /// Create a new [`BodyLimitService`].
72    pub const fn new(service: S, size: usize) -> Self {
73        Self {
74            inner: service,
75            size,
76        }
77    }
78
79    define_inner_service_accessors!();
80}
81
82impl<S, ReqBody> Service<Request<ReqBody>> for BodyLimitService<S>
83where
84    S: Service<Request<Body>>,
85    ReqBody: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
86{
87    type Output = S::Output;
88    type Error = S::Error;
89
90    async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
91        let req = req.map(|body| {
92            if self.size == 0 {
93                Body::new(body)
94            } else {
95                Body::new(Limited::new(body, self.size))
96            }
97        });
98        self.inner.serve(req).await
99    }
100}
101
102impl<S> fmt::Debug for BodyLimitService<S>
103where
104    S: fmt::Debug,
105{
106    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
107        f.debug_struct("BodyLimitService")
108            .field("inner", &self.inner)
109            .field("size", &self.size)
110            .finish()
111    }
112}
rama_http/layer/body_limit.rs

rama_http/layer/
body_limit.rs
