feat(body): make Body know about incoming Content-Length

When getting a `Body` from hyper, such as in a client response,
the method `Body::content_length()` now returns a value if the header
was present.

Closes #1545
This commit is contained in:
Laurențiu Nicola
2018-06-08 22:00:46 +02:00
committed by Sean McArthur
parent 396fe80e76
commit a0a0fcdd9b
5 changed files with 108 additions and 20 deletions

View File

@@ -114,7 +114,7 @@ where I: AsyncRead + AsyncWrite,
read_buf.len() >= 24 && read_buf[..24] == *H2_PREFACE
}
pub fn read_head(&mut self) -> Poll<Option<(MessageHead<T::Incoming>, bool)>, ::Error> {
pub fn read_head(&mut self) -> Poll<Option<(MessageHead<T::Incoming>, Option<BodyLength>)>, ::Error> {
debug_assert!(self.can_read_head());
trace!("Conn::read_head");
@@ -162,7 +162,6 @@ where I: AsyncRead + AsyncWrite,
continue;
}
};
debug!("incoming body is {}", decoder);
self.state.busy();
@@ -172,20 +171,23 @@ where I: AsyncRead + AsyncWrite,
}
let wants_keep_alive = msg.keep_alive;
self.state.keep_alive &= wants_keep_alive;
let (body, reading) = if decoder.is_eof() {
(false, Reading::KeepAlive)
} else {
(true, Reading::Body(decoder))
};
let content_length = decoder.content_length();
if let Reading::Closed = self.state.reading {
// actually want an `if not let ...`
} else {
self.state.reading = reading;
self.state.reading = if content_length.is_none() {
Reading::KeepAlive
} else {
Reading::Body(decoder)
};
}
if !body {
if content_length.is_none() {
self.try_keep_alive();
}
return Ok(Async::Ready(Some((head, body))));
return Ok(Async::Ready(Some((head, content_length))));
}
}

View File

@@ -7,6 +7,7 @@ use futures::{Async, Poll};
use bytes::Bytes;
use super::io::MemRead;
use super::BodyLength;
use self::Kind::{Length, Chunked, Eof};
@@ -84,6 +85,16 @@ impl Decoder {
}
}
pub fn content_length(&self) -> Option<BodyLength> {
match self.kind {
Length(0) |
Chunked(ChunkedState::End, _) |
Eof(true) => None,
Length(len) => Some(BodyLength::Known(len)),
_ => Some(BodyLength::Unknown),
}
}
pub fn decode<R: MemRead>(&mut self, body: &mut R) -> Poll<Bytes, io::Error> {
trace!("decode; state={:?}", self.kind);
match self.kind {

View File

@@ -190,9 +190,14 @@ where
}
// dispatch is ready for a message, try to read one
match self.conn.read_head() {
Ok(Async::Ready(Some((head, has_body)))) => {
let body = if has_body {
let (mut tx, rx) = Body::channel();
Ok(Async::Ready(Some((head, body_len)))) => {
let body = if let Some(body_len) = body_len {
let (mut tx, rx) =
Body::new_channel(if let BodyLength::Known(len) = body_len {
Some(len)
} else {
None
});
let _ = tx.poll_ready(); // register this task if rx is dropped
self.body_tx = Some(tx);
rx
@@ -201,7 +206,7 @@ where
};
self.dispatch.recv_msg(Ok((head, body)))?;
Ok(Async::Ready(()))
},
}
Ok(Async::Ready(None)) => {
// read eof, conn will start to shutdown automatically
Ok(Async::Ready(()))