Support bounded reads of 'expect close' bodies

Author: Lukasa

hyper/http11/response.py

@@ -112,8 +112,8 @@ def read(self, amt=None, decode_content=True):
         # Otherwise, we've been asked to do a bounded read. We should read no
         # more than the remaining length, obviously.
         # FIXME: Handle cases without _length
-        assert self._length is not None
-        amt = min(amt, self._length)
+        if self._length is not None:
+            amt = min(amt, self._length)
 
         # If we are now going to read nothing, exit early.
         if not amt:
@@ -132,28 +132,41 @@ def read(self, amt=None, decode_content=True):
             chunk = self._sock.recv(amt).tobytes()
 
             # If we got an empty read, but were expecting more, the remote end
-            # has hung up. Raise an exception.
+            # has hung up. Raise an exception if we were expecting more data,
+            # but if we were expecting the remote end to close then it's ok.
             if not chunk:
-                self.close()
-                raise ConnectionResetError("Remote end hung up!")
+                if self._length is not None or not self._expect_close:
+                    self.close()
+                    raise ConnectionResetError("Remote end hung up!")
+
+                break
 
             to_read -= len(chunk)
             chunks.append(chunk)
 
         data = b''.join(chunks)
-        self._length -= len(data)
+        if self._length is not None:
+            self._length -= len(data)
+
+        # If we're at the end of the request, we have some cleaning up to do.
+        # Close the stream, and if necessary flush the buffer. Checking that
+        # we're at the end is actually obscenely complex: either we've read the
+        # full content-length or, if we were expecting a closed connection,
+        # we've had a read shorter than the requested amount. We also have to
+        # do this before we try to decompress the body.
+        end_of_request = (self._length == 0 or
+                          (self._expect_close and len(data) < amt))
 
         # We may need to decode the body.
         if decode_content and self._decompressobj and data:
             data = self._decompressobj.decompress(data)
 
-        # If we're at the end of the request, we have some cleaning up to do.
-        # Close the stream, and if necessary flush the buffer.
-        if decode_content and self._decompressobj and not self._length:
+        if decode_content and self._decompressobj and end_of_request:
             data += self._decompressobj.flush()
 
-        # We're at the end. Close the connection.
-        if not self._length:
+        # We're at the end. Close the connection. Explicit check for zero here
+        # because self._length might be None.
+        if end_of_request:
             self.close()
 
         return data

test/test_http11.py

@@ -482,6 +482,56 @@ def test_chunk_length_read(self):
         assert r.read(20) == b'reabouts'
         assert r.read() == b''
 
+    def test_bounded_read_expect_close_no_content_length(self):
+        d = DummySocket()
+        r = HTTP11Response(200, 'OK', {b'connection': [b'close']}, d)
+        d._buffer = BytesIO(b'hello there sir')
+
+        assert r.read(5) == b'hello'
+        assert r.read(6) == b' there'
+        assert r.read(8) == b' sir'
+        assert r.read(9) == b''
+
+        assert r._sock is None
+
+    def test_bounded_read_expect_close_with_content_length(self):
+        headers = {b'connection': [b'close'], b'content-length': [b'15']}
+        d = DummySocket()
+        r = HTTP11Response(200, 'OK', headers, d)
+        d._buffer = BytesIO(b'hello there sir')
+
+        assert r.read(5) == b'hello'
+        assert r.read(6) == b' there'
+        assert r.read(8) == b' sir'
+        assert r.read(9) == b''
+
+        assert r._sock is None
+
+    def test_compressed_bounded_read_expect_close(self):
+        headers = {b'connection': [b'close'], b'content-encoding': [b'gzip']}
+
+        c = zlib_compressobj(wbits=24)
+        body = c.compress(b'hello there sir')
+        body += c.flush()
+
+        d = DummySocket()
+        r = HTTP11Response(200, 'OK', headers, d)
+        d._buffer = BytesIO(body)
+
+        response = b''
+        while True:
+            # 12 is magic here: it's the smallest read that actually
+            # decompresses immediately.
+            chunk = r.read(15)
+            if not chunk:
+                break
+
+            response += chunk
+
+        assert response == b'hello there sir'
+
+        assert r._sock is None
+
 class DummySocket(object):
     def __init__(self):
         self.queue = []