rama_http/layer/
body_limit.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
//! Apply a limit to the request body.
//!
//! # Example
//!
//! ```
//! use rama_http::{Body, Request, Response};
//! use std::convert::Infallible;
//! use rama_core::service::service_fn;
//! use rama_core::{Context, Layer, Service};
//! use rama_http::layer::body_limit::BodyLimitLayer;
//!
//! async fn handle<B>(_: Request<B>) -> Result<Response, Infallible> {
//!     // ...
//!     # Ok(Response::new(Body::default()))
//! }
//!
//! # #[tokio::main]
//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
//! let mut svc = (
//!      // Limit the request body to 2MB
//!     BodyLimitLayer::new(2*1024*1024),
//! ).layer(service_fn(handle));
//!
//! // Call the service
//! let request = Request::new(Body::default());
//!
//! svc.serve(Context::default(), request).await?;
//! # Ok(())
//! # }
//! ```

use crate::dep::http_body_util::Limited;
use crate::Request;
use bytes::Bytes;
use rama_core::{error::BoxError, Context, Layer, Service};
use rama_http_types::Body;
use rama_utils::macros::define_inner_service_accessors;
use std::fmt;

/// Apply a limit to the request body's size.
///
/// See the [module docs](crate::layer::body_limit) for an example.
#[derive(Debug, Clone)]
pub struct BodyLimitLayer {
    size: usize,
}

impl BodyLimitLayer {
    /// Create a new [`BodyLimitLayer`].
    pub const fn new(size: usize) -> Self {
        Self { size }
    }
}

impl<S> Layer<S> for BodyLimitLayer {
    type Service = BodyLimitService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        BodyLimitService::new(inner, self.size)
    }
}

/// Apply a transformation to the request body.
///
/// See the [module docs](crate::layer::body_limit) for an example.
#[derive(Clone)]
pub struct BodyLimitService<S> {
    inner: S,
    size: usize,
}

impl<S> BodyLimitService<S> {
    /// Create a new [`BodyLimitService`].
    pub const fn new(service: S, size: usize) -> Self {
        Self {
            inner: service,
            size,
        }
    }

    define_inner_service_accessors!();
}

impl<S, State, ReqBody> Service<State, Request<ReqBody>> for BodyLimitService<S>
where
    S: Service<State, Request<Body>>,
    State: Clone + Send + Sync + 'static,
    ReqBody: http_body::Body<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
{
    type Response = S::Response;
    type Error = S::Error;

    async fn serve(
        &self,
        ctx: Context<State>,
        req: Request<ReqBody>,
    ) -> Result<Self::Response, Self::Error> {
        let req = req.map(|body| {
            if self.size == 0 {
                Body::new(body)
            } else {
                Body::new(Limited::new(body, self.size))
            }
        });
        self.inner.serve(ctx, req).await
    }
}

impl<S> fmt::Debug for BodyLimitService<S>
where
    S: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("BodyLimitService")
            .field("inner", &self.inner)
            .field("size", &self.size)
            .finish()
    }
}