Skip to main content

mechanics_http_server/
body.rs

1//! HTTP/3 request body adapter.
2
3use std::future::Future;
4use std::pin::Pin;
5use std::task::{Context, Poll};
6
7use bytes::{Buf, Bytes};
8use h3::server::RequestStream;
9use http::HeaderMap;
10use http_body::Frame;
11
12type H3RecvStream = RequestStream<h3_quinn::RecvStream, Bytes>;
13type RecvDataFuture =
14    Pin<Box<dyn Future<Output = (H3RecvStream, Result<Option<Bytes>, H3RequestBodyError>)> + Send>>;
15type RecvTrailersFuture = Pin<
16    Box<dyn Future<Output = (H3RecvStream, Result<Option<HeaderMap>, H3RequestBodyError>)> + Send>,
17>;
18
19/// Streaming request body received from an HTTP/3 client.
20///
21/// `H3RequestBody` is the request-body type passed to services by
22/// [`crate::Http3Server::start`]. It implements [`http_body::Body`]
23/// and yields HTTP/3 DATA frames as [`Bytes`].
24pub struct H3RequestBody {
25    stream: Option<H3RecvStream>,
26    state: H3RequestBodyState,
27}
28
29enum H3RequestBodyState {
30    ReadingData,
31    DataFuture(RecvDataFuture),
32    ReadingTrailers,
33    TrailersFuture(RecvTrailersFuture),
34    Done,
35}
36
37impl H3RequestBody {
38    pub(crate) fn new(stream: H3RecvStream) -> Self {
39        Self {
40            stream: Some(stream),
41            state: H3RequestBodyState::ReadingData,
42        }
43    }
44}
45
46impl std::fmt::Debug for H3RequestBody {
47    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        f.debug_struct("H3RequestBody")
49            .field(
50                "is_end_stream",
51                &matches!(self.state, H3RequestBodyState::Done),
52            )
53            .finish_non_exhaustive()
54    }
55}
56
57/// Error returned while receiving an HTTP/3 request body.
58#[derive(Debug, thiserror::Error)]
59pub enum H3RequestBodyError {
60    /// The underlying HTTP/3 stream returned an error while receiving
61    /// data or trailers.
62    #[error("HTTP/3 request body failed: {0}")]
63    Stream(String),
64}
65
66impl H3RequestBodyError {
67    fn stream(error: h3::error::StreamError) -> Self {
68        Self::Stream(error.to_string())
69    }
70}
71
72impl http_body::Body for H3RequestBody {
73    type Data = Bytes;
74    type Error = H3RequestBodyError;
75
76    fn poll_frame(
77        mut self: Pin<&mut Self>,
78        cx: &mut Context<'_>,
79    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
80        loop {
81            match &mut self.state {
82                H3RequestBodyState::ReadingData => {
83                    let Some(mut stream) = self.stream.take() else {
84                        self.state = H3RequestBodyState::Done;
85                        return Poll::Ready(None);
86                    };
87                    self.state = H3RequestBodyState::DataFuture(Box::pin(async move {
88                        let result = stream.recv_data().await.map_err(H3RequestBodyError::stream);
89                        let result = result.map(|chunk| {
90                            chunk.map(|mut chunk| {
91                                let remaining = chunk.remaining();
92                                chunk.copy_to_bytes(remaining)
93                            })
94                        });
95                        (stream, result)
96                    }));
97                }
98                H3RequestBodyState::DataFuture(future) => {
99                    let (stream, result) = match future.as_mut().poll(cx) {
100                        Poll::Pending => return Poll::Pending,
101                        Poll::Ready(result) => result,
102                    };
103                    self.stream = Some(stream);
104                    match result {
105                        Ok(Some(bytes)) => {
106                            self.state = H3RequestBodyState::ReadingData;
107                            return Poll::Ready(Some(Ok(Frame::data(bytes))));
108                        }
109                        Ok(None) => {
110                            self.state = H3RequestBodyState::ReadingTrailers;
111                        }
112                        Err(error) => {
113                            self.state = H3RequestBodyState::Done;
114                            return Poll::Ready(Some(Err(error)));
115                        }
116                    }
117                }
118                H3RequestBodyState::ReadingTrailers => {
119                    let Some(mut stream) = self.stream.take() else {
120                        self.state = H3RequestBodyState::Done;
121                        return Poll::Ready(None);
122                    };
123                    self.state = H3RequestBodyState::TrailersFuture(Box::pin(async move {
124                        let result = stream
125                            .recv_trailers()
126                            .await
127                            .map_err(H3RequestBodyError::stream);
128                        (stream, result)
129                    }));
130                }
131                H3RequestBodyState::TrailersFuture(future) => {
132                    let (stream, result) = match future.as_mut().poll(cx) {
133                        Poll::Pending => return Poll::Pending,
134                        Poll::Ready(result) => result,
135                    };
136                    self.stream = Some(stream);
137                    self.state = H3RequestBodyState::Done;
138                    return Poll::Ready(match result {
139                        Ok(Some(trailers)) => Some(Ok(Frame::trailers(trailers))),
140                        Ok(None) => None,
141                        Err(error) => Some(Err(error)),
142                    });
143                }
144                H3RequestBodyState::Done => return Poll::Ready(None),
145            }
146        }
147    }
148
149    fn is_end_stream(&self) -> bool {
150        matches!(self.state, H3RequestBodyState::Done)
151    }
152}