Bound the buffered read to the range the server declared

remotezip.py

@@ -30,8 +30,30 @@ class PartialBuffer:
         however, any attempt to read data outside the partial data is going to fail
         with OutOfBound error.
     """
+    @staticmethod
+    def _read_up_to(buffer, size):
+        """Read at most `size` bytes into a new buffer.
+
+        A single read() is not enough: a socket-backed response can return
+        fewer bytes than requested while more are still coming. Data is written
+        straight into the result so that no intermediate copy of the whole
+        range is held.
+        """
+        result = io.BytesIO()
+        remaining = size
+        while remaining > 0:
+            chunk = buffer.read(remaining)
+            if not chunk:
+                break
+            result.write(chunk)
+            remaining -= len(chunk)
+        result.seek(0)
+        return result
+
     def __init__(self, buffer, offset, size, stream):
-        self.buffer = buffer if stream else io.BytesIO(buffer.read())
+        # Read at most `size` bytes: the declared range is what this buffer
+        # represents, and a server may send more than it announced.
+        self.buffer = buffer if stream else self._read_up_to(buffer, size)
         self._offset = offset
         self._size = size
         self._position = offset
@@ -223,7 +245,14 @@ def fetch(self, data_range, stream=False):
         kwargs = self.prepare_request(data_range)
         try:
             res, range_header = self._request(kwargs)
-            range_min, range_max = self.parse_range_header(range_header)
+            try:
+                range_min, range_max = self.parse_range_header(range_header)
+            except ValueError:
+                raise RemoteZipError(
+                    "Malformed Content-Range returned by the server: %s" % range_header)
+            if range_max is None or range_max < range_min:
+                raise RemoteZipError(
+                    "Invalid Content-Range returned by the server: %s" % range_header)
             return PartialBuffer(res, range_min, range_max - range_min + 1, stream)
         except IOError as e:
             raise RemoteIOError(str(e))

test_remotezip.py

@@ -65,6 +65,36 @@ def fetch(self, data_range, stream=False):
 
 
 class TestPartialBuffer(unittest.TestCase):
+    def test_handles_short_reads_from_the_stream(self):
+        """A socket-backed response may return less than asked for per read."""
+        class ShortReader(io.RawIOBase):
+            def __init__(self, data, chunk):
+                self._b = io.BytesIO(data)
+                self._chunk = chunk
+
+            def read(self, n=-1):
+                if n is None or n < 0:
+                    return self._b.read()
+                return self._b.read(min(n, self._chunk))
+
+        data = b'z' * 1000
+        for chunk in (1000, 512, 100, 1):
+            pb = rz.PartialBuffer(ShortReader(data, chunk), 0, len(data), stream=False)
+            self.assertEqual(pb.read(0), data)
+
+    def test_handles_a_server_sending_less_than_declared(self):
+        """A truncated response must not hang or raise; it yields what arrived."""
+        pb = rz.PartialBuffer(io.BytesIO(b'z' * 40), 0, 1000, stream=False)
+        self.assertEqual(pb.read(0), b'z' * 40)
+
+    def test_does_not_buffer_more_than_declared_size(self):
+        """A server sending more than it declared must not enlarge the buffer."""
+        oversized = io.BytesIO(b'x' * 10000)
+        pb = rz.PartialBuffer(oversized, 0, 100, stream=False)
+        self.assertEqual(len(pb.read(0)), 100)
+        # the rest of the response was never pulled into memory
+        self.assertEqual(oversized.tell(), 100)
+
     def setUp(self):
         if not hasattr(self, 'assertRaisesRegex'):
             self.assertRaisesRegex = self.assertRaisesRegexp
@@ -206,6 +236,33 @@ def test_build_range_header(self):
         header = rz.RemoteFetcher.build_range_header(-123, None)
         self.assertEqual(header, 'bytes=-123')
 
+    def test_fetch_rejects_invalid_content_range(self):
+        """A server must not be able to declare a range that ends before it starts."""
+        class Fetcher(rz.RemoteFetcher):
+            def __init__(self, header):
+                super(Fetcher, self).__init__('http://test.com/file.zip')
+                self.header = header
+
+            def _request(self, kwargs):
+                return io.BytesIO(b'x' * 100), self.header
+
+        with self.assertRaises(rz.RemoteZipError):
+            Fetcher('bytes 100-50/1000').fetch((0, 99))
+
+        with self.assertRaises(rz.RemoteZipError):
+            Fetcher('bytes -500/1000').fetch((0, 99))
+
+        # a malformed header must not leak a bare ValueError to the caller.
+        # 'bytes */1000' is the RFC 7233 unsatisfied-range form, so this is not
+        # only about hostile input.
+        for bad in ('bytes abc-def/1000', 'bytes -500-100/1000', 'bytes /1000',
+                    'bytes */1000', ''):
+            with self.assertRaises(rz.RemoteZipError):
+                Fetcher(bad).fetch((0, 99))
+
+        # an unknown total length is legitimate and must still be accepted
+        Fetcher('bytes 0-99/*').fetch((0, 99))
+
     def test_parse_range_header(self):
         range_min, range_max = rz.RemoteFetcher.parse_range_header('bytes 0-11/12')
         self.assertEqual(range_min, 0)