rama_net/stream/layer/http/
body_limit.rs

1use crate::stream::Stream;
2use rama_core::{Context, Layer, Service};
3use rama_http_types::BodyLimit;
4use rama_utils::macros::define_inner_service_accessors;
5use std::fmt;
6
7/// Limit the size of the request and/or response bodies.
8///
9/// As this layer operates on the transport layer ([`Stream`]),
10/// it only is used to add the [`BodyLimit`] value to the [`Context`],
11/// such that the L7 http service can apply the limit when found in that [`Context`].
12///
13/// [`Stream`]: crate::stream::Stream
14/// [`Context`]: rama_core::Context`
15#[derive(Debug, Clone)]
16pub struct BodyLimitLayer {
17    limit: BodyLimit,
18}
19
20impl BodyLimitLayer {
21    /// Create a new [`BodyLimitLayer`], with the given limit to be applied to the request only.
22    ///
23    /// See [`BodyLimitLayer`] for more information.
24    pub fn request_only(limit: usize) -> Self {
25        Self {
26            limit: BodyLimit::request_only(limit),
27        }
28    }
29
30    /// Create a new [`BodyLimitLayer`], with the given limit to be applied to the response only.
31    ///
32    /// See [`BodyLimitLayer`] for more information.
33    pub fn response_only(limit: usize) -> Self {
34        Self {
35            limit: BodyLimit::response_only(limit),
36        }
37    }
38
39    /// Create a new [`BodyLimitLayer`], with the given limit to be applied to both the request and response bodies.
40    ///
41    /// See [`BodyLimitLayer`] for more information.
42    pub fn symmetric(limit: usize) -> Self {
43        Self {
44            limit: BodyLimit::symmetric(limit),
45        }
46    }
47
48    /// Create a new [`BodyLimitLayer`], with the given limits
49    /// respectively to be applied to the request and response bodies.
50    ///
51    /// See [`BodyLimitLayer`] for more information.
52    pub fn asymmetric(request: usize, response: usize) -> Self {
53        Self {
54            limit: BodyLimit::asymmetric(request, response),
55        }
56    }
57}
58
59impl<S> Layer<S> for BodyLimitLayer {
60    type Service = BodyLimitService<S>;
61
62    fn layer(&self, inner: S) -> Self::Service {
63        BodyLimitService::new(inner, self.limit)
64    }
65}
66
67/// Communicate to the downstream http service to apply a limit to the body.
68///
69/// See [`BodyLimitService`] for more information.
70#[derive(Clone)]
71pub struct BodyLimitService<S> {
72    inner: S,
73    limit: BodyLimit,
74}
75
76impl<S> BodyLimitService<S> {
77    /// Create a new [`BodyLimitService`].
78    pub const fn new(service: S, limit: BodyLimit) -> Self {
79        Self {
80            inner: service,
81            limit,
82        }
83    }
84
85    define_inner_service_accessors!();
86
87    /// Create a new [`BodyLimitService`], with the given limit to be applied to the request only.
88    ///
89    /// See [`BodyLimitLayer`] for more information.
90    pub fn request_only(service: S, limit: usize) -> Self {
91        BodyLimitLayer::request_only(limit).into_layer(service)
92    }
93
94    /// Create a new [`BodyLimitService`], with the given limit to be applied to the response only.
95    ///
96    /// See [`BodyLimitLayer`] for more information.
97    pub fn response_only(service: S, limit: usize) -> Self {
98        BodyLimitLayer::response_only(limit).into_layer(service)
99    }
100
101    /// Create a new [`BodyLimitService`], with the given limit to be applied to both the request and response bodies.
102    ///
103    /// See [`BodyLimitLayer`] for more information.
104    pub fn symmetric(service: S, limit: usize) -> Self {
105        BodyLimitLayer::symmetric(limit).into_layer(service)
106    }
107
108    /// Create a new [`BodyLimitService`], with the given limits
109    /// respectively to be applied to the request and response bodies.
110    ///
111    /// See [`BodyLimitLayer`] for more information.
112    pub fn asymmetric(service: S, request: usize, response: usize) -> Self {
113        BodyLimitLayer::asymmetric(request, response).into_layer(service)
114    }
115}
116
117impl<S, State, IO> Service<State, IO> for BodyLimitService<S>
118where
119    S: Service<State, IO>,
120    State: Clone + Send + Sync + 'static,
121    IO: Stream,
122{
123    type Response = S::Response;
124    type Error = S::Error;
125
126    async fn serve(
127        &self,
128        mut ctx: Context<State>,
129        stream: IO,
130    ) -> Result<Self::Response, Self::Error> {
131        ctx.insert(self.limit);
132        self.inner.serve(ctx, stream).await
133    }
134}
135
136impl<S> fmt::Debug for BodyLimitService<S>
137where
138    S: fmt::Debug,
139{
140    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
141        f.debug_struct("BodyLimitService")
142            .field("inner", &self.inner)
143            .field("limit", &self.limit)
144            .finish()
145    }
146}
rama_net/stream/layer/http/body_limit.rs

rama_net/stream/layer/http/
body_limit.rs
