Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(http1): Add support for writing Trailer Fields #3375

Merged
merged 5 commits into from
Dec 15, 2023
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 41 additions & 1 deletion src/proto/h1/conn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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,
}
Expand Down Expand Up @@ -264,6 +265,18 @@ where
self.state.reading = Reading::Body(Decoder::new(msg.decode));
}

if msg
.head
.headers
.get(TE)
.map(|te_header| te_header == "trailers")
.unwrap_or(false)
hjr3 marked this conversation as resolved.
Show resolved Hide resolved
{
self.state.allow_trailer_fields = true;
} else {
self.state.allow_trailer_fields = false;
}

Poll::Ready(Some(Ok((msg.head, msg.decode, wants))))
}

Expand Down Expand Up @@ -640,6 +653,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
Expand Down Expand Up @@ -842,6 +880,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)]
Expand Down
42 changes: 24 additions & 18 deletions src/proto/h1/dispatch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading
Loading