mirror of
https://github.com/msgpack/msgpack-python.git
synced 2025-10-19 12:03:15 +00:00
Fix Unpacker max_buffer_length handling (#506)
This commit is contained in:
parent
b75e3412fb
commit
500a238028
3 changed files with 24 additions and 17 deletions
|
@ -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)
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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")
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue