Skip to main content

hyper_body_utils/
lib.rs

1#![doc = include_str!("../README.md")]
2#![deny(missing_docs)]
3#[cfg(feature = "http3")]
4use bytes::Buf;
5use bytes::Bytes;
6#[cfg(all(feature = "http3", feature = "compio"))]
7use compio_quic::RecvStream as CompioRecvStream;
8#[cfg(feature = "http3")]
9use futures::ready;
10use futures::{stream, FutureExt, Stream, TryStreamExt};
11#[cfg(feature = "http3")]
12use h3::{
13    client::RequestStream as ClientRequestStream, server::RequestStream as ServerRequestStream,
14};
15#[cfg(all(feature = "http3", feature = "generic"))]
16use h3_quinn::RecvStream as GenericRecvStream;
17use http_body_util::{combinators::BoxBody, StreamBody};
18use hyper::body::{Body, Frame, Incoming};
19use std::{
20    pin::Pin,
21    task::{Context, Poll},
22};
23
24pub use http_body_util::BodyExt;
25
26#[cfg(test)]
27mod tests;
28
29/// Enum to represent different types of HTTP bodies
30pub enum HttpBody {
31    /// Incoming body from hyper
32    Incoming(Incoming),
33    /// Boxed stream body
34    BoxedStream(BoxBody<Bytes, std::io::Error>),
35    /// QUIC client incoming stream
36    #[cfg(all(feature = "http3", feature = "generic"))]
37    GenericClient(ClientRequestStream<GenericRecvStream, Bytes>),
38    /// QUIC server incoming stream
39    #[cfg(all(feature = "http3", feature = "generic"))]
40    GenericServer(ServerRequestStream<GenericRecvStream, Bytes>),
41    /// QUIC client incoming stream
42    #[cfg(all(feature = "http3", feature = "compio"))]
43    CompioClient(ClientRequestStream<CompioRecvStream, Bytes>),
44    /// QUIC server incoming stream
45    #[cfg(all(feature = "http3", feature = "compio"))]
46    CompioServer(ServerRequestStream<CompioRecvStream, Bytes>),
47}
48
49impl HttpBody {
50    /// Create a new HttpBody from an Incoming body
51    pub fn from_incoming(incoming: Incoming) -> Self {
52        HttpBody::Incoming(incoming)
53    }
54
55    /// Create a new HttpBody from a QUIC client stream
56    #[cfg(all(feature = "http3", feature = "generic"))]
57    pub fn from_generic_client(stream: ClientRequestStream<GenericRecvStream, Bytes>) -> Self {
58        HttpBody::GenericClient(stream)
59    }
60
61    /// Create a new HttpBody from a QUIC server stream
62    #[cfg(all(feature = "http3", feature = "generic"))]
63    pub fn from_generic_server(stream: ServerRequestStream<GenericRecvStream, Bytes>) -> Self {
64        HttpBody::GenericServer(stream)
65    }
66
67    /// Create a new HttpBody from a QUIC client stream
68    #[cfg(all(feature = "http3", feature = "compio"))]
69    pub fn from_compio_client(stream: ClientRequestStream<CompioRecvStream, Bytes>) -> Self {
70        HttpBody::CompioClient(stream)
71    }
72
73    /// Create a new HttpBody from a QUIC server stream
74    #[cfg(all(feature = "http3", feature = "compio"))]
75    pub fn from_compio_server(stream: ServerRequestStream<CompioRecvStream, Bytes>) -> Self {
76        HttpBody::CompioServer(stream)
77    }
78
79    /// Create a new HttpBody from a text string
80    pub fn from_text(text: &str) -> Self {
81        Self::from_bytes(text.as_bytes())
82    }
83
84    /// Create a new empty HttpBody
85    pub fn empty() -> Self {
86        Self::from_bytes(&Bytes::new())
87    }
88
89    /// Create a new HttpBody from a stream
90    pub fn from_stream<S>(stream: S) -> Self
91    where
92        S: Stream<Item = Result<Frame<Bytes>, std::io::Error>> + Send + Sync + 'static,
93    {
94        let body = StreamBody::new(stream);
95        HttpBody::BoxedStream(BodyExt::boxed(body))
96    }
97
98    /// Create a new HttpBody from bytes
99    pub fn from_bytes(bytes: &[u8]) -> Self {
100        let all_bytes = Bytes::copy_from_slice(bytes);
101        let content = stream::iter(vec![Ok(all_bytes)]).map_ok(Frame::data);
102        let body = StreamBody::new(content);
103        HttpBody::BoxedStream(BodyExt::boxed(body))
104    }
105}
106
107impl Body for HttpBody {
108    type Data = Bytes;
109
110    type Error = std::io::Error;
111
112    fn poll_frame(
113        self: Pin<&mut Self>,
114        cx: &mut Context<'_>,
115    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
116        match self.get_mut() {
117            HttpBody::Incoming(incoming) => incoming
118                .frame()
119                .poll_unpin(cx)
120                .map_err(std::io::Error::other),
121
122            HttpBody::BoxedStream(stream) => {
123                stream.frame().poll_unpin(cx).map_err(std::io::Error::other)
124            }
125
126            #[cfg(all(feature = "http3", feature = "generic"))]
127            HttpBody::GenericClient(stream) => match ready!(stream.poll_recv_data(cx)) {
128                Ok(frame) => match frame {
129                    Some(mut frame) => Poll::Ready(Some(Ok(Frame::data(
130                        frame.copy_to_bytes(frame.remaining()),
131                    )))),
132                    None => {
133                        cx.waker().wake_by_ref();
134                        Poll::Ready(None)
135                    }
136                },
137                Err(e) => {
138                    println!("Error polling frame: {}", e);
139                    Poll::Ready(Some(Err(std::io::Error::other(e))))
140                }
141            },
142
143            #[cfg(all(feature = "http3", feature = "generic"))]
144            HttpBody::GenericServer(stream) => match ready!(stream.poll_recv_data(cx)) {
145                Ok(frame) => match frame {
146                    Some(mut frame) => Poll::Ready(Some(Ok(Frame::data(
147                        frame.copy_to_bytes(frame.remaining()),
148                    )))),
149                    None => {
150                        cx.waker().wake_by_ref();
151                        Poll::Ready(None)
152                    }
153                },
154                Err(e) => Poll::Ready(Some(Err(std::io::Error::other(e)))),
155            },
156
157            #[cfg(all(feature = "http3", feature = "compio"))]
158            HttpBody::CompioClient(stream) => match ready!(stream.poll_recv_data(cx)) {
159                Ok(frame) => match frame {
160                    Some(mut frame) => Poll::Ready(Some(Ok(Frame::data(
161                        frame.copy_to_bytes(frame.remaining()),
162                    )))),
163                    None => {
164                        cx.waker().wake_by_ref();
165                        Poll::Ready(None)
166                    }
167                },
168                Err(e) => {
169                    println!("Error polling frame: {}", e);
170                    Poll::Ready(Some(Err(std::io::Error::other(e))))
171                }
172            },
173
174            #[cfg(all(feature = "http3", feature = "compio"))]
175            HttpBody::CompioServer(stream) => match ready!(stream.poll_recv_data(cx)) {
176                Ok(frame) => match frame {
177                    Some(mut frame) => Poll::Ready(Some(Ok(Frame::data(
178                        frame.copy_to_bytes(frame.remaining()),
179                    )))),
180                    None => {
181                        cx.waker().wake_by_ref();
182                        Poll::Ready(None)
183                    }
184                },
185                Err(e) => Poll::Ready(Some(Err(std::io::Error::other(e)))),
186            },
187        }
188    }
189}
190
191impl Stream for HttpBody {
192    type Item = Result<Frame<Bytes>, std::io::Error>;
193
194    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
195        self.poll_frame(cx)
196    }
197}