From 4ca0914f9a41d778a7302dc289fb94497628438d Mon Sep 17 00:00:00 2001
From: Martin Algesten <martin@algesten.se>
Date: Mon, 11 Jun 2018 21:59:37 +0200
Subject: [PATCH] content limited and bogus chunk

---
 src/response.rs       | 44 +++++++++++++++++++++++++++++++++++++++++--
 src/test/body_read.rs | 42 +++++++++++++++++++++++++++++++++++++++++
 src/test/mod.rs       |  1 +
 3 files changed, 85 insertions(+), 2 deletions(-)
 create mode 100644 src/test/body_read.rs

diff --git a/src/response.rs b/src/response.rs
index 046fb32..af90bcc 100644
--- a/src/response.rs
+++ b/src/response.rs
@@ -132,12 +132,18 @@ impl Response {
 
     pub fn into_reader(self) -> impl Read {
         let is_chunked = self.get("transfer-encoding")
-            .map(|enc| enc.eq_ignore_ascii_case("chunked"))
+            .map(|enc| enc.len() > 0) // whatever it says, do chunked
             .unwrap_or(false);
+        let len = self.get("content-length").and_then(|l| l.parse::<usize>().ok());
         let reader = self.reader.expect("No reader in response?!");
         match is_chunked {
             true => Box::new(chunked_transfer::Decoder::new(reader)),
-            false => reader,
+            false => {
+                match len {
+                    Some(len) => Box::new(LimitedRead::new(reader, len)),
+                    None => reader,
+                }
+            },
         }
     }
 
@@ -273,3 +279,37 @@ fn read_next_line<R: Read>(reader: &mut R) -> IoResult<AsciiString> {
         buf.push(byte);
     }
 }
+
+struct LimitedRead {
+    reader: Box<Read + Send>,
+    limit: usize,
+    position: usize,
+}
+
+impl LimitedRead {
+    fn new(reader: Box<Read + Send>, limit: usize) -> Self {
+        LimitedRead {
+            reader,
+            limit,
+            position: 0,
+        }
+    }
+}
+
+impl Read for LimitedRead {
+    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
+        let left = self.limit - self.position;
+        let from = if left < buf.len() {
+            &mut buf[0..left]
+        } else {
+            buf
+        };
+        match self.reader.read(from) {
+            Ok(amount) => {
+                self.position += amount;
+                Ok(amount)
+            },
+            Err(e) => Err(e)
+        }
+    }
+}
diff --git a/src/test/body_read.rs b/src/test/body_read.rs
new file mode 100644
index 0000000..2f6848b
--- /dev/null
+++ b/src/test/body_read.rs
@@ -0,0 +1,42 @@
+use std::io::Read;
+use test;
+
+use super::super::*;
+
+#[test]
+fn transfer_encoding_bogus() {
+    test::set_handler("/transfer_encoding_bogus", |_req, _url| {
+        test::make_stream(
+            200,
+            "OK",
+            vec![
+                "transfer-encoding: bogus", // whatever it says here, we should chunk
+            ],
+            "3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n"
+                .to_string()
+                .into_bytes(),
+        )
+    });
+    let resp = get("test://host/transfer_encoding_bogus").call();
+    let mut reader = resp.into_reader();
+    let mut text = String::new();
+    reader.read_to_string(&mut text).unwrap();
+    assert_eq!(text, "hello world!!!");
+}
+
+#[test]
+fn content_length_limited() {
+    test::set_handler("/content_length_limited", |_req, _url| {
+        test::make_stream(
+            200,
+            "OK",
+            vec!["Content-Length: 4"],
+            "abcdefgh".to_string().into_bytes(),
+        )
+    });
+    let resp = get("test://host/content_length_limited").call();
+    let mut reader = resp.into_reader();
+    let mut text = String::new();
+    reader.read_to_string(&mut text).unwrap();
+    assert_eq!(text, "abcd");
+}
diff --git a/src/test/mod.rs b/src/test/mod.rs
index e5a73ea..1c0f425 100644
--- a/src/test/mod.rs
+++ b/src/test/mod.rs
@@ -9,6 +9,7 @@ use url::Url;
 use util::vecread::VecRead;
 
 mod simple;
+mod body_read;
 
 type RequestHandler = Fn(&Request, &Url) -> Result<Stream, Error> + Send + 'static;