Chunked responses

Author: Lukasa

hyper/common/exceptions.py

@@ -0,0 +1,12 @@
+# -*- coding: utf-8 -*-
+"""
+hyper/common/exceptions
+~~~~~~~~~~~~~~~~~~~~~~~
+
+Contains hyper's exceptions.
+"""
+class ChunkedDecodeError(Exception):
+    """
+    An error was encountered while decoding a chunked response.
+    """
+    pass

hyper/http11/response.py

@@ -10,6 +10,7 @@
 import zlib
 
 from ..common.decoder import DeflateDecoder
+from ..common.exceptions import ChunkedDecodeError
 from ..http20.exceptions import ConnectionResetError
 
 log = logging.getLogger(__name__)
@@ -52,7 +53,7 @@ def __init__(self, code, reason, headers, sock):
             self._length = None
 
         # Whether we expect a chunked response.
-        self._chunked = b'chunked' in self.headers.get('transfer-encoding', [])
+        self._chunked = b'chunked' in self.headers.get(b'transfer-encoding', [])
 
         # One of the following must be true: we must expect that the connection
         # will be closed following the body, or that a content-length was sent,
@@ -153,6 +154,44 @@ def read(self, amt=None, decode_content=True):
 
         return data
 
+    def read_chunked(self):
+        """
+        Reads chunked transfer encoded bodies. Each read returns a single
+        chunk.
+        """
+        if not self._chunked:
+            raise ChunkedDecodeError(
+                "Attempted chunked read of non-chunked body."
+            )
+
+        # Return early if possible.
+        if self._sock is None:
+            return b''
+
+        # Read to the newline to get the chunk length. This is a hexadecimal
+        # integer.
+        chunk_length = int(self._sock.readline().tobytes().strip(), 16)
+
+        # If the chunk length is zero, consume the newline and then we're done.
+        if not chunk_length:
+            self._sock.readline()
+            return b''
+
+        # Then read that many bytes.
+        data = b''
+
+        while chunk_length > 0:
+            chunk = self._sock.recv(chunk_length).tobytes()
+            data += chunk
+            chunk_length -= len(chunk)
+
+        assert chunk_length == 0
+
+        # Now, consume the newline.
+        self._sock.readline()
+
+        return data
+
     def close(self):
         """
         Close the response. In effect this closes the backing socket.

test/test_http11.py

@@ -18,6 +18,7 @@
 from hyper.http11.response import HTTP11Response
 from hyper.http20.exceptions import ConnectionResetError
 from hyper.common.headers import HTTPHeaderMap
+from hyper.common.exceptions import ChunkedDecodeError
 from hyper.compat import bytes, zlib_compressobj
 
 
@@ -383,6 +384,43 @@ def test_response_transparently_decrypts_wrong_deflate(self):
 
         assert r.read() == b'this is test data'
 
+    def test_basic_chunked_read(self):
+        d = DummySocket()
+        r = HTTP11Response(200, 'OK', {b'transfer-encoding': [b'chunked']}, d)
+
+        data = (
+            b'4\r\nwell\r\n'
+            b'4\r\nwell\r\n'
+            b'4\r\nwhat\r\n'
+            b'4\r\nhave\r\n'
+            b'2\r\nwe\r\n'
+            b'a\r\nhereabouts\r\n'
+            b'0\r\n\r\n'
+        )
+        d._buffer = BytesIO(data)
+        results = [
+            b'well', b'well', b'what', b'have', b'we', b'hereabouts'
+        ]
+
+        for c in results:
+            assert r.read_chunked() == c
+
+        assert not r.read_chunked()
+
+    def test_chunked_read_of_non_chunked(self):
+        r = HTTP11Response(200, 'OK', {b'content-length': [b'0']}, None)
+
+        with pytest.raises(ChunkedDecodeError):
+            r.read_chunked()
+
+    def test_chunked_read_aborts_early(self):
+        r = HTTP11Response(
+            200, 'OK', {b'transfer-encoding': [b'chunked']}, None
+        )
+
+        assert not r.read_chunked()
+
+
 class DummySocket(object):
     def __init__(self):
         self.queue = []