diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx index 27facc0a..157291ec 100644 --- a/msgpack/_unpacker.pyx +++ b/msgpack/_unpacker.pyx @@ -441,12 +441,17 @@ cdef class Unpacker(object): self.buf_tail = tail + _buf_len cdef read_from_file(self): + current_size = self.buf_tail - self.buf_head next_bytes = self.file_like_read( - min(self.read_size, - self.max_buffer_size - (self.buf_tail - self.buf_head) - )) + min(self.read_size, self.max_buffer_size - current_size + 1) + ) if next_bytes: - self.append_buffer(PyBytes_AsString(next_bytes), PyBytes_Size(next_bytes)) + next_bytes_size = PyBytes_Size(next_bytes) + if next_bytes_size + current_size > self.max_buffer_size: + raise ValueError( + "object exceeds max_buffer_size(%s)" % self.max_buffer_size + ) + self.append_buffer(PyBytes_AsString(next_bytes), next_bytes_size) else: self.file_like = None diff --git a/test/test_unpack.py b/test/test_unpack.py index aa4c01f8..c7070763 100644 --- a/test/test_unpack.py +++ b/test/test_unpack.py @@ -1,6 +1,6 @@ from io import BytesIO import sys -from msgpack import Unpacker, packb, OutOfData, ExtType +from msgpack import Unpacker, pack, packb, OutOfData, ExtType from pytest import raises, mark try: @@ -90,3 +90,20 @@ def test_unpacker_tell_read_bytes(): assert obj == unp assert pos == unpacker.tell() assert unpacker.read_bytes(n) == raw + + +def test_unpacker_raise_max_buffer_size(): + max_buffer_size = 1024 + small_value = b"a" * max_buffer_size + + f = BytesIO() + pack(small_value, f) + f.seek(0) + assert list(Unpacker(f, max_buffer_size=max_buffer_size)) == [small_value] + + large_value = b"a" * (max_buffer_size + 1) + f = BytesIO() + pack(large_value, f) + f.seek(0) + with raises(ValueError): + list(Unpacker(f, max_buffer_size=max_buffer_size))