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(feature = "http3")]
7use futures::ready;
8use futures::{stream, FutureExt, Stream, TryStreamExt};
9#[cfg(feature = "http3")]
10use h3::client::RequestStream as ClientRequestStream;
11#[cfg(feature = "http3")]
12use h3::server::RequestStream as ServerRequestStream;
13#[cfg(all(feature = "http3", not(feature = "compio")))]
14use h3_quinn::RecvStream;
15#[cfg(all(feature = "http3", feature = "compio"))]
16use compio_quic::RecvStream;
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(feature = "http3")]
37    QuicClientIncoming(ClientRequestStream<RecvStream, Bytes>),
38    /// QUIC server incoming stream
39    #[cfg(feature = "http3")]
40    QuicServerIncoming(ServerRequestStream<RecvStream, Bytes>),
41}
42
43impl HttpBody {
44    /// Create a new HttpBody from an Incoming body
45    pub fn from_incoming(incoming: Incoming) -> Self {
46        HttpBody::Incoming(incoming)
47    }
48
49    /// Create a new HttpBody from a QUIC client stream
50    #[cfg(feature = "http3")]
51    pub fn from_quic_client(stream: ClientRequestStream<RecvStream, Bytes>) -> Self {
52        HttpBody::QuicClientIncoming(stream)
53    }
54
55    /// Create a new HttpBody from a QUIC server stream
56    #[cfg(feature = "http3")]
57    pub fn from_quic_server(stream: ServerRequestStream<RecvStream, Bytes>) -> Self {
58        HttpBody::QuicServerIncoming(stream)
59    }
60
61    /// Create a new HttpBody from a text string
62    pub fn from_text(text: &str) -> Self {
63        Self::from_bytes(text.as_bytes())
64    }
65
66    /// Create a new empty HttpBody
67    pub fn empty() -> Self {
68        Self::from_bytes(&Bytes::new())
69    }
70
71    /// Create a new HttpBody from a stream
72    pub fn from_stream<S>(stream: S) -> Self
73    where
74        S: Stream<Item = Result<Frame<Bytes>, std::io::Error>> + Send + Sync + 'static,
75    {
76        let body = StreamBody::new(stream);
77        HttpBody::BoxedStream(BodyExt::boxed(body))
78    }
79
80    /// Create a new HttpBody from bytes
81    pub fn from_bytes(bytes: &[u8]) -> Self {
82        let all_bytes = Bytes::copy_from_slice(bytes);
83        let content = stream::iter(vec![Ok(all_bytes)]).map_ok(Frame::data);
84        let body = StreamBody::new(content);
85        HttpBody::BoxedStream(BodyExt::boxed(body))
86    }
87}
88
89impl Body for HttpBody {
90    type Data = Bytes;
91
92    type Error = std::io::Error;
93
94    fn poll_frame(
95        self: Pin<&mut Self>,
96        cx: &mut Context<'_>,
97    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
98        match self.get_mut() {
99            HttpBody::Incoming(incoming) => incoming
100                .frame()
101                .poll_unpin(cx)
102                .map_err(std::io::Error::other),
103
104            HttpBody::BoxedStream(stream) => {
105                stream.frame().poll_unpin(cx).map_err(std::io::Error::other)
106            }
107
108            #[cfg(feature = "http3")]
109            HttpBody::QuicClientIncoming(stream) => match ready!(stream.poll_recv_data(cx)) {
110                Ok(frame) => match frame {
111                    Some(mut frame) => Poll::Ready(Some(Ok(Frame::data(
112                        frame.copy_to_bytes(frame.remaining()),
113                    )))),
114                    None => {
115                        cx.waker().wake_by_ref();
116                        Poll::Ready(None)
117                    }
118                },
119                Err(e) => {
120                    println!("Error polling frame: {}", e);
121                    Poll::Ready(Some(Err(std::io::Error::other(e))))
122                }
123            },
124
125            #[cfg(feature = "http3")]
126            HttpBody::QuicServerIncoming(stream) => match ready!(stream.poll_recv_data(cx)) {
127                Ok(frame) => match frame {
128                    Some(mut frame) => Poll::Ready(Some(Ok(Frame::data(
129                        frame.copy_to_bytes(frame.remaining()),
130                    )))),
131                    None => {
132                        cx.waker().wake_by_ref();
133                        Poll::Ready(None)
134                    }
135                },
136                Err(e) => Poll::Ready(Some(Err(std::io::Error::other(e)))),
137            },
138        }
139    }
140}
141
142impl Stream for HttpBody {
143    type Item = Result<Frame<Bytes>, std::io::Error>;
144
145    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
146        self.poll_frame(cx)
147    }
148}