Unpacker raises BufferFull when max_buffer_size is not enough.

msgpack · methane · May 24, 2022 · May 24, 2022 · May 24, 2022 · May 24, 2022
commit f233d0bcec6ac0f6ccf72512374ad47479a25f30
diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx
@@ -440,15 +440,17 @@ 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
@@ -471,11 +473,10 @@ cdef class Unpacker(object):
                if self.file_like is not None:
                    self.read_from_file()
                    continue
-                if prev_head == self.buf_tail:
-                    if iter:
-                        raise StopIteration("No more data to unpack.")
-                    else:
-                        raise OutOfData("No more data to unpack.")
+                if iter:
+                    raise StopIteration("No more data to unpack.")
+                else:
+                    raise OutOfData("No more data to unpack.")
            elif ret == -2:
                raise FormatError
            elif ret == -3:

diff --git 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"*100))
+    unpacker = Unpacker(buff, read_size=1, max_buffer_size=99)
+    assert unpacker.unpack() == b"a"*10
+    #assert unpacker.unpack() == b"a"*100
+    with raises(BufferFull):
+        unpacker.unpack() == b"a"*100
+
+
 def test_readbytes():
    unpacker = Unpacker(read_size=3)
    unpacker.feed(b"foobar")