feat(http1): add support for receiving trailer fields

Closes hyperium#2703
hjr3 · Apr 17, 2024 · b730382 · b730382
1 parent c78379e
commit b730382
Show file tree

Hide file tree

Showing 6 changed files with 467 additions and 67 deletions.
diff --git a/src/body/incoming.rs b/src/body/incoming.rs
@@ -403,6 +403,19 @@ impl Sender {
             .map_err(|err| err.into_inner().expect("just sent Ok"))
     }
 
+    #[cfg(feature = "http1")]
+    pub(crate) fn try_send_trailers(
+        &mut self,
+        trailers: HeaderMap,
+    ) -> Result<(), Option<HeaderMap>> {
+        let tx = match self.trailers_tx.take() {
+            Some(tx) => tx,
+            None => return Err(None),
+        };
+
+        tx.send(trailers).map_err(|err| Some(err))
+    }
+
     #[cfg(test)]
     pub(crate) fn abort(mut self) {
         self.send_error(crate::Error::new_body_write_aborted());

diff --git a/src/proto/h1/conn.rs b/src/proto/h1/conn.rs
@@ -11,6 +11,7 @@ use bytes::{Buf, Bytes};
 use futures_util::ready;
 use http::header::{HeaderValue, CONNECTION, TE};
 use http::{HeaderMap, Method, Version};
+use http_body::Frame;
 use httparse::ParserConfig;
 
 use super::io::Buffered;
@@ -311,33 +312,41 @@ where
     pub(crate) fn poll_read_body(
         &mut self,
         cx: &mut Context<'_>,
-    ) -> Poll<Option<io::Result<Bytes>>> {
+    ) -> Poll<Option<io::Result<Frame<Bytes>>>> {
         debug_assert!(self.can_read_body());
 
         let (reading, ret) = match self.state.reading {
             Reading::Body(ref mut decoder) => {
                 match ready!(decoder.decode(cx, &mut self.io)) {
-                    Ok(slice) => {
-                        let (reading, chunk) = if decoder.is_eof() {
-                            debug!("incoming body completed");
-                            (
-                                Reading::KeepAlive,
-                                if !slice.is_empty() {
-                                    Some(Ok(slice))
-                                } else {
-                                    None
-                                },
-                            )
-                        } else if slice.is_empty() {
-                            error!("incoming body unexpectedly ended");
-                            // This should be unreachable, since all 3 decoders
-                            // either set eof=true or return an Err when reading
-                            // an empty slice...
-                            (Reading::Closed, None)
+                    Ok(frame) => {
+                        if frame.is_data() {
+                            let slice = frame.data_ref().unwrap_or_else(|| unreachable!());
+                            let (reading, maybe_frame) = if decoder.is_eof() {
+                                debug!("incoming body completed");
+                                (
+                                    Reading::KeepAlive,
+                                    if !slice.is_empty() {
+                                        Some(Ok(frame))
+                                    } else {
+                                        None
+                                    },
+                                )
+                            } else if slice.is_empty() {
+                                error!("incoming body unexpectedly ended");
+                                // This should be unreachable, since all 3 decoders
+                                // either set eof=true or return an Err when reading
+                                // an empty slice...
+                                (Reading::Closed, None)
+                            } else {
+                                return Poll::Ready(Some(Ok(frame)));
+                            };
+                            (reading, Poll::Ready(maybe_frame))
+                        } else if frame.is_trailers() {
+                            (Reading::Closed, Poll::Ready(Some(Ok(frame))))
                         } else {
-                            return Poll::Ready(Some(Ok(slice)));
-                        };
-                        (reading, Poll::Ready(chunk))
+                            trace!("discarding unknown frame");
+                            (Reading::Closed, Poll::Ready(None))
+                        }
                     }
                     Err(e) => {
                         debug!("incoming body decode error: {}", e);