http_pool/body/
variant_body.rs

1use hyper::body::{Body, Buf, Frame, SizeHint};
2use std::error::Error;
3use std::pin::Pin;
4use std::task::{Context, Poll};
5
6/// 擦除掉具体的body和error
7trait ErasedBody: Send + Sync {
8    fn poll_frame(
9        self: Pin<&mut Self>,
10        cx: &mut Context<'_>,
11    ) -> Poll<Option<Result<Frame<Box<dyn Buf + Send>>, Box<dyn Error + Send + Sync>>>>;
12
13    fn is_end_stream(&self) -> bool;
14
15    fn size_hint(&self) -> SizeHint;
16}
17
18/// 为所有满足条件的 Body 实现 ErasedBody
19impl<B> ErasedBody for B
20where
21    B: Body + Send + Sync + 'static,
22    B::Data: Buf + Send,
23    B::Error: Into<Box<dyn Error + Send + Sync>>,
24{
25    fn poll_frame(
26        self: Pin<&mut Self>,
27        cx: &mut Context<'_>,
28    ) -> Poll<Option<Result<Frame<Box<dyn Buf + Send>>, Box<dyn Error + Send + Sync>>>> {
29        match <Self as Body>::poll_frame(self, cx) {
30            Poll::Pending => Poll::Pending,
31            Poll::Ready(None) => Poll::Ready(None),
32            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e.into()))),
33            Poll::Ready(Some(Ok(frame))) => {
34                let mapped = frame.map_data(|d| Box::new(d) as Box<dyn Buf + Send>);
35                Poll::Ready(Some(Ok(mapped)))
36            }
37        }
38    }
39
40    fn is_end_stream(&self) -> bool {
41        <Self as Body>::is_end_stream(self)
42    }
43
44    fn size_hint(&self) -> SizeHint {
45        <Self as Body>::size_hint(self)
46    }
47}
48
49/// 通用型http body
50pub struct VariantBody {
51    inner: Pin<Box<dyn ErasedBody + 'static>>,
52}
53
54impl VariantBody {
55    pub fn new<B>(body: B) -> Self
56    where
57        B: Body + Send + Sync + 'static,
58        B::Data: Buf + Send,
59        B::Error: Into<Box<dyn Error + Send + Sync>>,
60    {
61        Self {
62            inner: Box::pin(body),
63        }
64    }
65}
66
67/// 为VariantBody实现body trait
68impl Body for VariantBody {
69    type Data = Box<dyn Buf + Send>;
70    type Error = Box<dyn Error + Send + Sync>;
71
72    fn poll_frame(
73        self: Pin<&mut Self>,
74        cx: &mut Context<'_>,
75    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
76        self.get_mut().inner.as_mut().poll_frame(cx)
77    }
78
79    fn is_end_stream(&self) -> bool {
80        self.inner.is_end_stream()
81    }
82
83    fn size_hint(&self) -> SizeHint {
84        self.inner.size_hint()
85    }
86}
87
88pub fn variant_body<B>(body: B) -> VariantBody
89where
90    B: Body + Send + Sync + 'static,
91    B::Data: Buf + Send,
92    B::Error: Into<Box<dyn Error + Send + Sync>>,
93{
94    VariantBody::new(body)
95}