Support writing continuation frames. (#198)

Large header sets might require being split up across multiple frames.
This patch adds support for doing so.
This commit is contained in:
Carl Lerche
2017-12-20 17:24:29 -08:00
committed by GitHub
parent a89401dd91
commit fc75311fae
4 changed files with 294 additions and 94 deletions

View File

@@ -56,11 +56,7 @@ pub struct Continuation {
/// Stream ID of continuation frame
stream_id: StreamId,
/// Argument to pass to the HPACK encoder to resume encoding
hpack: hpack::EncodeState,
/// remaining headers to encode
headers: Iter,
header_block: EncodingHeaderBlock,
}
// TODO: These fields shouldn't be `pub`
@@ -85,7 +81,7 @@ pub struct Iter {
fields: header::IntoIter<HeaderValue>,
}
#[derive(PartialEq, Eq)]
#[derive(Debug, PartialEq, Eq)]
struct HeaderBlock {
/// The decoded header fields
fields: HeaderMap,
@@ -95,6 +91,15 @@ struct HeaderBlock {
pseudo: Pseudo,
}
#[derive(Debug)]
struct EncodingHeaderBlock {
/// Argument to pass to the HPACK encoder to resume encoding
hpack: Option<hpack::EncodeState>,
/// remaining headers to encode
headers: Iter,
}
const END_STREAM: u8 = 0x1;
const END_HEADERS: u8 = 0x4;
const PADDED: u8 = 0x8;
@@ -200,6 +205,10 @@ impl Headers {
self.flags.is_end_headers()
}
pub fn set_end_headers(&mut self) {
self.flags.set_end_headers();
}
pub fn is_end_stream(&self) -> bool {
self.flags.is_end_stream()
}
@@ -226,21 +235,15 @@ impl Headers {
}
pub fn encode(self, encoder: &mut hpack::Encoder, dst: &mut BytesMut) -> Option<Continuation> {
// At this point, the `is_end_headers` flag should always be set
debug_assert!(self.flags.is_end_headers());
// Get the HEADERS frame head
let head = self.head();
let pos = dst.len();
// At this point, we don't know how big the h2 frame will be.
// So, we write the head with length 0, then write the body, and
// finally write the length once we know the size.
head.encode(0, dst);
// Encode the frame
let (len, cont) = self.header_block.encode(self.stream_id, encoder, dst);
// Write the frame length
BigEndian::write_uint(&mut dst[pos..pos + 3], len, 3);
cont
self.header_block.into_encoding()
.encode(&head, encoder, dst, |_| {
})
}
fn head(&self) -> Head {
@@ -325,25 +328,23 @@ impl PushPromise {
self.flags.is_end_headers()
}
pub fn set_end_headers(&mut self) {
self.flags.set_end_headers();
}
pub fn encode(self, encoder: &mut hpack::Encoder, dst: &mut BytesMut) -> Option<Continuation> {
use bytes::BufMut;
// At this point, the `is_end_headers` flag should always be set
debug_assert!(self.flags.is_end_headers());
let head = self.head();
let pos = dst.len();
let promised_id = self.promised_id;
// At this point, we don't know how big the h2 frame will be.
// So, we write the head with length 0, then write the body, and
// finally write the length once we know the size.
head.encode(0, dst);
// Encode the frame
dst.put_u32::<BigEndian>(self.promised_id.into());
let (len, cont) = self.header_block.encode(self.stream_id, encoder, dst);
// Write the frame length
BigEndian::write_uint(&mut dst[pos..pos + 3], len + 4, 3);
cont
self.header_block.into_encoding()
.encode(&head, encoder, dst, |dst| {
dst.put_u32::<BigEndian>(promised_id.into());
})
}
fn head(&self) -> Head {
@@ -400,6 +401,23 @@ impl fmt::Debug for PushPromise {
}
}
// ===== impl Continuation =====
impl Continuation {
fn head(&self) -> Head {
Head::new(Kind::Continuation, END_HEADERS, self.stream_id)
}
pub fn encode(self, encoder: &mut hpack::Encoder, dst: &mut BytesMut) -> Option<Continuation> {
// Get the CONTINUATION frame head
let head = self.head();
self.header_block
.encode(&head, encoder, dst, |_| {
})
}
}
// ===== impl Pseudo =====
impl Pseudo {
@@ -458,6 +476,58 @@ fn to_string(src: Bytes) -> String<Bytes> {
unsafe { String::from_utf8_unchecked(src) }
}
// ===== impl EncodingHeaderBlock =====
impl EncodingHeaderBlock {
fn encode<F>(mut self,
head: &Head,
encoder: &mut hpack::Encoder,
dst: &mut BytesMut,
f: F)
-> Option<Continuation>
where F: FnOnce(&mut BytesMut),
{
let head_pos = dst.len();
// At this point, we don't know how big the h2 frame will be.
// So, we write the head with length 0, then write the body, and
// finally write the length once we know the size.
head.encode(0, dst);
let payload_pos = dst.len();
f(dst);
// Now, encode the header payload
let continuation = match encoder.encode(self.hpack, &mut self.headers, dst) {
hpack::Encode::Full => None,
hpack::Encode::Partial(state) => Some(Continuation {
stream_id: head.stream_id(),
header_block: EncodingHeaderBlock {
hpack: Some(state),
headers: self.headers,
},
}),
};
// Compute the header block length
let payload_len = (dst.len() - payload_pos) as u64;
// Write the frame length
BigEndian::write_uint(&mut dst[head_pos..head_pos + 3], payload_len, 3);
if continuation.is_some() {
// There will be continuation frames, so the `is_end_headers` flag
// must be unset
debug_assert!(dst[head_pos + 4] & END_HEADERS == END_HEADERS);
dst[head_pos + 4] -= END_HEADERS;
}
continuation
}
}
// ===== impl Iter =====
impl Iterator for Iter {
@@ -515,13 +585,17 @@ impl HeadersFlag {
}
pub fn set_end_stream(&mut self) {
self.0 |= END_STREAM
self.0 |= END_STREAM;
}
pub fn is_end_headers(&self) -> bool {
self.0 & END_HEADERS == END_HEADERS
}
pub fn set_end_headers(&mut self) {
self.0 |= END_HEADERS;
}
pub fn is_padded(&self) -> bool {
self.0 & PADDED == PADDED
}
@@ -570,6 +644,10 @@ impl PushPromiseFlag {
self.0 & END_HEADERS == END_HEADERS
}
pub fn set_end_headers(&mut self) {
self.0 |= END_HEADERS;
}
pub fn is_padded(&self) -> bool {
self.0 & PADDED == PADDED
}
@@ -624,7 +702,10 @@ impl HeaderBlock {
// contain the entire payload. Later, we need to check for stream
// priority.
//
// TODO: Provide a way to abort decoding if an error is hit.
// If the header frame is malformed, we still have to continue decoding
// the headers. A malformed header frame is a stream level error, but
// the hpack state is connection level. In order to maintain correct
// state for other streams, the hpack decoding process must complete.
let res = decoder.decode(&mut src, |header| {
use hpack::Header::*;
@@ -673,30 +754,13 @@ impl HeaderBlock {
Ok(())
}
fn encode(
self,
stream_id: StreamId,
encoder: &mut hpack::Encoder,
dst: &mut BytesMut,
) -> (u64, Option<Continuation>) {
let pos = dst.len();
let mut headers = Iter {
pseudo: Some(self.pseudo),
fields: self.fields.into_iter(),
};
let cont = match encoder.encode(None, &mut headers, dst) {
hpack::Encode::Full => None,
hpack::Encode::Partial(state) => Some(Continuation {
stream_id: stream_id,
hpack: state,
headers: headers,
}),
};
// Compute the header block length
let len = (dst.len() - pos) as u64;
(len, cont)
fn into_encoding(self) -> EncodingHeaderBlock {
EncodingHeaderBlock {
hpack: None,
headers: Iter {
pseudo: Some(self.pseudo),
fields: self.fields.into_iter(),
},
}
}
}