Refuse to export the Packer buffer during pack

msgpack/_packer.pyx

@@ -106,6 +106,7 @@ cdef class Packer:
     cdef object _berrors
     cdef const char *unicode_errors
     cdef size_t exports  # number of exported buffers
+    cdef int _packing  # public pack methods currently running
     cdef bint strict_types
     cdef bint use_float
     cdef bint autoreset
@@ -118,6 +119,7 @@ cdef class Packer:
         self.pk.buf_size = buf_size
         self.pk.length = 0
         self.exports = 0
+        self._packing = 0
 
     def __dealloc__(self):
         PyMem_Free(self.pk.buf)
@@ -273,29 +275,40 @@ cdef class Packer:
     def pack(self, object obj):
         cdef int ret
         self._check_exports()
+        # User code runs inside _pack (default, items(), iteration) and can
+        # export this buffer. Hold the flag until every write from this call
+        # has finished, including a nested pack.
+        self._packing += 1
         try:
-            ret = self._pack(obj, DEFAULT_RECURSE_LIMIT)
-        except:
-            self.pk.length = 0
-            raise
-        if ret:  # should not happen.
-            raise RuntimeError("internal error")
-        if self.autoreset:
-            buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
-            self.pk.length = 0
-            return buf
+            try:
+                ret = self._pack(obj, DEFAULT_RECURSE_LIMIT)
+            except:
+                self.pk.length = 0
+                raise
+            if ret:  # should not happen.
+                raise RuntimeError("internal error")
+            if self.autoreset:
+                buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
+                self.pk.length = 0
+                return buf
+        finally:
+            self._packing -= 1
 
     @cython.critical_section
     def pack_ext_type(self, typecode, data):
         self._check_exports()
-        if len(data) > ITEM_LIMIT:
-            raise ValueError("ext data too large")
-        msgpack_pack_ext(&self.pk, typecode, len(data))
-        msgpack_pack_raw_body(&self.pk, data, len(data))
-        if self.autoreset:
-            buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
-            self.pk.length = 0
-            return buf
+        self._packing += 1
+        try:
+            if len(data) > ITEM_LIMIT:
+                raise ValueError("ext data too large")
+            msgpack_pack_ext(&self.pk, typecode, len(data))
+            msgpack_pack_raw_body(&self.pk, data, len(data))
+            if self.autoreset:
+                buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
+                self.pk.length = 0
+                return buf
+        finally:
+            self._packing -= 1
 
     @cython.critical_section
     def pack_array_header(self, long long size):
@@ -332,17 +345,21 @@ cdef class Packer:
         (`len(pairs)` and `for k, v in pairs:` should be supported.)
         """
         self._check_exports()
-        size = len(pairs)
-        if size > ITEM_LIMIT:
-            raise ValueError("map too large")
-        msgpack_pack_map(&self.pk, size)
-        for k, v in pairs:
-            self._pack(k)
-            self._pack(v)
-        if self.autoreset:
-            buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
-            self.pk.length = 0
-            return buf
+        self._packing += 1
+        try:
+            size = len(pairs)
+            if size > ITEM_LIMIT:
+                raise ValueError("map too large")
+            msgpack_pack_map(&self.pk, size)
+            for k, v in pairs:
+                self._pack(k)
+                self._pack(v)
+            if self.autoreset:
+                buf = PyBytes_FromStringAndSize(self.pk.buf, self.pk.length)
+                self.pk.length = 0
+                return buf
+        finally:
+            self._packing -= 1
 
     @cython.critical_section
     def reset(self):
@@ -362,11 +379,16 @@ cdef class Packer:
         """Return memoryview of internal buffer.
 
         Note: Packer now supports buffer protocol. You can use memoryview(packer).
+        Raises BufferError while a pack method is still running.
         """
         return memoryview(self)
 
     @cython.critical_section
     def __getbuffer__(self, Py_buffer *buffer, int flags):
+        if self._packing:
+            raise BufferError(
+                "Packer cannot export its buffer while packing is in progress"
+            )
         PyBuffer_FillInfo(buffer, self, self.pk.buf, self.pk.length, 1, flags)
         self.exports += 1
 

test/test_buffer.py

@@ -47,3 +47,64 @@ def test_packer_getbuffer():
         buffer.release()
         packer.pack(42)
         assert bytes(packer) == b"\x92*\xa5hello*"
+
+
+def _c_extension():
+    return Packer.__module__ == "msgpack._cmsgpack"
+
+
+def test_getbuffer_during_default_is_rejected():
+    # The default callback runs in the middle of pack(). Exporting the
+    # internal buffer there lets a later realloc free the view.
+    if not _c_extension():
+        return
+
+    class Unsupported:
+        pass
+
+    def default(obj):
+        packer.getbuffer()
+        return b"ok"
+
+    packer = Packer(default=default, autoreset=False)
+    with raises(BufferError):
+        packer.pack(Unsupported())
+    packer.pack(b"ok")
+    assert bytes(packer.getbuffer()) == b"\xc4\x02ok"
+
+
+def test_getbuffer_during_dict_items_is_rejected():
+    if not _c_extension():
+        return
+
+    packer = Packer(autoreset=False)
+
+    class EvilDict(dict):
+        def items(self):
+            packer.getbuffer()
+            return [(b"k", b"v")]
+
+    with raises(BufferError):
+        packer.pack(EvilDict({b"x": 1}))
+    assert bytes(packer.getbuffer()) == b""
+
+
+def test_getbuffer_during_map_pairs_is_rejected():
+    if not _c_extension():
+        return
+
+    packer = Packer(autoreset=False)
+
+    class EvilPairs:
+        def __len__(self):
+            return 1
+
+        def __iter__(self):
+            packer.getbuffer()
+            yield (b"k", b"v")
+
+    with raises(BufferError):
+        packer.pack_map_pairs(EvilPairs())
+    packer.reset()
+    packer.pack(b"ok")
+    assert bytes(packer.bytes()) == b"\xc4\x02ok"