diff options
author | Inada Naoki <songofacandy@gmail.com> | 2022-05-24 19:46:51 +0900 |
---|---|---|
committer | GitHub <noreply@github.com> | 2022-05-24 19:46:51 +0900 |
commit | 500a238028bdebe123b502b07769578b5f0e8a3a (patch) | |
tree | 9c5796a126485a13c3963a84e30a4542df9b2fe4 | |
parent | b75e3412fb8a2b6d6cd1da1b7063e14f6bfc0337 (diff) | |
download | msgpack-python-500a238028bdebe123b502b07769578b5f0e8a3a.tar.gz |
Fix Unpacker max_buffer_length handling (#506)
-rw-r--r-- | msgpack/_unpacker.pyx | 28 | ||||
-rw-r--r-- | msgpack/fallback.py | 2 | ||||
-rw-r--r-- | test/test_sequnpack.py | 11 |
3 files changed, 24 insertions, 17 deletions
diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx index 27facc0..8b06661 100644 --- a/msgpack/_unpacker.pyx +++ b/msgpack/_unpacker.pyx @@ -440,34 +440,30 @@ cdef class Unpacker(object): self.buf_size = buf_size self.buf_tail = tail + _buf_len - cdef read_from_file(self): - next_bytes = self.file_like_read( - min(self.read_size, - self.max_buffer_size - (self.buf_tail - self.buf_head) - )) + cdef int read_from_file(self) except -1: + cdef Py_ssize_t remains = self.max_buffer_size - (self.buf_tail - self.buf_head) + if remains <= 0: + raise BufferFull + + next_bytes = self.file_like_read(min(self.read_size, remains)) if next_bytes: self.append_buffer(PyBytes_AsString(next_bytes), PyBytes_Size(next_bytes)) else: self.file_like = None + return 0 cdef object _unpack(self, execute_fn execute, bint iter=0): cdef int ret cdef object obj cdef Py_ssize_t prev_head - if self.buf_head >= self.buf_tail and self.file_like is not None: - self.read_from_file() - while 1: prev_head = self.buf_head - if prev_head >= self.buf_tail: - if iter: - raise StopIteration("No more data to unpack.") - else: - raise OutOfData("No more data to unpack.") - - ret = execute(&self.ctx, self.buf, self.buf_tail, &self.buf_head) - self.stream_offset += self.buf_head - prev_head + if prev_head < self.buf_tail: + ret = execute(&self.ctx, self.buf, self.buf_tail, &self.buf_head) + self.stream_offset += self.buf_head - prev_head + else: + ret = 0 if ret == 1: obj = unpack_data(&self.ctx) diff --git a/msgpack/fallback.py b/msgpack/fallback.py index 5f215e9..f560c7b 100644 --- a/msgpack/fallback.py +++ b/msgpack/fallback.py @@ -423,6 +423,8 @@ class Unpacker(object): # Read from file remain_bytes = -remain_bytes + if remain_bytes + len(self._buffer) > self._max_buffer_size: + raise BufferFull while remain_bytes > 0: to_read_bytes = max(self._read_size, remain_bytes) read_data = self.file_like.read(to_read_bytes) diff --git a/test/test_sequnpack.py b/test/test_sequnpack.py index 9f20c07..c091076 100644 --- a/test/test_sequnpack.py +++ b/test/test_sequnpack.py @@ -2,7 +2,7 @@ # coding: utf-8 import io from msgpack import Unpacker, BufferFull -from msgpack import pack +from msgpack import pack, packb from msgpack.exceptions import OutOfData from pytest import raises @@ -78,6 +78,15 @@ def test_maxbuffersize(): assert ord("b") == next(unpacker) +def test_maxbuffersize_file(): + buff = io.BytesIO(packb(b"a" * 10) + packb([b"a" * 20] * 2)) + unpacker = Unpacker(buff, read_size=1, max_buffer_size=19, max_bin_len=20) + assert unpacker.unpack() == b"a" * 10 + # assert unpacker.unpack() == [b"a" * 20]*2 + with raises(BufferFull): + print(unpacker.unpack()) + + def test_readbytes(): unpacker = Unpacker(read_size=3) unpacker.feed(b"foobar") |