Skip to main content

leyline/core/
body.rs

1use std::fmt;
2use std::io;
3use std::pin::Pin;
4use std::task::{Context, Poll};
5
6use bytes::Bytes;
7use futures_util::Stream;
8
9pub(crate) type BoxedStream = Pin<Box<dyn Stream<Item = io::Result<Bytes>> + Send + 'static>>;
10
11#[derive(Default)]
12pub(crate) enum BodyKind {
13    #[default]
14    Empty,
15    Bytes(Bytes),
16    Stream {
17        stream: BoxedStream,
18        length_hint: Option<u64>,
19    },
20}
21
22#[derive(Default)]
23pub struct Body(pub(crate) BodyKind);
24
25impl fmt::Debug for Body {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        let mut out = f.debug_struct("Body");
28        match &self.0 {
29            BodyKind::Empty => out.field("kind", &"empty"),
30            BodyKind::Bytes(b) => out.field("kind", &"bytes").field("len", &b.len()),
31            BodyKind::Stream { length_hint, .. } => out
32                .field("kind", &"stream")
33                .field("length_hint", length_hint),
34        }
35        .finish()
36    }
37}
38
39impl Stream for Body {
40    type Item = io::Result<Bytes>;
41
42    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
43        match &mut self.0 {
44            BodyKind::Empty => Poll::Ready(None),
45            BodyKind::Bytes(_) => match std::mem::take(&mut self.0) {
46                BodyKind::Bytes(b) => Poll::Ready(Some(Ok(b))),
47                _ => Poll::Ready(None),
48            },
49            BodyKind::Stream { stream, .. } => stream.as_mut().poll_next(cx),
50        }
51    }
52}
53
54impl Body {
55    pub(crate) fn bytes(bytes: Bytes) -> Self {
56        Body(BodyKind::Bytes(bytes))
57    }
58
59    pub fn stream<S>(stream: S, length: Option<u64>) -> Self
60    where
61        S: Stream<Item = io::Result<Bytes>> + Send + 'static,
62    {
63        Body(BodyKind::Stream {
64            stream: Box::pin(stream),
65            length_hint: length,
66        })
67    }
68
69    pub fn len_hint(&self) -> Option<u64> {
70        match &self.0 {
71            BodyKind::Empty => Some(0),
72            BodyKind::Bytes(b) => Some(b.len() as u64),
73            BodyKind::Stream { length_hint, .. } => *length_hint,
74        }
75    }
76
77    #[cfg(feature = "http3")]
78    pub(crate) fn is_stream(&self) -> bool {
79        matches!(self.0, BodyKind::Stream { .. })
80    }
81
82    #[cfg(feature = "http3")]
83    pub(crate) fn into_parts(self) -> (Option<Bytes>, Option<BoxedStream>) {
84        match self.0 {
85            BodyKind::Empty => (None, None),
86            BodyKind::Bytes(b) => (Some(b), None),
87            BodyKind::Stream { stream, .. } => (None, Some(stream)),
88        }
89    }
90
91    pub(crate) fn into_h2(self) -> crate::h2::client::RequestBody {
92        use crate::h2::client::RequestBody;
93        match self.0 {
94            BodyKind::Empty => RequestBody::None,
95            BodyKind::Bytes(b) => RequestBody::Buffered(b),
96            BodyKind::Stream {
97                stream,
98                length_hint,
99            } => RequestBody::Streaming {
100                stream,
101                length_hint,
102            },
103        }
104    }
105
106    pub(crate) fn replay(&self) -> Option<Body> {
107        match &self.0 {
108            BodyKind::Empty => Some(Body::default()),
109            BodyKind::Bytes(b) => Some(Body::bytes(b.clone())),
110            BodyKind::Stream { .. } => None,
111        }
112    }
113}
114
115impl From<Bytes> for Body {
116    fn from(b: Bytes) -> Self {
117        if b.is_empty() {
118            Body::default()
119        } else {
120            Body::bytes(b)
121        }
122    }
123}
124
125impl From<Vec<u8>> for Body {
126    fn from(v: Vec<u8>) -> Self {
127        Body::from(Bytes::from(v))
128    }
129}
130
131impl From<&'static [u8]> for Body {
132    fn from(v: &'static [u8]) -> Self {
133        Body::from(Bytes::from_static(v))
134    }
135}
136
137impl From<String> for Body {
138    fn from(s: String) -> Self {
139        Body::from(Bytes::from(s.into_bytes()))
140    }
141}
142
143impl From<&'static str> for Body {
144    fn from(s: &'static str) -> Self {
145        Body::from(Bytes::from_static(s.as_bytes()))
146    }
147}
148
149impl From<()> for Body {
150    fn from(_: ()) -> Self {
151        Body::default()
152    }
153}
154
155#[cfg(test)]
156mod tests;