1 use self::private::DefaultBodyLimitService;
2 use tower_layer::Layer;
3 
4 /// Layer for configuring the default request body limit.
5 ///
6 /// For security reasons, [`Bytes`] will, by default, not accept bodies larger than 2MB. This also
7 /// applies to extractors that uses [`Bytes`] internally such as `String`, [`Json`], and [`Form`].
8 ///
9 /// This middleware provides ways to configure that.
10 ///
11 /// Note that if an extractor consumes the body directly with [`Body::data`], or similar, the
12 /// default limit is _not_ applied.
13 ///
14 /// # Difference between `DefaultBodyLimit` and [`RequestBodyLimit`]
15 ///
16 /// `DefaultBodyLimit` and [`RequestBodyLimit`] serve similar functions but in different ways.
17 ///
18 /// `DefaultBodyLimit` is local in that it only applies to [`FromRequest`] implementations that
19 /// explicitly apply it (or call another extractor that does). You can apply the limit with
20 /// [`RequestExt::with_limited_body`] or [`RequestExt::into_limited_body`]
21 ///
22 /// [`RequestBodyLimit`] is applied globally to all requests, regardless of which extractors are
23 /// used or how the body is consumed.
24 ///
25 /// `DefaultBodyLimit` is also easier to integrate into an existing setup since it doesn't change
26 /// the request body type:
27 ///
28 /// ```
29 /// use axum::{
30 ///     Router,
31 ///     routing::post,
32 ///     body::Body,
33 ///     extract::{DefaultBodyLimit, RawBody},
34 ///     http::Request,
35 /// };
36 ///
37 /// let app = Router::new()
38 ///     .route(
39 ///         "/",
40 ///         // even with `DefaultBodyLimit` the request body is still just `Body`
41 ///         post(|request: Request<Body>| async {}),
42 ///     )
43 ///     .layer(DefaultBodyLimit::max(1024));
44 /// # let _: Router<(), _> = app;
45 /// ```
46 ///
47 /// ```
48 /// use axum::{Router, routing::post, body::Body, extract::RawBody, http::Request};
49 /// use tower_http::limit::RequestBodyLimitLayer;
50 /// use http_body::Limited;
51 ///
52 /// let app = Router::new()
53 ///     .route(
54 ///         "/",
55 ///         // `RequestBodyLimitLayer` changes the request body type to `Limited<Body>`
56 ///         // extracting a different body type wont work
57 ///         post(|request: Request<Limited<Body>>| async {}),
58 ///     )
59 ///     .layer(RequestBodyLimitLayer::new(1024));
60 /// # let _: Router<(), _> = app;
61 /// ```
62 ///
63 /// In general using `DefaultBodyLimit` is recommended but if you need to use third party
64 /// extractors and want to sure a limit is also applied there then [`RequestBodyLimit`] should be
65 /// used.
66 ///
67 /// [`Body::data`]: http_body::Body::data
68 /// [`Bytes`]: bytes::Bytes
69 /// [`Json`]: https://docs.rs/axum/0.6.0/axum/struct.Json.html
70 /// [`Form`]: https://docs.rs/axum/0.6.0/axum/struct.Form.html
71 /// [`FromRequest`]: crate::extract::FromRequest
72 /// [`RequestBodyLimit`]: tower_http::limit::RequestBodyLimit
73 /// [`RequestExt::with_limited_body`]: crate::RequestExt::with_limited_body
74 /// [`RequestExt::into_limited_body`]: crate::RequestExt::into_limited_body
75 #[derive(Debug, Clone)]
76 #[must_use]
77 pub struct DefaultBodyLimit {
78     kind: DefaultBodyLimitKind,
79 }
80 
81 #[derive(Debug, Clone, Copy)]
82 pub(crate) enum DefaultBodyLimitKind {
83     Disable,
84     Limit(usize),
85 }
86 
87 impl DefaultBodyLimit {
88     /// Disable the default request body limit.
89     ///
90     /// This must be used to receive bodies larger than the default limit of 2MB using [`Bytes`] or
91     /// an extractor built on it such as `String`, [`Json`], [`Form`].
92     ///
93     /// Note that if you're accepting data from untrusted remotes it is recommend to add your own
94     /// limit such as [`tower_http::limit`].
95     ///
96     /// # Example
97     ///
98     /// ```
99     /// use axum::{
100     ///     Router,
101     ///     routing::get,
102     ///     body::{Bytes, Body},
103     ///     extract::DefaultBodyLimit,
104     /// };
105     /// use tower_http::limit::RequestBodyLimitLayer;
106     /// use http_body::Limited;
107     ///
108     /// let app: Router<(), Limited<Body>> = Router::new()
109     ///     .route("/", get(|body: Bytes| async {}))
110     ///     // Disable the default limit
111     ///     .layer(DefaultBodyLimit::disable())
112     ///     // Set a different limit
113     ///     .layer(RequestBodyLimitLayer::new(10 * 1000 * 1000));
114     /// ```
115     ///
116     /// [`Bytes`]: bytes::Bytes
117     /// [`Json`]: https://docs.rs/axum/0.6.0/axum/struct.Json.html
118     /// [`Form`]: https://docs.rs/axum/0.6.0/axum/struct.Form.html
disable() -> Self119     pub fn disable() -> Self {
120         Self {
121             kind: DefaultBodyLimitKind::Disable,
122         }
123     }
124 
125     /// Set the default request body limit.
126     ///
127     /// By default the limit of request body sizes that [`Bytes::from_request`] (and other
128     /// extractors built on top of it such as `String`, [`Json`], and [`Form`]) is 2MB. This method
129     /// can be used to change that limit.
130     ///
131     /// # Example
132     ///
133     /// ```
134     /// use axum::{
135     ///     Router,
136     ///     routing::get,
137     ///     body::{Bytes, Body},
138     ///     extract::DefaultBodyLimit,
139     /// };
140     /// use tower_http::limit::RequestBodyLimitLayer;
141     /// use http_body::Limited;
142     ///
143     /// let app: Router<(), Limited<Body>> = Router::new()
144     ///     .route("/", get(|body: Bytes| async {}))
145     ///     // Replace the default of 2MB with 1024 bytes.
146     ///     .layer(DefaultBodyLimit::max(1024));
147     /// ```
148     ///
149     /// [`Bytes::from_request`]: bytes::Bytes
150     /// [`Json`]: https://docs.rs/axum/0.6.0/axum/struct.Json.html
151     /// [`Form`]: https://docs.rs/axum/0.6.0/axum/struct.Form.html
max(limit: usize) -> Self152     pub fn max(limit: usize) -> Self {
153         Self {
154             kind: DefaultBodyLimitKind::Limit(limit),
155         }
156     }
157 }
158 
159 impl<S> Layer<S> for DefaultBodyLimit {
160     type Service = DefaultBodyLimitService<S>;
161 
layer(&self, inner: S) -> Self::Service162     fn layer(&self, inner: S) -> Self::Service {
163         DefaultBodyLimitService {
164             inner,
165             kind: self.kind,
166         }
167     }
168 }
169 
170 mod private {
171     use super::DefaultBodyLimitKind;
172     use http::Request;
173     use std::task::Context;
174     use tower_service::Service;
175 
176     #[derive(Debug, Clone, Copy)]
177     pub struct DefaultBodyLimitService<S> {
178         pub(super) inner: S,
179         pub(super) kind: DefaultBodyLimitKind,
180     }
181 
182     impl<B, S> Service<Request<B>> for DefaultBodyLimitService<S>
183     where
184         S: Service<Request<B>>,
185     {
186         type Response = S::Response;
187         type Error = S::Error;
188         type Future = S::Future;
189 
190         #[inline]
poll_ready(&mut self, cx: &mut Context<'_>) -> std::task::Poll<Result<(), Self::Error>>191         fn poll_ready(&mut self, cx: &mut Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
192             self.inner.poll_ready(cx)
193         }
194 
195         #[inline]
call(&mut self, mut req: Request<B>) -> Self::Future196         fn call(&mut self, mut req: Request<B>) -> Self::Future {
197             req.extensions_mut().insert(self.kind);
198             self.inner.call(req)
199         }
200     }
201 }
202