xref: /tonic/tonic/src/body.rs (revision 60b131d2)
1 //! HTTP specific body utilities.
2 
3 use std::{pin::Pin, task::Poll};
4 
5 use http_body_util::BodyExt as _;
6 
7 // A type erased HTTP body.
8 type BoxBody = http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, crate::Status>;
9 
10 /// A body type used in `tonic`.
11 #[derive(Debug)]
12 pub struct Body {
13     kind: Kind,
14 }
15 
16 #[derive(Debug)]
17 enum Kind {
18     Empty,
19     Wrap(BoxBody),
20 }
21 
22 impl Body {
from_kind(kind: Kind) -> Self23     fn from_kind(kind: Kind) -> Self {
24         Self { kind }
25     }
26 
27     /// Create a new empty `Body`.
empty() -> Self28     pub const fn empty() -> Self {
29         Self { kind: Kind::Empty }
30     }
31 
32     /// Create a new `Body` from an existing `Body`.
new<B>(body: B) -> Self where B: http_body::Body<Data = bytes::Bytes> + Send + 'static, B::Error: Into<crate::BoxError>,33     pub fn new<B>(body: B) -> Self
34     where
35         B: http_body::Body<Data = bytes::Bytes> + Send + 'static,
36         B::Error: Into<crate::BoxError>,
37     {
38         if body.is_end_stream() {
39             return Self::empty();
40         }
41 
42         let mut body = Some(body);
43 
44         if let Some(body) = <dyn std::any::Any>::downcast_mut::<Option<Body>>(&mut body) {
45             return body.take().unwrap();
46         }
47 
48         if let Some(body) = <dyn std::any::Any>::downcast_mut::<Option<BoxBody>>(&mut body) {
49             return Self::from_kind(Kind::Wrap(body.take().unwrap()));
50         }
51 
52         let body = body
53             .unwrap()
54             .map_err(crate::Status::map_error)
55             .boxed_unsync();
56 
57         Self::from_kind(Kind::Wrap(body))
58     }
59 }
60 
61 impl Default for Body {
default() -> Self62     fn default() -> Self {
63         Self::empty()
64     }
65 }
66 
67 impl http_body::Body for Body {
68     type Data = bytes::Bytes;
69     type Error = crate::Status;
70 
poll_frame( mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>>71     fn poll_frame(
72         mut self: std::pin::Pin<&mut Self>,
73         cx: &mut std::task::Context<'_>,
74     ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
75         match &mut self.kind {
76             Kind::Empty => Poll::Ready(None),
77             Kind::Wrap(body) => Pin::new(body).poll_frame(cx),
78         }
79     }
80 
size_hint(&self) -> http_body::SizeHint81     fn size_hint(&self) -> http_body::SizeHint {
82         match &self.kind {
83             Kind::Empty => http_body::SizeHint::with_exact(0),
84             Kind::Wrap(body) => body.size_hint(),
85         }
86     }
87 
is_end_stream(&self) -> bool88     fn is_end_stream(&self) -> bool {
89         match &self.kind {
90             Kind::Empty => true,
91             Kind::Wrap(body) => body.is_end_stream(),
92         }
93     }
94 }
95