xitca_web/handler/types/
body.rs

1//! type extractor for request body stream.
2
3use core::{cmp, convert::Infallible, pin::pin};
4
5use crate::{
6    body::{BodyExt, BodyStream, BoxBody, ResponseBody, SizeHint},
7    bytes::{Bytes, BytesMut},
8    context::WebContext,
9    error::{BodyOverFlow, Error},
10    handler::{FromRequest, Responder},
11    http::{IntoResponse, WebResponse},
12};
13
14pub struct Body<B>(pub B);
15
16impl<'a, 'r, C, B> FromRequest<'a, WebContext<'r, C, B>> for Body<B>
17where
18    B: BodyStream + Default,
19{
20    type Type<'b> = Body<B>;
21    type Error = Error;
22
23    #[inline]
24    async fn from_request(ctx: &'a WebContext<'r, C, B>) -> Result<Self, Self::Error> {
25        Ok(Body(ctx.take_body_ref()))
26    }
27}
28
29/// helper type for limiting body size.
30/// when LIMIT > 0 body size is limited to LIMIT in bytes.
31/// when LIMIT == 0 body size is unlimited.
32pub struct Limit<const LIMIT: usize>;
33
34macro_rules! from_bytes_impl {
35    ($type: ty, $original: tt) => {
36        impl<'a, 'r, C, B, const LIMIT: usize> FromRequest<'a, WebContext<'r, C, B>> for ($type, Limit<LIMIT>)
37        where
38            B: BodyStream + Default,
39        {
40            type Type<'b> = ($type, Limit<LIMIT>);
41            type Error = Error;
42
43            async fn from_request(ctx: &'a WebContext<'r, C, B>) -> Result<Self, Self::Error> {
44                let body = ctx.take_body_ref();
45
46                let limit = match body.size_hint() {
47                    SizeHint::Exact(size) if LIMIT > 0 => cmp::min(size as usize, LIMIT),
48                    SizeHint::Exact(size) => size as usize,
49                    _ => LIMIT,
50                };
51
52                let mut body = pin!(body);
53
54                let mut buf = <$type>::with_capacity(limit);
55
56                while let Some(data) = body.as_mut().data().await {
57                    let data = data.map_err(Into::into)?;
58                    buf.extend_from_slice(data.as_ref());
59                    if limit > 0 && buf.len() > limit {
60                        return Err(Error::from(BodyOverFlow { limit }));
61                    }
62                }
63
64                Ok((buf, Limit))
65            }
66        }
67
68        from_bytes_impl!($type);
69    };
70    ($type: ty) => {
71        impl<'a, 'r, C, B> FromRequest<'a, WebContext<'r, C, B>> for $type
72        where
73            B: BodyStream + Default,
74        {
75            type Type<'b> = $type;
76            type Error = Error;
77
78            #[inline]
79            async fn from_request(ctx: &'a WebContext<'r, C, B>) -> Result<Self, Self::Error> {
80                <($type, Limit<0>)>::from_request(ctx)
81                    .await
82                    .map(|(bytes, _)| bytes)
83            }
84        }
85    };
86}
87
88from_bytes_impl!(BytesMut, _);
89from_bytes_impl!(Vec<u8>, _);
90
91impl<'a, 'r, C, B, const LIMIT: usize> FromRequest<'a, WebContext<'r, C, B>> for (Bytes, Limit<LIMIT>)
92where
93    B: BodyStream + Default,
94{
95    type Type<'b> = (Bytes, Limit<LIMIT>);
96    type Error = Error;
97
98    #[inline]
99    async fn from_request(ctx: &'a WebContext<'r, C, B>) -> Result<Self, Self::Error> {
100        <(BytesMut, Limit<LIMIT>)>::from_request(ctx)
101            .await
102            .map(|(bytes, limit)| (bytes.into(), limit))
103    }
104}
105
106from_bytes_impl!(Bytes);
107
108macro_rules! responder_impl {
109    ($type: ty) => {
110        impl<'r, C, B> Responder<WebContext<'r, C, B>> for $type {
111            type Response = WebResponse;
112            type Error = Infallible;
113
114            #[inline]
115            async fn respond(self, ctx: WebContext<'r, C, B>) -> Result<Self::Response, Self::Error> {
116                Ok(ctx.into_response(self))
117            }
118
119            #[inline]
120            fn map(self, res: Self::Response) -> Result<Self::Response, Self::Error> {
121                Ok(res.map(|_| self.into()))
122            }
123        }
124    };
125}
126
127responder_impl!(Bytes);
128responder_impl!(BytesMut);
129responder_impl!(Vec<u8>);
130
131impl<'r, C, B, ResB> Responder<WebContext<'r, C, B>> for ResponseBody<ResB> {
132    type Response = WebResponse<ResponseBody<ResB>>;
133    type Error = Error;
134
135    #[inline]
136    async fn respond(self, ctx: WebContext<'r, C, B>) -> Result<Self::Response, Self::Error> {
137        Ok(ctx.req.as_response(self))
138    }
139
140    #[inline]
141    fn map(self, res: Self::Response) -> Result<Self::Response, Self::Error> {
142        Ok(res.map(|_| self))
143    }
144}
145
146impl<'r, C, B> Responder<WebContext<'r, C, B>> for BoxBody {
147    type Response = WebResponse;
148    type Error = Error;
149
150    #[inline]
151    async fn respond(self, ctx: WebContext<'r, C, B>) -> Result<Self::Response, Self::Error> {
152        ResponseBody::body(self).respond(ctx).await
153    }
154
155    #[inline]
156    fn map(self, res: Self::Response) -> Result<Self::Response, Self::Error> {
157        Responder::<WebContext<'r, C, B>>::map(ResponseBody::body(self), res)
158    }
159}
xitca_web/handler/types/body.rs

xitca_web/handler/types/
body.rs
