Bound the buffered read to the range the server declared

PartialBuffer represents a declared byte range of a remote file, but on the
non-stream path it called buffer.read() with no argument, consuming whatever the
server chose to send before the declared size was consulted. How much got
buffered was decided by the server rather than by the range requested: a client
asking for 100 bytes buffered 20 MB when the server streamed that much, which
defeats the purpose of fetching ranges at all.

Read at most `size` bytes instead, in a loop that tolerates short reads. A
single read() call is not enough: a socket-backed response can return fewer
bytes than requested while more are still coming, so reading once would silently
truncate. The data is written straight into the result buffer so no intermediate
copy of the whole range is held. The existing tests all use BytesIO, which never
short-reads, so these cases need their own tests.

RemoteFetcher.fetch also turned the server's Content-Range straight into that
size without checking it. A header whose end precedes its start, such as
"bytes 100-50/1000", produced a negative size, and PartialBuffer.read(0) then
computed a negative length, which for a file object means read everything.
Malformed values such as "bytes abc-def/1000", "bytes */1000" or an empty header
raised a bare ValueError out of the library. Both are now RemoteZipError.

Adds tests that a server sending more than it declared does not enlarge the
buffer, that a server sending less does not hang or raise, that short reads are
handled without truncation, and that invalid or malformed Content-Range values
are rejected.

Author: Gares95

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)