Merge pull request #645 from hyperium/640-invalid-response

fix(client): be resilient to invalid response bodies

src/client/mod.rs

@@ -446,8 +446,10 @@ fn get_host_and_port(url: &Url) -> ::Result<(String, u16)> {
 
 #[cfg(test)]
 mod tests {
+    use std::io::Read;
     use header::Server;
     use super::{Client, RedirectPolicy};
+    use super::pool::Pool;
     use url::Url;
 
     mock_connector!(MockRedirectPolicy {
@@ -494,4 +496,31 @@ mod tests {
         let res = client.get("http://127.0.0.1").send().unwrap();
         assert_eq!(res.headers.get(), Some(&Server("mock2".to_owned())));
     }
+
+    mock_connector!(Issue640Connector {
+        b"HTTP/1.1 200 OK\r\nContent-Length: 3\r\n\r\n",
+        b"GET",
+        b"HTTP/1.1 200 OK\r\nContent-Length: 4\r\n\r\n",
+        b"HEAD",
+        b"HTTP/1.1 200 OK\r\nContent-Length: 4\r\n\r\n",
+        b"POST"
+    });
+
+    // see issue #640
+    #[test]
+    fn test_head_response_body_keep_alive() {
+        let client = Client::with_connector(Pool::with_connector(Default::default(), Issue640Connector));
+
+        let mut s = String::new();
+        client.get("http://127.0.0.1").send().unwrap().read_to_string(&mut s).unwrap();
+        assert_eq!(s, "GET");
+
+        let mut s = String::new();
+        client.head("http://127.0.0.1").send().unwrap().read_to_string(&mut s).unwrap();
+        assert_eq!(s, "");
+
+        let mut s = String::new();
+        client.post("http://127.0.0.1").send().unwrap().read_to_string(&mut s).unwrap();
+        assert_eq!(s, "POST");
+    }
 }

src/client/pool.rs

@@ -34,7 +34,7 @@ impl Default for Config {
 
 #[derive(Debug)]
 struct PoolImpl<S> {
-    conns: HashMap<Key, Vec<S>>,
+    conns: HashMap<Key, Vec<PooledStreamInner<S>>>,
     config: Config,
 }
 
@@ -90,7 +90,7 @@ impl<C: NetworkConnector> Pool<C> {
 }
 
 impl<S> PoolImpl<S> {
-    fn reuse(&mut self, key: Key, conn: S) {
+    fn reuse(&mut self, key: Key, conn: PooledStreamInner<S>) {
         trace!("reuse {:?}", key);
         let conns = self.conns.entry(key).or_insert(vec![]);
         if conns.len() < self.config.max_idle {
@@ -105,73 +105,97 @@ impl<C: NetworkConnector<Stream=S>, S: NetworkStream + Send> NetworkConnector fo
         let key = key(host, port, scheme);
         let mut locked = self.inner.lock().unwrap();
         let mut should_remove = false;
-        let conn = match locked.conns.get_mut(&key) {
+        let inner = match locked.conns.get_mut(&key) {
             Some(ref mut vec) => {
                 trace!("Pool had connection, using");
                 should_remove = vec.len() == 1;
                 vec.pop().unwrap()
             }
-            _ => try!(self.connector.connect(host, port, scheme))
+            _ => PooledStreamInner {
+                key: key.clone(),
+                stream: try!(self.connector.connect(host, port, scheme)),
+                previous_response_expected_no_content: false,
+            }
         };
         if should_remove {
             locked.conns.remove(&key);
         }
         Ok(PooledStream {
-            inner: Some((key, conn)),
+            inner: Some(inner),
             is_closed: false,
-            pool: self.inner.clone()
+            pool: self.inner.clone(),
         })
     }
 }
 
 /// A Stream that will try to be returned to the Pool when dropped.
 pub struct PooledStream<S> {
-    inner: Option<(Key, S)>,
+    inner: Option<PooledStreamInner<S>>,
     is_closed: bool,
-    pool: Arc<Mutex<PoolImpl<S>>>
+    pool: Arc<Mutex<PoolImpl<S>>>,
+}
+
+#[derive(Debug)]
+struct PooledStreamInner<S> {
+    key: Key,
+    stream: S,
+    previous_response_expected_no_content: bool,
 }
 
 impl<S: NetworkStream> Read for PooledStream<S> {
     #[inline]
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        self.inner.as_mut().unwrap().1.read(buf)
+        self.inner.as_mut().unwrap().stream.read(buf)
     }
 }
 
 impl<S: NetworkStream> Write for PooledStream<S> {
     #[inline]
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        self.inner.as_mut().unwrap().1.write(buf)
+        self.inner.as_mut().unwrap().stream.write(buf)
     }
 
     #[inline]
     fn flush(&mut self) -> io::Result<()> {
-        self.inner.as_mut().unwrap().1.flush()
+        self.inner.as_mut().unwrap().stream.flush()
     }
 }
 
 impl<S: NetworkStream> NetworkStream for PooledStream<S> {
     #[inline]
     fn peer_addr(&mut self) -> io::Result<SocketAddr> {
-        self.inner.as_mut().unwrap().1.peer_addr()
+        self.inner.as_mut().unwrap().stream.peer_addr()
     }
 
     #[cfg(feature = "timeouts")]
     #[inline]
     fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
-        self.inner.as_ref().unwrap().1.set_read_timeout(dur)
+        self.inner.as_ref().unwrap().stream.set_read_timeout(dur)
     }
 
     #[cfg(feature = "timeouts")]
     #[inline]
     fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
-        self.inner.as_ref().unwrap().1.set_write_timeout(dur)
+        self.inner.as_ref().unwrap().stream.set_write_timeout(dur)
     }
 
     #[inline]
     fn close(&mut self, how: Shutdown) -> io::Result<()> {
         self.is_closed = true;
-        self.inner.as_mut().unwrap().1.close(how)
+        self.inner.as_mut().unwrap().stream.close(how)
+    }
+
+    #[inline]
+    fn set_previous_response_expected_no_content(&mut self, expected: bool) {
+        trace!("set_previous_response_expected_no_content {}", expected);
+        self.inner.as_mut().unwrap().previous_response_expected_no_content = expected;
+    }
+
+    #[inline]
+    fn previous_response_expected_no_content(&self) -> bool {
+        let answer = self.inner.as_ref().unwrap().previous_response_expected_no_content;
+        trace!("previous_response_expected_no_content {}", answer);
+        answer
     }
 }
 
@@ -179,9 +203,9 @@ impl<S> Drop for PooledStream<S> {
     fn drop(&mut self) {
         trace!("PooledStream.drop, is_closed={}", self.is_closed);
         if !self.is_closed {
-            self.inner.take().map(|(key, conn)| {
+            self.inner.take().map(|inner| {
                 if let Ok(mut pool) = self.pool.lock() {
-                    pool.reuse(key, conn);
+                    pool.reuse(inner.key.clone(), inner);
                 }
                 // else poisoned, give up
             });

src/client/response.rs

@@ -64,7 +64,6 @@ impl Response {
     pub fn status_raw(&self) -> &RawStatus {
         &self.status_raw
     }
-
 }
 
 impl Read for Response {
@@ -91,11 +90,11 @@ impl Drop for Response {
         //
         // otherwise, the response has been drained. we should check that the
         // server has agreed to keep the connection open
-        trace!("Response.is_drained = {:?}", self.is_drained);
+        trace!("Response.drop is_drained={}", self.is_drained);
         if !(self.is_drained && http::should_keep_alive(self.version, &self.headers)) {
-            trace!("closing connection");
+            trace!("Response.drop closing connection");
             if let Err(e) = self.message.close_connection() {
-                error!("error closing connection: {}", e);
+                error!("Response.drop error closing connection: {}", e);
             }
         }
     }

src/http/h1.rs

@@ -33,6 +33,8 @@ use http::{
 use header;
 use version;
 
+const MAX_INVALID_RESPONSE_BYTES: usize = 1024 * 128;
+
 /// An implementation of the `HttpMessage` trait for HTTP/1.1.
 #[derive(Debug)]
 pub struct Http11Message {
@@ -169,19 +171,38 @@ impl HttpMessage for Http11Message {
             }
         };
 
+        let expected_no_content = stream.previous_response_expected_no_content();
+        trace!("previous_response_expected_no_content = {}", expected_no_content);
+
         let mut stream = BufReader::new(stream);
 
-        let head = match parse_response(&mut stream) {
-            Ok(head) => head,
-            Err(e) => {
-                self.stream = Some(stream.into_inner());
-                return Err(e);
-            }
-        };
+        let mut invalid_bytes_read = 0;
+        let head;
+        loop {
+            head = match parse_response(&mut stream) {
+                Ok(head) => head,
+                Err(::Error::Version)
+                if expected_no_content && invalid_bytes_read < MAX_INVALID_RESPONSE_BYTES => {
+                    trace!("expected_no_content, found content");
+                    invalid_bytes_read += 1;
+                    stream.consume(1);
+                    continue;
+                }
+                Err(e) => {
+                    self.stream = Some(stream.into_inner());
+                    return Err(e);
+                }
+            };
+            break;
+        }
+
         let raw_status = head.subject;
         let headers = head.headers;
 
         let method = self.method.take().unwrap_or(Method::Get);
+
+        let is_empty = !should_have_response_body(&method, raw_status.0);
+        stream.get_mut().set_previous_response_expected_no_content(is_empty);
         // According to https://tools.ietf.org/html/rfc7230#section-3.3.3
         // 1. HEAD reponses, and Status 1xx, 204, and 304 cannot have a body.
         // 2. Status 2xx to a CONNECT cannot have a body.
@@ -190,27 +211,24 @@ impl HttpMessage for Http11Message {
         // 5. Content-Length header has a sized body.
         // 6. Not Client.
         // 7. Read till EOF.
-        self.reader = Some(match (method, raw_status.0) {
-            (Method::Head, _) => EmptyReader(stream),
-            (_, 100...199) | (_, 204) | (_, 304) => EmptyReader(stream),
-            (Method::Connect, 200...299) => EmptyReader(stream),
-            _ => {
-                 if let Some(&TransferEncoding(ref codings)) = headers.get() {
-                    if codings.last() == Some(&Chunked) {
-                        ChunkedReader(stream, None)
-                    } else {
-                        trace!("not chuncked. read till eof");
-                        EofReader(stream)
-                    }
-                } else if let Some(&ContentLength(len)) =  headers.get() {
-                    SizedReader(stream, len)
-                } else if headers.has::<ContentLength>() {
-                    trace!("illegal Content-Length: {:?}", headers.get_raw("Content-Length"));
-                    return Err(Error::Header);
+        self.reader = Some(if is_empty {
+            EmptyReader(stream)
+        } else {
+             if let Some(&TransferEncoding(ref codings)) = headers.get() {
+                if codings.last() == Some(&Chunked) {
+                    ChunkedReader(stream, None)
                 } else {
-                    trace!("neither Transfer-Encoding nor Content-Length");
+                    trace!("not chuncked. read till eof");
                     EofReader(stream)
                 }
+            } else if let Some(&ContentLength(len)) =  headers.get() {
+                SizedReader(stream, len)
+            } else if headers.has::<ContentLength>() {
+                trace!("illegal Content-Length: {:?}", headers.get_raw("Content-Length"));
+                return Err(Error::Header);
+            } else {
+                trace!("neither Transfer-Encoding nor Content-Length");
+                EofReader(stream)
             }
         });
 
@@ -226,7 +244,9 @@ impl HttpMessage for Http11Message {
 
     fn has_body(&self) -> bool {
         match self.reader {
-            Some(EmptyReader(..)) => false,
+            Some(EmptyReader(..)) |
+            Some(SizedReader(_, 0)) |
+            Some(ChunkedReader(_, Some(0))) => false,
             _ => true
         }
     }
@@ -597,6 +617,18 @@ fn read_chunk_size<R: Read>(rdr: &mut R) -> io::Result<u64> {
     Ok(size)
 }
 
+fn should_have_response_body(method: &Method, status: u16) -> bool {
+    trace!("should_have_response_body({:?}, {})", method, status);
+    match (method, status) {
+        (&Method::Head, _) |
+        (_, 100...199) |
+        (_, 204) |
+        (_, 304) |
+        (&Method::Connect, 200...299) => false,
+        _ => true
+    }
+}
+
 /// Writers to handle different Transfer-Encodings.
 pub enum HttpWriter<W: Write> {
     /// A no-op Writer, used initially before Transfer-Encoding is determined.

src/mock.rs

@@ -19,6 +19,7 @@ use net::{NetworkStream, NetworkConnector};
 #[derive(Clone, Debug)]
 pub struct MockStream {
     pub read: Cursor<Vec<u8>>,
+    next_reads: Vec<Vec<u8>>,
     pub write: Vec<u8>,
     pub is_closed: bool,
     pub error_on_write: bool,
@@ -40,21 +41,15 @@ impl MockStream {
         MockStream::with_input(b"")
     }
 
-    #[cfg(not(feature = "timeouts"))]
     pub fn with_input(input: &[u8]) -> MockStream {
-        MockStream {
-            read: Cursor::new(input.to_vec()),
-            write: vec![],
-            is_closed: false,
-            error_on_write: false,
-            error_on_read: false,
-        }
+        MockStream::with_responses(vec![input])
     }
 
     #[cfg(feature = "timeouts")]
-    pub fn with_input(input: &[u8]) -> MockStream {
+    pub fn with_responses(mut responses: Vec<&[u8]>) -> MockStream {
         MockStream {
-            read: Cursor::new(input.to_vec()),
+            read: Cursor::new(responses.remove(0).to_vec()),
+            next_reads: responses.into_iter().map(|arr| arr.to_vec()).collect(),
             write: vec![],
             is_closed: false,
             error_on_write: false,
@@ -63,6 +58,18 @@ impl MockStream {
             write_timeout: Cell::new(None),
         }
     }
+
+    #[cfg(not(feature = "timeouts"))]
+    pub fn with_responses(mut responses: Vec<&[u8]>) -> MockStream {
+        MockStream {
+            read: Cursor::new(responses.remove(0).to_vec()),
+            next_reads: responses.into_iter().map(|arr| arr.to_vec()).collect(),
+            write: vec![],
+            is_closed: false,
+            error_on_write: false,
+            error_on_read: false,
+        }
+    }
 }
 
 impl Read for MockStream {
@@ -70,7 +77,17 @@ impl Read for MockStream {
         if self.error_on_read {
             Err(io::Error::new(io::ErrorKind::Other, "mock error"))
         } else {
-            self.read.read(buf)
+            match self.read.read(buf) {
+                Ok(n) => {
+                    if self.read.position() as usize == self.read.get_ref().len() {
+                        if self.next_reads.len() > 0 {
+                            self.read = Cursor::new(self.next_reads.remove(0));
+                        }
+                    }
+                    Ok(n)
+                },
+                r => r
+            }
         }
     }
 }
@@ -191,7 +208,7 @@ macro_rules! mock_connector (
 
         struct $name;
 
-        impl ::net::NetworkConnector for $name {
+        impl $crate::net::NetworkConnector for $name {
             type Stream = ::mock::MockStream;
             fn connect(&self, host: &str, port: u16, scheme: &str)
                     -> $crate::Result<::mock::MockStream> {
@@ -210,7 +227,21 @@ macro_rules! mock_connector (
             }
         }
 
-    )
+    );
+
+    ($name:ident { $($response:expr),+ }) => (
+        struct $name;
+
+        impl $crate::net::NetworkConnector for $name {
+            type Stream = $crate::mock::MockStream;
+            fn connect(&self, _: &str, _: u16, _: &str)
+                    -> $crate::Result<$crate::mock::MockStream> {
+                Ok($crate::mock::MockStream::with_responses(vec![
+                    $($response),+
+                ]))
+            }
+        }
+    );
 );
 
 impl TransportStream for MockStream {

src/net.rs

@@ -62,6 +62,17 @@ pub trait NetworkStream: Read + Write + Any + Send + Typeable {
     fn close(&mut self, _how: Shutdown) -> io::Result<()> {
         Ok(())
     }
+
+    // Unsure about name and implementation...
+
+    #[doc(hidden)]
+    fn set_previous_response_expected_no_content(&mut self, _expected: bool) {
+    
+    }
+    #[doc(hidden)]
+    fn previous_response_expected_no_content(&self) -> bool {
+        false
+    }
 }
 
 /// A connector creates a NetworkStream.