tower_http/limit/
body.rs

1use bytes::Bytes;
2use http::{HeaderValue, Response, StatusCode};
3use http_body::{Body, SizeHint};
4use http_body_util::Full;
5use pin_project_lite::pin_project;
6use std::pin::Pin;
7use std::task::{Context, Poll};
8
9pin_project! {
10    /// Response body for [`RequestBodyLimit`].
11    ///
12    /// [`RequestBodyLimit`]: super::RequestBodyLimit
13    pub struct ResponseBody<B> {
14        #[pin]
15        inner: ResponseBodyInner<B>
16    }
17}
18
19impl<B> ResponseBody<B> {
20    fn payload_too_large() -> Self {
21        Self {
22            inner: ResponseBodyInner::PayloadTooLarge {
23                body: Full::from(BODY),
24            },
25        }
26    }
27
28    pub(crate) fn new(body: B) -> Self {
29        Self {
30            inner: ResponseBodyInner::Body { body },
31        }
32    }
33}
34
35pin_project! {
36    #[project = BodyProj]
37    enum ResponseBodyInner<B> {
38        PayloadTooLarge {
39            #[pin]
40            body: Full<Bytes>,
41        },
42        Body {
43            #[pin]
44            body: B
45        }
46    }
47}
48
49impl<B> Body for ResponseBody<B>
50where
51    B: Body<Data = Bytes>,
52{
53    type Data = Bytes;
54    type Error = B::Error;
55
56    fn poll_frame(
57        self: Pin<&mut Self>,
58        cx: &mut Context<'_>,
59    ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
60        match self.project().inner.project() {
61            BodyProj::PayloadTooLarge { body } => body.poll_frame(cx).map_err(|err| match err {}),
62            BodyProj::Body { body } => body.poll_frame(cx),
63        }
64    }
65
66    fn is_end_stream(&self) -> bool {
67        match &self.inner {
68            ResponseBodyInner::PayloadTooLarge { body } => body.is_end_stream(),
69            ResponseBodyInner::Body { body } => body.is_end_stream(),
70        }
71    }
72
73    fn size_hint(&self) -> SizeHint {
74        match &self.inner {
75            ResponseBodyInner::PayloadTooLarge { body } => body.size_hint(),
76            ResponseBodyInner::Body { body } => body.size_hint(),
77        }
78    }
79}
80
81const BODY: &[u8] = b"length limit exceeded";
82
83pub(crate) fn create_error_response<B>() -> Response<ResponseBody<B>>
84where
85    B: Body,
86{
87    let mut res = Response::new(ResponseBody::payload_too_large());
88    *res.status_mut() = StatusCode::PAYLOAD_TOO_LARGE;
89
90    #[allow(clippy::declare_interior_mutable_const)]
91    const TEXT_PLAIN: HeaderValue = HeaderValue::from_static("text/plain; charset=utf-8");
92    res.headers_mut()
93        .insert(http::header::CONTENT_TYPE, TEXT_PLAIN);
94
95    res
96}