feat(http1): Add support for sending HTTP/1.1 Chunked Trailer Fields (h…

…yperium#3375)

Closes hyperium#2719

Signed-off-by: Sven Pfennig <[email protected]>

src/proto/h1/conn.rs

@@ -8,7 +8,7 @@ use std::time::Duration;
 
 use crate::rt::{Read, Write};
 use bytes::{Buf, Bytes};
-use http::header::{HeaderValue, CONNECTION};
+use http::header::{HeaderValue, CONNECTION, TE};
 use http::{HeaderMap, Method, Version};
 use httparse::ParserConfig;
 
@@ -75,6 +75,7 @@ where
                 // We assume a modern world where the remote speaks HTTP/1.1.
                 // If they tell us otherwise, we'll downgrade in `read_head`.
                 version: Version::HTTP_11,
+                allow_trailer_fields: false,
             },
             _marker: PhantomData,
         }
@@ -264,6 +265,13 @@ where
             self.state.reading = Reading::Body(Decoder::new(msg.decode));
         }
 
+        self.state.allow_trailer_fields = msg
+            .head
+            .headers
+            .get(TE)
+            .map(|te_header| te_header == "trailers")
+            .unwrap_or(false);
+
         Poll::Ready(Some(Ok((msg.head, msg.decode, wants))))
     }
 
@@ -640,6 +648,31 @@ where
         self.state.writing = state;
     }
 
+    pub(crate) fn write_trailers(&mut self, trailers: HeaderMap) {
+        if T::is_server() && self.state.allow_trailer_fields == false {
+            debug!("trailers not allowed to be sent");
+            return;
+        }
+        debug_assert!(self.can_write_body() && self.can_buffer_body());
+
+        match self.state.writing {
+            Writing::Body(ref encoder) => {
+                if let Some(enc_buf) =
+                    encoder.encode_trailers(trailers, self.state.title_case_headers)
+                {
+                    self.io.buffer(enc_buf);
+
+                    self.state.writing = if encoder.is_last() || encoder.is_close_delimited() {
+                        Writing::Closed
+                    } else {
+                        Writing::KeepAlive
+                    };
+                }
+            }
+            _ => unreachable!("write_trailers invalid state: {:?}", self.state.writing),
+        }
+    }
+
     pub(crate) fn write_body_and_end(&mut self, chunk: B) {
         debug_assert!(self.can_write_body() && self.can_buffer_body());
         // empty chunks should be discarded at Dispatcher level
@@ -842,6 +875,8 @@ struct State {
     upgrade: Option<crate::upgrade::Pending>,
     /// Either HTTP/1.0 or 1.1 connection
     version: Version,
+    /// Flag to track if trailer fields are allowed to be sent
+    allow_trailer_fields: bool,
 }
 
 #[derive(Debug)]

src/proto/h1/dispatch.rs

@@ -351,27 +351,33 @@ where
                             *clear_body = true;
                             crate::Error::new_user_body(e)
                         })?;
-                        let chunk = if let Ok(data) = frame.into_data() {
-                            data
-                        } else {
-                            trace!("discarding non-data frame");
-                            continue;
-                        };
-                        let eos = body.is_end_stream();
-                        if eos {
-                            *clear_body = true;
-                            if chunk.remaining() == 0 {
-                                trace!("discarding empty chunk");
-                                self.conn.end_body()?;
+
+                        if frame.is_data() {
+                            let chunk = frame.into_data().unwrap_or_else(|_| unreachable!());
+                            let eos = body.is_end_stream();
+                            if eos {
+                                *clear_body = true;
+                                if chunk.remaining() == 0 {
+                                    trace!("discarding empty chunk");
+                                    self.conn.end_body()?;
+                                } else {
+                                    self.conn.write_body_and_end(chunk);
+                                }
                             } else {
-                                self.conn.write_body_and_end(chunk);
+                                if chunk.remaining() == 0 {
+                                    trace!("discarding empty chunk");
+                                    continue;
+                                }
+                                self.conn.write_body(chunk);
                             }
+                        } else if frame.is_trailers() {
+                            *clear_body = true;
+                            self.conn.write_trailers(
+                                frame.into_trailers().unwrap_or_else(|_| unreachable!()),
+                            );
                         } else {
-                            if chunk.remaining() == 0 {
-                                trace!("discarding empty chunk");
-                                continue;
-                            }
-                            self.conn.write_body(chunk);
+                            trace!("discarding unknown frame");
+                            continue;
                         }
                     } else {
                         *clear_body = true;