scuffle_http/
body.rs

1//! Types for working with HTTP bodies.
2
3use std::fmt::Debug;
4use std::pin::Pin;
5use std::task::{Context, Poll};
6
7use bytes::{Buf, Bytes};
8use http_body::Frame;
9
10/// An error that can occur when reading the body of an incoming request.
11#[derive(thiserror::Error, Debug)]
12pub enum IncomingBodyError {
13    /// An error that occurred while reading a hyper body.
14    #[error("hyper error: {0}")]
15    #[cfg(any(feature = "http1", feature = "http2"))]
16    Hyper(#[from] hyper::Error),
17    /// An error that occurred while reading a h3 body.
18    #[error("h3 body error: {0}")]
19    #[cfg(feature = "http3")]
20    H3(#[from] crate::backend::h3::body::H3BodyError),
21}
22
23/// The body of an incoming request.
24///
25/// This enum is used to abstract away the differences between the body types of HTTP/1, HTTP/2 and HTTP/3.
26/// It implements the [`http_body::Body`] trait.
27pub enum IncomingBody {
28    /// The body of an incoming hyper request.
29    #[cfg(any(feature = "http1", feature = "http2"))]
30    Hyper(hyper::body::Incoming),
31    /// The body of an incoming h3 request.
32    #[cfg(feature = "http3")]
33    Quic(crate::backend::h3::body::QuicIncomingBody<h3_quinn::RecvStream>),
34}
35
36#[cfg(any(feature = "http1", feature = "http2"))]
37impl From<hyper::body::Incoming> for IncomingBody {
38    fn from(body: hyper::body::Incoming) -> Self {
39        IncomingBody::Hyper(body)
40    }
41}
42
43#[cfg(feature = "http3")]
44impl From<crate::backend::h3::body::QuicIncomingBody<h3_quinn::RecvStream>> for IncomingBody {
45    fn from(body: crate::backend::h3::body::QuicIncomingBody<h3_quinn::RecvStream>) -> Self {
46        IncomingBody::Quic(body)
47    }
48}
49
50impl http_body::Body for IncomingBody {
51    type Data = Bytes;
52    type Error = IncomingBodyError;
53
54    fn is_end_stream(&self) -> bool {
55        match self {
56            #[cfg(any(feature = "http1", feature = "http2"))]
57            IncomingBody::Hyper(body) => body.is_end_stream(),
58            #[cfg(feature = "http3")]
59            IncomingBody::Quic(body) => body.is_end_stream(),
60            #[cfg(not(any(feature = "http1", feature = "http2", feature = "http3")))]
61            _ => false,
62        }
63    }
64
65    fn poll_frame(
66        self: std::pin::Pin<&mut Self>,
67        _cx: &mut std::task::Context<'_>,
68    ) -> std::task::Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
69        match self.get_mut() {
70            #[cfg(any(feature = "http1", feature = "http2"))]
71            IncomingBody::Hyper(body) => std::pin::Pin::new(body).poll_frame(_cx).map_err(Into::into),
72            #[cfg(feature = "http3")]
73            IncomingBody::Quic(body) => std::pin::Pin::new(body).poll_frame(_cx).map_err(Into::into),
74            #[cfg(not(any(feature = "http1", feature = "http2", feature = "http3")))]
75            _ => std::task::Poll::Ready(None),
76        }
77    }
78
79    fn size_hint(&self) -> http_body::SizeHint {
80        match self {
81            #[cfg(any(feature = "http1", feature = "http2"))]
82            IncomingBody::Hyper(body) => body.size_hint(),
83            #[cfg(feature = "http3")]
84            IncomingBody::Quic(body) => body.size_hint(),
85            #[cfg(not(any(feature = "http1", feature = "http2", feature = "http3")))]
86            _ => http_body::SizeHint::default(),
87        }
88    }
89}
90
91pin_project_lite::pin_project! {
92    /// A wrapper around an HTTP body that tracks the size of the data that is read from it.
93    pub struct TrackedBody<B, T> {
94        #[pin]
95        body: B,
96        tracker: T,
97    }
98}
99
100impl<B, T> TrackedBody<B, T> {
101    /// Create a new [`TrackedBody`] with the given body and tracker.
102    pub fn new(body: B, tracker: T) -> Self {
103        Self { body, tracker }
104    }
105}
106
107/// An error that can occur when tracking the body of an incoming request.
108#[derive(thiserror::Error)]
109pub enum TrackedBodyError<B, T>
110where
111    B: http_body::Body,
112    T: Tracker,
113{
114    /// An error that occurred while reading the body.
115    #[error("body error: {0}")]
116    Body(B::Error),
117    /// An error that occurred while calling [`Tracker::on_data`].
118    #[error("tracker error: {0}")]
119    Tracker(T::Error),
120}
121
122impl<B, T> Debug for TrackedBodyError<B, T>
123where
124    B: http_body::Body,
125    B::Error: Debug,
126    T: Tracker,
127    T::Error: Debug,
128{
129    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
130        match self {
131            TrackedBodyError::Body(err) => f.debug_tuple("TrackedBodyError::Body").field(err).finish(),
132            TrackedBodyError::Tracker(err) => f.debug_tuple("TrackedBodyError::Tracker").field(err).finish(),
133        }
134    }
135}
136
137/// A trait for tracking the size of the data that is read from an HTTP body.
138pub trait Tracker: Send + Sync + 'static {
139    /// The error type that can occur when [`Tracker::on_data`] is called.
140    type Error;
141
142    /// Called when data is read from the body.
143    ///
144    /// The `size` parameter is the size of the data that is remaining to be read from the body.
145    fn on_data(&self, size: usize) -> Result<(), Self::Error> {
146        let _ = size;
147        Ok(())
148    }
149}
150
151impl<B, T> http_body::Body for TrackedBody<B, T>
152where
153    B: http_body::Body,
154    T: Tracker,
155{
156    type Data = B::Data;
157    type Error = TrackedBodyError<B, T>;
158
159    fn poll_frame(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
160        let this = self.project();
161
162        match this.body.poll_frame(cx) {
163            Poll::Pending => Poll::Pending,
164            Poll::Ready(frame) => {
165                if let Some(Ok(frame)) = &frame {
166                    if let Some(data) = frame.data_ref() {
167                        if let Err(err) = this.tracker.on_data(data.remaining()) {
168                            return Poll::Ready(Some(Err(TrackedBodyError::Tracker(err))));
169                        }
170                    }
171                }
172
173                Poll::Ready(frame.transpose().map_err(TrackedBodyError::Body).transpose())
174            }
175        }
176    }
177
178    fn is_end_stream(&self) -> bool {
179        self.body.is_end_stream()
180    }
181
182    fn size_hint(&self) -> http_body::SizeHint {
183        self.body.size_hint()
184    }
185}
186
187#[cfg(test)]
188#[cfg_attr(all(test, coverage_nightly), coverage(off))]
189mod tests {
190    use std::convert::Infallible;
191
192    use crate::body::TrackedBodyError;
193
194    #[test]
195    fn tracked_body_error_debug() {
196        struct TestTracker;
197
198        impl super::Tracker for TestTracker {
199            type Error = Infallible;
200        }
201
202        struct TestBody;
203
204        impl http_body::Body for TestBody {
205            type Data = bytes::Bytes;
206            type Error = ();
207
208            fn poll_frame(
209                self: std::pin::Pin<&mut Self>,
210                _cx: &mut std::task::Context<'_>,
211            ) -> std::task::Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
212                std::task::Poll::Ready(None)
213            }
214        }
215
216        let err = TrackedBodyError::<TestBody, TestTracker>::Body(());
217        assert_eq!(format!("{err:?}"), "TrackedBodyError::Body(())",);
218    }
219}
scuffle_http/body.rs

scuffle_http/
body.rs
