handle chunked decoding in the client and fix some issues with it

Author: dignifiedquire

Cargo.toml

@@ -20,6 +20,7 @@ pin-project-lite = "0.1.1"
 byte-pool = "0.2.1"
 lazy_static = "1.4.0"
 futures-core = "0.3.1"
+log = "0.4"
 
 [dev-dependencies]
 pretty_assertions = "0.6.1"

src/chunked.rs

@@ -59,7 +59,7 @@ impl<R: Read> ChunkedDecoder<R> {
 
 fn decode_init(buffer: Block<'static>, pos: &Position) -> io::Result<DecodeResult> {
     dbg!(pos);
-    dbg!(std::str::from_utf8(&buffer[pos.start..pos.end]).unwrap());
+    dbg!(std::str::from_utf8(&buffer[pos.start..pos.end]));
 
     use httparse::Status;
     match httparse::parse_chunk_size(&buffer[pos.start..pos.end]) {
@@ -91,7 +91,7 @@ fn decode_init(buffer: Block<'static>, pos: &Position) -> io::Result<DecodeResul
 
 fn decode_chunk_end(buffer: Block<'static>, pos: &Position) -> io::Result<DecodeResult> {
     dbg!(pos);
-    dbg!(std::str::from_utf8(&buffer[pos.start..pos.end]).unwrap());
+    dbg!(std::str::from_utf8(&buffer[pos.start..pos.end]));
 
     if pos.len() < 2 {
         return Ok(DecodeResult::None(buffer));
@@ -120,7 +120,7 @@ fn decode_trailer(buffer: Block<'static>, pos: &Position) -> io::Result<DecodeRe
     // read headers
     let mut headers = [httparse::EMPTY_HEADER; 16];
     dbg!(pos);
-    dbg!(std::str::from_utf8(&buffer[pos.start..pos.end]).unwrap());
+    dbg!(std::str::from_utf8(&buffer[pos.start..pos.end]));
 
     match httparse::parse_headers(&buffer[pos.start..pos.end], &mut headers) {
         Ok(Status::Complete((used, headers))) => {
@@ -162,61 +162,100 @@ impl<R: Read + Unpin> ChunkedDecoder<R> {
         current: u64,
         len: u64,
     ) -> io::Result<DecodeResult> {
+        dbg!("poll_read_chunk");
         dbg!(pos);
-        dbg!(std::str::from_utf8(&buffer[pos.start..pos.end]).unwrap());
+        dbg!(std::str::from_utf8(&buffer[pos.start..pos.end]));
+        dbg!(len);
+        dbg!(current);
 
         let mut new_pos = pos.clone();
+        let remaining = (len - current) as usize;
+        let mut to_read = std::cmp::min(remaining, buf.len());
+
+        let mut new_current = current;
+        let mut start = 0;
+
+        // first drain the buffer
+        if new_pos.len() > 0 {
+            let to_read_buf = std::cmp::min(to_read, pos.len());
+            dbg!(to_read_buf);
+            dbg!(to_read);
+            dbg!(new_pos);
+            dbg!(std::str::from_utf8(
+                &buffer[new_pos.start..new_pos.start + to_read_buf]
+            ));
+            buf[..to_read_buf].copy_from_slice(&buffer[new_pos.start..new_pos.start + to_read_buf]);
+
+            // to_read -= to_read_buf;
+            new_pos.start += to_read_buf;
+            new_current += to_read_buf as u64;
+            // start += to_read_buf;
+
+            let read = new_current as usize - current as usize;
+            let new_state = if new_current == len {
+                State::ChunkEnd
+            } else {
+                State::Chunk(new_current, len)
+            };
 
-        loop {
-            let remaining = (len - current) as usize;
-            let mut to_read = std::cmp::min(remaining, buf.len());
-            let mut new_current = current;
-            let mut start = 0;
-
-            // first drain the buffer
-            if pos.len() > 0 {
-                let to_read_buf = std::cmp::min(to_read, pos.len());
-                dbg!(std::str::from_utf8(&buffer[pos.start..pos.start + to_read_buf]).unwrap());
-                buf[..to_read_buf].copy_from_slice(&buffer[pos.start..pos.start + to_read_buf]);
-                to_read -= to_read_buf;
-                new_pos.start += to_read_buf;
-                new_current += to_read_buf as u64;
-                start += to_read_buf;
-            }
-
-            if to_read > 0 {
-                dbg!("reading");
-                dbg!(to_read);
-                dbg!(start);
-                let n = match Pin::new(&mut self.inner)
-                    .poll_read(cx, &mut buf[start..start + to_read])
-                {
-                    Poll::Ready(val) => val?,
-                    Poll::Pending => {
-                        return Ok(DecodeResult::Some {
-                            read: 0,
-                            new_state: self.state.clone(),
-                            new_pos,
-                            buffer,
-                            pending: true,
-                        })
-                    }
-                };
+            return Ok(DecodeResult::Some {
+                read,
+                new_state,
+                new_pos,
+                buffer,
+                pending: false,
+            });
+        }
 
-                new_current += n as u64;
-            }
+        while to_read > 0 {
+            dbg!("poll_read_chunk_loop");
+
+            dbg!(to_read);
+            dbg!(start);
+            dbg!(new_pos);
+            dbg!(new_current);
+
+            dbg!("reading");
+            let n = match Pin::new(&mut self.inner).poll_read(cx, &mut buf[start..start + to_read])
+            {
+                Poll::Ready(val) => val?,
+                Poll::Pending => {
+                    dbg!("pending");
+                    return Ok(DecodeResult::Some {
+                        read: 0,
+                        new_state: State::Chunk(new_current, len),
+                        new_pos,
+                        buffer,
+                        pending: true,
+                    });
+                }
+            };
 
-            if new_current == len {
-                return Ok(DecodeResult::Some {
-                    read: new_current as usize - current as usize,
-                    new_state: State::ChunkEnd,
-                    new_pos,
-                    buffer,
-                    pending: false,
-                });
-            }
-            self.state = State::Chunk(new_current, len);
+            dbg!(n);
+            to_read -= n;
+            new_current += n as u64;
+            start += n;
         }
+
+        let read = new_current as usize - current as usize;
+        let new_state = if new_current == len {
+            State::ChunkEnd
+        } else {
+            State::Chunk(new_current, len)
+        };
+        dbg!(read);
+        dbg!(&new_state);
+        dbg!(new_current);
+        dbg!(len);
+        dbg!(current);
+
+        Ok(DecodeResult::Some {
+            read,
+            new_state,
+            new_pos,
+            buffer,
+            pending: false,
+        })
     }
 
     fn poll_read_inner(
@@ -226,6 +265,8 @@ impl<R: Read + Unpin> ChunkedDecoder<R> {
         pos: &Position,
         buf: &mut [u8],
     ) -> io::Result<DecodeResult> {
+        dbg!(&self.state);
+
         match self.state {
             State::Init => {
                 // Initial read
@@ -289,8 +330,13 @@ impl<R: Read + Unpin> Read for ChunkedDecoder<R> {
 
         let mut n = std::mem::replace(&mut this.current, Position::default());
         let buffer = std::mem::replace(&mut this.buffer, POOL.alloc(INITIAL_CAPACITY));
-        let mut needs_read = true;
+        let mut needs_read = if let State::Chunk(_, _) = this.state {
+            false // Do not attempt to fill the buffer when we are reading a chunk
+        } else {
+            true
+        };
 
+        dbg!(n);
         let mut buffer = if n.len() > 0 && this.initial_decode {
             match this.poll_read_inner(cx, buffer, &n, buf)? {
                 DecodeResult::Some {
@@ -329,6 +375,7 @@ impl<R: Read + Unpin> Read for ChunkedDecoder<R> {
             dbg!("loop");
             if n.len() >= buffer.capacity() {
                 if buffer.capacity() + 1024 <= MAX_CAPACITY {
+                    dbg!("resizing buffer");
                     buffer.realloc(buffer.capacity() + 1024);
                 } else {
                     std::mem::replace(&mut this.buffer, buffer);
@@ -341,6 +388,7 @@ impl<R: Read + Unpin> Read for ChunkedDecoder<R> {
             }
 
             if needs_read {
+                dbg!("reading");
                 let bytes_read = match Pin::new(&mut this.inner).poll_read(cx, &mut buffer[n.end..])
                 {
                     Poll::Ready(result) => result?,
@@ -354,6 +402,7 @@ impl<R: Read + Unpin> Read for ChunkedDecoder<R> {
                         return Poll::Pending;
                     }
                 };
+                dbg!(bytes_read);
                 n.end += bytes_read;
             }
             match this.poll_read_inner(cx, buffer, &n, buf)? {
@@ -364,6 +413,7 @@ impl<R: Read + Unpin> Read for ChunkedDecoder<R> {
                     new_state,
                     pending,
                 } => {
+                    dbg!("some");
                     // current buffer might now contain more data inside, so we need to attempt
                     // to decode it next time
                     this.initial_decode = true;
@@ -386,6 +436,8 @@ impl<R: Read + Unpin> Read for ChunkedDecoder<R> {
                     continue;
                 }
                 DecodeResult::None(buf) => {
+                    dbg!("none");
+
                     buffer = buf;
 
                     if this.buffer.is_empty() || n.is_zero() {
@@ -394,17 +446,10 @@ impl<R: Read + Unpin> Read for ChunkedDecoder<R> {
 
                         std::mem::replace(&mut this.buffer, buffer);
                         this.current = n;
-                        return Poll::Ready(Ok(0));
-                    } else if n.len() == 0 {
-                        // "logical buffer" is empty, there is nothing to decode on the next step
-                        this.initial_decode = false;
 
-                        std::mem::replace(&mut this.buffer, buffer);
-                        this.current = n;
-                        return Poll::Ready(Err(io::Error::new(
-                            io::ErrorKind::UnexpectedEof,
-                            "bytes remaining in stream",
-                        )));
+                        return Poll::Ready(Ok(0));
+                    } else {
+                        needs_read = true;
                     }
                 }
             }

src/client.rs

@@ -5,13 +5,15 @@ use async_std::prelude::*;
 use async_std::task::{Context, Poll};
 use futures_core::ready;
 use http_types::{
-    headers::{HeaderName, HeaderValue, CONTENT_LENGTH},
+    headers::{HeaderName, HeaderValue, CONTENT_LENGTH, TRANSFER_ENCODING},
     Body, Request, Response, StatusCode,
 };
 
 use std::pin::Pin;
 use std::str::FromStr;
 
+use crate::chunked::ChunkedDecoder;
+use crate::error::HttpError;
 use crate::{Exception, MAX_HEADERS};
 
 /// An HTTP encoder.
@@ -60,12 +62,23 @@ pub async fn encode(req: Request) -> Result<Encoder, std::io::Error> {
     }
 
     let val = format!("{} {} HTTP/1.1\r\n", req.method(), url);
+    log::trace!("> {}", &val);
+    buf.write_all(val.as_bytes()).await?;
+
+    let val = format!(
+        "Host: {}\r\n",
+        req.url()
+            .host_str()
+            .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "Missing hostname"))?
+    );
+    log::trace!("> {}", &val);
     buf.write_all(val.as_bytes()).await?;
 
     // If the body isn't streaming, we can set the content-length ahead of time. Else we need to
     // send all items in chunks.
     if let Some(len) = req.len() {
         let val = format!("Content-Length: {}\r\n", len);
+        log::trace!("> {}", &val);
         buf.write_all(val.as_bytes()).await?;
     } else {
         // write!(&mut buf, "Transfer-Encoding: chunked\r\n")?;
@@ -76,6 +89,7 @@ pub async fn encode(req: Request) -> Result<Encoder, std::io::Error> {
     for (header, values) in req.iter() {
         for value in values.iter() {
             let val = format!("{}: {}\r\n", header, value);
+            log::trace!("> {}", &val);
             buf.write_all(val.as_bytes()).await?;
         }
     }
@@ -131,21 +145,39 @@ where
         res.insert_header(name, value)?;
     }
 
-    // Process the body if `Content-Length` was passed.
-    if let Some(content_length) = res.header(&CONTENT_LENGTH) {
-        let length = content_length
-            .last()
-            .unwrap()
-            .as_str()
-            .parse::<usize>()
-            .ok();
-
-        if let Some(len) = length {
-            res.set_body(Body::from_reader(reader, Some(len)));
-        } else {
-            return Err("Invalid value for Content-Length".into());
+    let content_length = res.header(&CONTENT_LENGTH);
+    let transfer_encoding = res.header(&TRANSFER_ENCODING);
+
+    if content_length.is_some() && transfer_encoding.is_some() {
+        // This is always an error.
+        return Err(HttpError::UnexpectedContentLengthHeader.into());
+    }
+
+    // Check for Transfer-Encoding
+    match transfer_encoding {
+        Some(encoding) if !encoding.is_empty() => {
+            if encoding.last().unwrap().as_str() == "chunked" {
+                res.set_body(Body::from_reader(
+                    BufReader::new(ChunkedDecoder::new(reader)),
+                    None,
+                ));
+                return Ok(res);
+            }
+            // Fall through to Content-Length
         }
-    };
+        _ => {
+            // Fall through to Content-Length
+        }
+    }
+
+    // Check for Content-Length.
+    match content_length {
+        Some(len) => {
+            let len = len.last().unwrap().as_str().parse::<usize>()?;
+            res.set_body(Body::from_reader(reader.take(len as u64), Some(len)));
+        }
+        None => {}
+    }
 
     // Return the response.
     Ok(res)