Skip to main content

hyper_body_utils/
lib.rs

1#![doc = include_str!("../README.md")]
2#![deny(missing_docs)]
3use bytes::Buf;
4use futures::ready;
5
6use bytes::Bytes;
7use futures::{stream, FutureExt, Stream, TryStreamExt};
8use std::pin::Pin;
9use std::task::{Context, Poll};
10
11use h3::client::RequestStream as ClientRequestStream;
12use h3::server::RequestStream as ServerRequestStream;
13use h3_quinn::RecvStream;
14
15pub use http_body_util::BodyExt;
16
17use http_body_util::StreamBody;
18use hyper::body::{Body, Frame, Incoming};
19
20use http_body_util::combinators::BoxBody;
21
22#[cfg(test)]
23mod tests;
24
25/// Enum to represent different types of HTTP bodies
26pub enum HttpBody {
27    /// Incoming body from hyper
28    Incoming(Incoming),
29    /// Boxed stream body
30    BoxedStream(BoxBody<Bytes, std::io::Error>),
31    /// QUIC client incoming stream
32    QuicClientIncoming(ClientRequestStream<RecvStream, Bytes>),
33    /// QUIC server incoming stream
34    QuicServerIncoming(ServerRequestStream<RecvStream, Bytes>),
35}
36
37impl HttpBody {
38    /// Create a new HttpBody from an Incoming body
39    pub fn from_incoming(incoming: Incoming) -> Self {
40        HttpBody::Incoming(incoming)
41    }
42
43    /// Create a new HttpBody from a QUIC client stream
44    pub fn from_quic_client(stream: ClientRequestStream<RecvStream, Bytes>) -> Self {
45        HttpBody::QuicClientIncoming(stream)
46    }
47
48    /// Create a new HttpBody from a QUIC server stream
49    pub fn from_quic_server(stream: ServerRequestStream<RecvStream, Bytes>) -> Self {
50        HttpBody::QuicServerIncoming(stream)
51    }
52
53    /// Create a new HttpBody from a text string
54    pub fn from_text(text: &str) -> Self {
55        Self::from_bytes(text.as_bytes())
56    }
57
58    /// Create a new empty HttpBody
59    pub fn empty() -> Self {
60        Self::from_bytes(&Bytes::new())
61    }    
62
63    /// Create a new HttpBody from a stream
64    pub fn from_stream<S>(stream: S) -> Self
65    where
66        S: Stream<Item = Result<Frame<Bytes>, std::io::Error>> + Send + Sync + 'static,
67    {
68        let body = StreamBody::new(stream);
69        HttpBody::BoxedStream(BodyExt::boxed(body))
70    }
71
72    /// Create a new HttpBody from bytes
73    pub fn from_bytes(bytes: &[u8]) -> Self {
74        let all_bytes = Bytes::copy_from_slice(bytes);
75        let content = stream::iter(vec![Ok(all_bytes)]).map_ok(Frame::data);
76        let body = StreamBody::new(content);
77        HttpBody::BoxedStream(BodyExt::boxed(body))
78    }
79}
80
81impl Body for HttpBody {
82    type Data = Bytes;
83
84    type Error = std::io::Error;
85
86    fn poll_frame(
87        self: Pin<&mut Self>,
88        cx: &mut Context<'_>,
89    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
90        match self.get_mut() {
91            HttpBody::Incoming(incoming) => incoming
92                .frame()
93                .poll_unpin(cx)
94                .map_err(std::io::Error::other),
95
96            HttpBody::BoxedStream(stream) => {
97                stream.frame().poll_unpin(cx).map_err(std::io::Error::other)
98            }
99
100            HttpBody::QuicClientIncoming(stream) => match ready!(stream.poll_recv_data(cx)) {
101                Ok(frame) => match frame {
102                    Some(mut frame) => Poll::Ready(Some(Ok(Frame::data(
103                        frame.copy_to_bytes(frame.remaining()),
104                    )))),
105                    None => {
106                        cx.waker().wake_by_ref();
107                        Poll::Ready(None)
108                    }
109                },
110                Err(e) => {
111                    println!("Error polling frame: {}", e);
112                    Poll::Ready(Some(Err(std::io::Error::other(e))))
113                }
114            },
115            HttpBody::QuicServerIncoming(stream) => match ready!(stream.poll_recv_data(cx)) {
116                Ok(frame) => match frame {
117                    Some(mut frame) => Poll::Ready(Some(Ok(Frame::data(
118                        frame.copy_to_bytes(frame.remaining()),
119                    )))),
120                    None => {
121                        cx.waker().wake_by_ref();
122                        Poll::Ready(None)
123                    }
124                },
125                Err(e) => Poll::Ready(Some(Err(std::io::Error::other(e)))),
126            },
127        }
128    }
129}
130
131impl Stream for HttpBody {
132    type Item = Result<Frame<Bytes>, std::io::Error>;
133
134    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
135        self.poll_frame(cx)
136    }
137}