diff --git a/remotezip.py b/remotezip.py index 03d1e11..b43704d 100755 --- a/remotezip.py +++ b/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)) diff --git a/test_remotezip.py b/test_remotezip.py index 186a055..78afc9b 100644 --- a/test_remotezip.py +++ b/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)