diff --git a/src/common/buf.rs b/src/common/buf.rs new file mode 100644 index 00000000..5cd5a173 --- /dev/null +++ b/src/common/buf.rs @@ -0,0 +1,34 @@ +use bytes::Buf; +use iovec::IoVec; + +/// A `Buf` wrapping a static byte slice. +#[derive(Debug)] +pub(crate) struct StaticBuf(pub(crate) &'static [u8]); + +impl Buf for StaticBuf { + #[inline] + fn remaining(&self) -> usize { + self.0.len() + } + + #[inline] + fn bytes(&self) -> &[u8] { + self.0 + } + + #[inline] + fn advance(&mut self, cnt: usize) { + self.0 = &self.0[cnt..]; + } + + #[inline] + fn bytes_vec<'t>(&'t self, dst: &mut [&'t IoVec]) -> usize { + if dst.is_empty() { + return 0; + } else { + dst[0] = self.0.into(); + return 1; + } + } +} + diff --git a/src/common/mod.rs b/src/common/mod.rs index cd8ea4d5..124b2100 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,5 +1,7 @@ +mod buf; mod exec; mod never; +pub(crate) use self::buf::StaticBuf; pub(crate) use self::exec::Exec; pub use self::never::Never; diff --git a/src/proto/h1/conn.rs b/src/proto/h1/conn.rs index a4124fd6..04459d8b 100644 --- a/src/proto/h1/conn.rs +++ b/src/proto/h1/conn.rs @@ -482,20 +482,13 @@ where I: AsyncRead + AsyncWrite, debug_assert!(chunk.remaining() != 0); let state = match self.state.writing { - Writing::Body(ref mut encoder) => { - self.io.buffer(encoder.encode(chunk)); - match encoder.end() { - Ok(end) => { - if let Some(end) = end { - self.io.buffer(end); - } - if encoder.is_last() { - Writing::Closed - } else { - Writing::KeepAlive - } - }, - Err(_not_eof) => Writing::Closed, + Writing::Body(ref encoder) => { + let (encoded, can_keep_alive) = encoder.encode_and_end(chunk); + self.io.buffer(encoded); + if can_keep_alive { + Writing::KeepAlive + } else { + Writing::Closed } }, _ => unreachable!("write_body invalid state: {:?}", self.state.writing), diff --git a/src/proto/h1/encode.rs b/src/proto/h1/encode.rs index 2072f3a5..0f47561d 100644 --- a/src/proto/h1/encode.rs +++ b/src/proto/h1/encode.rs @@ -1,10 +1,11 @@ -//use std::cmp; use std::fmt; use bytes::{Buf, IntoBuf}; use bytes::buf::{Chain, Take}; use iovec::IoVec; +use common::StaticBuf; + /// Encoders to handle different Transfer-Encodings. #[derive(Debug, Clone)] pub struct Encoder { @@ -39,8 +40,8 @@ enum Kind { enum BufKind { Exact(B), Limited(Take), - Chunked(Chain>), - ChunkedEnd(CrLf), + Chunked(Chain, StaticBuf>), + ChunkedEnd(StaticBuf), } impl Encoder { @@ -81,7 +82,7 @@ impl Encoder { match self.kind { Kind::Length(0) => Ok(None), Kind::Chunked => Ok(Some(EncodedBuf { - kind: BufKind::ChunkedEnd(CrLf(b"0\r\n\r\n")), + kind: BufKind::ChunkedEnd(StaticBuf(b"0\r\n\r\n")), })), _ => Err(NotEof), } @@ -93,13 +94,15 @@ impl Encoder { { let msg = msg.into_buf(); let len = msg.remaining(); - assert!(len > 0, "encode() called with empty buf"); + debug_assert!(len > 0, "encode() called with empty buf"); - let buf = match self.kind { + let kind = match self.kind { Kind::Chunked => { trace!("encoding chunked {}B", len); - BufKind::Chunked(ChunkSize::new(len) - .chain(msg.chain(CrLf(b"\r\n")))) + let buf = ChunkSize::new(len) + .chain(msg) + .chain(StaticBuf(b"\r\n")); + BufKind::Chunked(buf) }, Kind::Length(ref mut remaining) => { trace!("sized write, len = {}", len); @@ -118,9 +121,52 @@ impl Encoder { } }; EncodedBuf { - kind: buf, + kind, } } + + pub fn encode_and_end(&self, msg: B) -> (EncodedBuf, bool) + where + B: IntoBuf, + { + let msg = msg.into_buf(); + let len = msg.remaining(); + debug_assert!(len > 0, "encode() called with empty buf"); + + let (kind, eof) = match self.kind { + Kind::Chunked => { + trace!("encoding chunked {}B", len); + let buf = ChunkSize::new(len) + .chain(msg) + .chain(StaticBuf(b"\r\n0\r\n\r\n")); + (BufKind::Chunked(buf), !self.is_last) + }, + Kind::Length(remaining) => { + use std::cmp::Ordering; + + trace!("sized write, len = {}", len); + match (len as u64).cmp(&remaining) { + Ordering::Greater => { + (BufKind::Limited(msg.take(remaining as usize)), !self.is_last) + }, + Ordering::Equal => { + (BufKind::Exact(msg), !self.is_last) + }, + Ordering::Less => { + (BufKind::Exact(msg), false) + } + } + }, + Kind::CloseDelimited => { + trace!("close delimited write {}B", len); + (BufKind::Exact(msg), false) + } + }; + + (EncodedBuf { + kind, + }, eof) + } } impl Buf for EncodedBuf @@ -236,36 +282,6 @@ impl fmt::Write for ChunkSize { } } -#[derive(Debug)] -struct CrLf(&'static [u8]); - -impl Buf for CrLf { - #[inline] - fn remaining(&self) -> usize { - self.0.len() - } - - #[inline] - fn bytes(&self) -> &[u8] { - self.0 - } - - #[inline] - fn advance(&mut self, cnt: usize) { - self.0 = &self.0[cnt..]; - } - - #[inline] - fn bytes_vec<'t>(&'t self, dst: &mut [&'t IoVec]) -> usize { - if dst.is_empty() { - return 0; - } else { - dst[0] = self.0.into(); - return 1; - } - } -} - #[cfg(test)] mod tests { use bytes::{BufMut};