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 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}