Skip to content

Commit 77395c1

Browse files
authored
Harden Unpacker.__init__ re-entry cleanup to prevent buffer/context leaks (#687)
1 parent 7df7136 commit 77395c1

2 files changed

Lines changed: 42 additions & 3 deletions

File tree

msgpack/_unpacker.pyx

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -317,9 +317,6 @@ cdef class Unpacker:
317317
cdef Py_ssize_t max_buffer_size
318318
cdef uint64_t stream_offset
319319

320-
def __cinit__(self):
321-
self.buf = NULL
322-
323320
def __dealloc__(self):
324321
unpack_clear(&self.ctx)
325322
PyMem_Free(self.buf)
@@ -338,6 +335,12 @@ cdef class Unpacker:
338335
Py_ssize_t max_ext_len=-1):
339336
cdef const char *cerr=NULL
340337

338+
unpack_clear(&self.ctx)
339+
unpack_init(&self.ctx)
340+
if self.buf != NULL:
341+
PyMem_Free(self.buf)
342+
self.buf = NULL
343+
341344
self.object_hook = object_hook
342345
self.object_pairs_hook = object_pairs_hook
343346
self.list_hook = list_hook

test/test_unpack.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1+
import gc
12
import sys
3+
import weakref
24
from io import BytesIO
35

46
from pytest import mark, raises
@@ -87,3 +89,37 @@ def test_unpacker_tell_read_bytes():
8789
assert obj == unp
8890
assert pos == unpacker.tell()
8991
assert unpacker.read_bytes(n) == raw
92+
93+
94+
@mark.skipif(
95+
Unpacker.__module__ == "msgpack.fallback",
96+
reason="specific to C extension reinit leak",
97+
)
98+
def test_unpacker_reinit_clears_partial_state():
99+
refs = []
100+
101+
class Marker:
102+
pass
103+
104+
def hook(code, data):
105+
obj = Marker()
106+
refs.append(weakref.ref(obj))
107+
return obj
108+
109+
unpacker = Unpacker(ext_hook=hook, strict_map_key=False)
110+
# Keep parser state mid-map with a live key object from ext_hook.
111+
# Encodes: [ {ExtType(1, b"a"): <missing value>} ].
112+
unpacker.feed(b"\x91\x81\xd4\x01a")
113+
with raises(OutOfData):
114+
unpacker.unpack()
115+
assert len(refs) == 1
116+
assert refs[0]() is not None
117+
118+
unpacker.__init__()
119+
gc.collect()
120+
assert refs[0]() is None
121+
with raises(OutOfData):
122+
unpacker.unpack()
123+
124+
unpacker.feed(packb({"a": 1}))
125+
assert unpacker.unpack() == {"a": 1}

0 commit comments

Comments
 (0)