tonic/
body.rs

1//! HTTP specific body utilities.
2
3use std::{pin::Pin, task::Poll};
4
5use http_body_util::BodyExt as _;
6
7// A type erased HTTP body.
8type BoxBody = http_body_util::combinators::UnsyncBoxBody<bytes::Bytes, crate::Status>;
9
10/// A body type used in `tonic`.
11#[derive(Debug)]
12pub struct Body {
13    kind: Kind,
14}
15
16#[derive(Debug)]
17enum Kind {
18    Empty,
19    Wrap(BoxBody),
20}
21
22impl Body {
23    fn from_kind(kind: Kind) -> Self {
24        Self { kind }
25    }
26
27    /// Create a new empty `Body`.
28    pub const fn empty() -> Self {
29        Self { kind: Kind::Empty }
30    }
31
32    /// Create a new `Body` from an existing `Body`.
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
61impl Default for Body {
62    fn default() -> Self {
63        Self::empty()
64    }
65}
66
67impl http_body::Body for Body {
68    type Data = bytes::Bytes;
69    type Error = crate::Status;
70
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
81    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
88    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}