ts_webapi/middleware/
response_body.rs

1//! A response body wrapper to allow injecting custom response body types in middleware.
2
3use core::{
4    pin::Pin,
5    task::{Context, Poll},
6};
7
8use bytes::Bytes;
9use http_body::{Body, SizeHint};
10use http_body_util::{Empty, Full};
11use pin_project_lite::pin_project;
12
13pin_project! {
14    /// A response body wrapper to allow injecting custom response body types in middleware.
15    pub struct ResponseBody<B> {
16        #[pin]
17        inner: ResponseBodyInner<B>
18    }
19}
20
21impl<B> ResponseBody<B> {
22    /// Create a new response body from the canonical body `B`
23    pub(crate) fn new(body: B) -> Self {
24        Self {
25            inner: ResponseBodyInner::Body { body },
26        }
27    }
28
29    /// Create a new response body from some bytes.
30    pub(crate) fn full(body: Bytes) -> Self {
31        Self {
32            inner: ResponseBodyInner::Full {
33                body: Full::new(body),
34            },
35        }
36    }
37
38    /// Create an empty response body.
39    pub(crate) fn empty() -> Self {
40        Self {
41            inner: ResponseBodyInner::Empty { body: Empty::new() },
42        }
43    }
44}
45
46pin_project! {
47    #[project = BodyProj]
48    enum ResponseBodyInner<B> {
49        Body {
50            #[pin]
51            body: B
52        },
53        Full {
54            #[pin]
55            body: Full<Bytes>
56        },
57        Empty {
58            #[pin]
59            body: Empty<Bytes>
60        }
61    }
62}
63
64impl<B> Body for ResponseBody<B>
65where
66    B: Body<Data = Bytes>,
67{
68    type Data = Bytes;
69    type Error = B::Error;
70
71    fn poll_frame(
72        self: Pin<&mut Self>,
73        cx: &mut Context<'_>,
74    ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
75        match self.project().inner.project() {
76            BodyProj::Body { body } => body.poll_frame(cx),
77            BodyProj::Full { body } => body.poll_frame(cx).map_err(|err| match err {}),
78            BodyProj::Empty { body } => body.poll_frame(cx).map_err(|err| match err {}),
79        }
80    }
81
82    fn is_end_stream(&self) -> bool {
83        match &self.inner {
84            ResponseBodyInner::Body { body } => body.is_end_stream(),
85            ResponseBodyInner::Full { body } => body.is_end_stream(),
86            ResponseBodyInner::Empty { body } => body.is_end_stream(),
87        }
88    }
89
90    fn size_hint(&self) -> SizeHint {
91        match &self.inner {
92            ResponseBodyInner::Body { body } => body.size_hint(),
93            ResponseBodyInner::Full { body } => body.size_hint(),
94            ResponseBodyInner::Empty { body } => body.size_hint(),
95        }
96    }
97}