8000 Auto limit configuration (#342) · loude/msgpack-python@28b5f46 · GitHub
[go: up one dir, main page]
Skip to content

Commit 28b5f46

Browse files
methanemethane
authored
Auto limit configuration (msgpack#342)
1 parent f46523b commit 28b5f46

File tree

3 files changed

+92
-28
lines changed

3 files changed

+92
-28
lines changed

‎msgpack/_unpacker.pyx

Lines changed: 41 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -145,11 +145,11 @@ def unpackb(object packed, object object_hook=None, object list_hook=None,
145145
bint use_list=True, bint raw=True, bint strict_map_key=False,
146146
encoding=None, unicode_errors=None,
147147
object_pairs_hook=None, ext_hook=ExtType,
148-
Py_ssize_t max_str_len=1024*1024,
149-
Py_ssize_t max_bin_len=1024*1024,
150-
Py_ssize_t max_array_len=128*1024,
151-
Py_ssize_t max_map_len=32*1024,
152-
Py_ssize_t max_ext_len=1024*1024):
148+
Py_ssize_t max_str_len=-1,
149+
Py_ssize_t max_bin_len=-1,
150+
Py_ssize_t max_array_len=-1,
151+
Py_ssize_t max_map_len=-1,
152+
Py_ssize_t max_ext_len=-1):
153153
"""
154154
Unpack packed_bytes to object. Returns an unpacked object.
155155
@@ -160,6 +160,8 @@ def unpackb(object packed, object object_hook=None, object list_hook=None,
160160
Other exceptions can be raised during unpacking.
161161
162162
See :class:`Unpacker` for options.
163+
164+
*max_xxx_len* options are configured automatically from ``len(packed)``.
163165
"""
164166
cdef unpack_context ctx
165167
cdef Py_ssize_t off = 0
@@ -180,6 +182,18 @@ def unpackb(object packed, object object_hook=None, object list_hook=None,
180182
cerr = unicode_errors
181183

182184
get_data_from_buffer(packed, &view, &buf, &buf_len, &new_protocol)
185+
186+
if max_str_len == -1:
187+
max_str_len = buf_len
188+
if max_bin_len == -1:
189+
max_bin_len = buf_len
190+
if max_array_len == -1:
191+
max_array_len = buf_len
192+
if max_map_len == -1:
193+
max_map_len = buf_len//2
194+
if max_ext_len == -1:
195+
max_ext_len = buf_len
196+
183197
try:
184198
init_ctx(&ctx, object_hook, object_pairs_hook, list_hook, ext_hook,
185199
use_list, raw, strict_map_key, cenc, cerr,
@@ -259,19 +273,19 @@ cdef class Unpacker(object):
259273
You should set this parameter when unpacking data from untrusted source.
260274
261275
:param int max_str_len:
262-
Limits max length of str. (default: 1024*1024)
276+
Limits max length of str. (default: max_buffer_size or 1024*1024)
263277
264278
:param int max_bin_len:
265-
Limits max length of bin. (default: 1024*1024)
279+
Limits max length of bin. (default: max_buffer_size or 1024*1024)
266280
267281
:param int max_array_len:
268-
Limits max length of array. (default: 128*1024)
282+
Limits max length of array. (default: max_buffer_size or 128*1024)
269283
270284
:param int max_map_len:
271-
Limits max length of map. (default: 32*1024)
285+
Limits max length of map. (default: max_buffer_size//2 or 32*1024)
272286
273287
:param int max_ext_len:
274-
Limits max size of ext type. (default: 1024*1024)
288+
Limits max size of ext type. (default: max_buffer_size or 1024*1024)
275289
276290
:param str encoding:
277291
Deprecated, use raw instead.
@@ -329,11 +343,11 @@ cdef class Unpacker(object):
329343
object object_hook=None, object object_pairs_hook=None, object list_hook=None,
330344
encoding=None, unicode_errors=None, Py_ssize_t max_buffer_size=0,
331345
object ext_hook=ExtType,
332-
Py_ssize_t max_str_len=1024*1024,
333-
Py_ssize_t max_bin_len=1024*1024,
334-
Py_ssize_t max_array_len=128*1024,
335-
Py_ssize_t max_map_len=32*1024,
336-
Py_ssize_t max_ext_len=1024*1024):
346+
Py_ssize_t max_str_len=-1,
347+
Py_ssize_t max_bin_len=-1,
348+
Py_ssize_t max_array_len=-1,
349+
Py_ssize_t max_map_len=-1,
350+
Py_ssize_t max_ext_len=-1):
337351
cdef const char *cenc=NULL,
338352
cdef const char *cerr=NULL
339353

@@ -347,6 +361,18 @@ cdef class Unpacker(object):
347361
self.file_like_read = file_like.read
348362
if not PyCallable_Check(self.file_like_read):
349363
raise TypeError("`file_like.read` must be a callable.")
364+
365+
if max_str_len == -1:
366+
max_str_len = max_buffer_size or 1024*1024
367+
if max_bin_len == -1:
368+
max_bin_len = max_buffer_size or 1024*1024
369+
if max_array_len == -1:
370+
max_array_len = max_buffer_size or 128*1024
371+
if max_map_len == -1:
372+
max_map_len = max_buffer_size//2 or 32*1024
373+
if max_ext_len == -1:
374+
max_ext_len = max_buffer_size or 1024*1024
375+
350376
if not max_buffer_size:
351377
max_buffer_size = INT_MAX
352378
if read_size > max_buffer_size:

‎msgpack/fallback.py

Lines changed: 27 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,7 @@ def unpackb(packed, **kwargs):
130130
131131
See :class:`Unpacker` for options.
132132
"""
133-
unpacker = Unpacker(None, **kwargs)
133+
unpacker = Unpacker(None, max_buffer_size=len(packed), **kwargs)
134134
unpacker.feed(packed)
135135
try:
136136
ret = unpacker._unpack()
@@ -208,19 +208,24 @@ class Unpacker(object):
208208
You should set this parameter when unpacking data from untrusted source.
209209
210210
:param int max_str_len:
211-
Limits max length of str. (default: 1024*1024)
211+
(deprecated) Limits max length of str.
212+
(default: max_buffer_size or 1024*1024)
212213
213214
:param int max_bin_len:
214-
Limits max length of bin. (default: 1024*1024)
215+
(deprecated) Limits max length of bin.
216+
(default: max_buffer_size or 1024*1024)
215217
216218
:param int max_array_len:
217-
Limits max length of array. (default: 128*1024)
219+
Limits max length of array.
220+
(default: max_buffer_size or 128*1024)
218221
219222
:param int max_map_len:
220-
Limits max length of map. (default: 32*1024)
223+
Limits max length of map.
224+
(default: max_buffer_size//2 or 32*1024)
221225
222226
:param int max_ext_len:
223-
Limits max size of ext type. (default: 1024*1024)
227+
(deprecated) Limits max size of ext type.
228+
(default: max_buffer_size or 1024*1024)
224229
225230
example of streaming deserialize from file-like object::
226231
@@ -250,12 +255,11 @@ def __init__(self, file_like=None, read_size=0, use_list=True, raw=True, strict_
250255
object_hook=None, object_pairs_hook=None, list_hook=None,
251256
encoding=None, unicode_errors=None, max_buffer_size=0,
252257
ext_hook=ExtType,
253-
max_str_len=1024*1024,
254-
max_bin_len=1024*1024,
255-
max_array_len=128*1024,
256-
max_map_len=32*1024,
257-
max_ext_len=1024*1024):
258-
258+
max_str_len=-1,
259+
max_bin_len=-1,
260+
max_array_len=-1,
261+
max_map_len=-1,
262+
max_ext_len=-1):
259263
if encoding is not None:
260264
warnings.warn(
261265
"encoding is deprecated, Use raw=False instead.",
@@ -286,6 +290,17 @@ def __init__(self, file_like=None, read_size=0, use_list=True, raw=True, strict_
286290
# state, which _buf_checkpoint records.
287291
self._buf_checkpoint = 0
288292

293+
if max_str_len == -1:
294+
max_str_len = max_buffer_size or 1024*1024
295+
if max_bin_len == -1:
296+
max_bin_len = max_buffer_size or 1024*1024
297+
if max_array_len == -1:
298+
max_array_len = max_buffer_size or 128*1024
299+
if max_map_len == -1:
300+
max_map_len = max_buffer_size//2 or 32*1024
301+
if max_ext_len == -1:
302+
max_ext_len = max_buffer_size or 1024*1024
303+
289304
self._max_buffer_size = max_buffer_size or 2**31-1
290305
if read_size > self._max_buffer_size:
291306
raise ValueError("read_size must be smaller than max_buffer_size")

‎test/test_limits.py

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,6 @@ def test_max_ext_len():
105105
unpacker.unpack()
106106

107107

108-
109108
# PyPy fails following tests because of constant folding?
110109
# https://bugs.pypy.org/issue1721
111110
#@pytest.mark.skipif(True, reason="Requires very large memory.")
@@ -134,3 +133,27 @@ def test_max_ext_len():
134133
# x.append(0)
135134
# with pytest.raises(ValueError):
136135
# packb(x)
136+
137+
138+
# auto max len
139+
140+
def test_auto_max_array_len():
141+
packed = b'\xde\x00\x06zz'
142+
with pytest.raises(UnpackValueError):
143+
unpackb(packed, raw=False)
144+
145+
unpacker = Unpacker(max_buffer_size=5, raw=False)
146+
unpacker.feed(packed)
147+
with pytest.raises(UnpackValueError):
148+
unpacker.unpack()
149+
150+
def test_auto_max_map_len():
151+
# len(packed) == 6 -> max_map_len == 3
152+
packed = b'\xde\x00\x04zzz'
153+
with pytest.raises(UnpackValueError):
154+
unpackb(packed, raw=False)
155+
156+
unpacker = Unpacker(max_buffer_size=6, raw=False)
157+
unpacker.feed(packed)
158+
with pytest.raises(UnpackValueError):
159+
unpacker.unpack()

0 commit comments

Comments
 (0)
0