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::{Context, 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(Context::default(), request).await?;
28//! # Ok(())
29//! # }
30//! ```
31
32use crate::Request;
33use crate::dep::http_body_util::Limited;
34use bytes::Bytes;
35use rama_core::{Context, Layer, Service, error::BoxError};
36use rama_http_types::Body;
37use rama_utils::macros::define_inner_service_accessors;
38use std::fmt;
39
40/// Apply a limit to the request body's size.
41///
42/// See the [module docs](crate::layer::body_limit) for an example.
43#[derive(Debug, Clone)]
44pub struct BodyLimitLayer {
45    size: usize,
46}
47
48impl BodyLimitLayer {
49    /// Create a new [`BodyLimitLayer`].
50    pub const fn new(size: usize) -> Self {
51        Self { size }
52    }
53}
54
55impl<S> Layer<S> for BodyLimitLayer {
56    type Service = BodyLimitService<S>;
57
58    fn layer(&self, inner: S) -> Self::Service {
59        BodyLimitService::new(inner, self.size)
60    }
61}
62
63/// Apply a transformation to the request body.
64///
65/// See the [module docs](crate::layer::body_limit) for an example.
66#[derive(Clone)]
67pub struct BodyLimitService<S> {
68    inner: S,
69    size: usize,
70}
71
72impl<S> BodyLimitService<S> {
73    /// Create a new [`BodyLimitService`].
74    pub const fn new(service: S, size: usize) -> Self {
75        Self {
76            inner: service,
77            size,
78        }
79    }
80
81    define_inner_service_accessors!();
82}
83
84impl<S, State, ReqBody> Service<State, Request<ReqBody>> for BodyLimitService<S>
85where
86    S: Service<State, Request<Body>>,
87    State: Clone + Send + Sync + 'static,
88    ReqBody: rama_http_types::dep::http_body::Body<Data = Bytes, Error: Into<BoxError>>
89        + Send
90        + Sync
91        + 'static,
92{
93    type Response = S::Response;
94    type Error = S::Error;
95
96    async fn serve(
97        &self,
98        ctx: Context<State>,
99        req: Request<ReqBody>,
100    ) -> Result<Self::Response, Self::Error> {
101        let req = req.map(|body| {
102            if self.size == 0 {
103                Body::new(body)
104            } else {
105                Body::new(Limited::new(body, self.size))
106            }
107        });
108        self.inner.serve(ctx, req).await
109    }
110}
111
112impl<S> fmt::Debug for BodyLimitService<S>
113where
114    S: fmt::Debug,
115{
116    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
117        f.debug_struct("BodyLimitService")
118            .field("inner", &self.inner)
119            .field("size", &self.size)
120            .finish()
121    }
122}
rama_http/layer/body_limit.rs

rama_http/layer/
body_limit.rs
