diff --git a/test_umsgpack.py b/test_umsgpack.py index 1652d60..0d1b672 100644 --- a/test_umsgpack.py +++ b/test_umsgpack.py @@ -381,6 +381,7 @@ "ReservedCodeException", "UnhashableKeyException", "DuplicateKeyException", + "MaximumDepthException", "KeyNotPrimitiveException", "KeyDuplicateException", "ext_serializable", @@ -764,12 +765,33 @@ def test_streaming_reader(self): reader = io.BytesIO(data) self.assertEqual(umsgpack.unpack(reader), obj) + def test_max_depth_default_wraps_recursion(self): + # Default (no max_depth): deeply nested input must raise a + # MaximumDepthException (an UnpackException), not a bare RecursionError. + payload = b"\x91" * 5000 + b"\xc0" + self.assertRaises(umsgpack.MaximumDepthException, umsgpack.unpackb, payload) + self.assertRaises(umsgpack.UnpackException, umsgpack.unpackb, payload) + # Moderate nesting still works. + self.assertEqual(umsgpack.unpackb(b"\x91" * 100 + b"\xc0"), + umsgpack.unpackb(b"\x91" * 100 + b"\xc0")) + + def test_max_depth_cap_arrays_and_maps(self): + self.assertTrue( + umsgpack.unpackb(b"\x91" * 32 + b"\xc0", max_depth=32) is not None) + self.assertRaises(umsgpack.MaximumDepthException, + umsgpack.unpackb, b"\x91" * 33 + b"\xc0", max_depth=32) + self.assertRaises(umsgpack.MaximumDepthException, + umsgpack.unpackb, b"\x81\xa1k" * 17 + b"\xc0", max_depth=16) + # Siblings are not counted, only nesting depth. + self.assertEqual( + len(umsgpack.unpackb(b"\xdc\x00\xc8" + b"\xc0" * 200, max_depth=1)), 200) + def test_namespacing(self): # Get a list of global variables from umsgpack module exported_vars = list([x for x in dir(umsgpack) if not x.startswith("_")]) # Ignore imports exported_vars = list([x for x in exported_vars if x != "struct" and x != "collections" and x != "datetime" and x != - "sys" and x != "io" and x != "xrange" and x != "Hashable"]) + "sys" and x != "io" and x != "xrange" and x != "Hashable" and x != "functools"]) self.assertTrue(len(exported_vars) == len(exported_vars_test_vector)) for var in exported_vars_test_vector: diff --git a/umsgpack/__init__.py b/umsgpack/__init__.py index e7e194f..55b827e 100644 --- a/umsgpack/__init__.py +++ b/umsgpack/__init__.py @@ -45,6 +45,7 @@ """ import struct import collections +import functools import datetime import sys import io @@ -238,6 +239,10 @@ class DuplicateKeyException(UnpackException): "Duplicate key encountered during map unpacking." +class MaximumDepthException(UnpackException): + "Maximum nesting depth exceeded during unpacking." + + # Backwards compatibility KeyNotPrimitiveException = UnhashableKeyException KeyDuplicateException = DuplicateKeyException @@ -868,7 +873,7 @@ def _unpack_ext_timestamp(ext_data, options): microseconds=microseconds) -def _unpack_array(code, fp, options): +def _unpack_array(code, fp, options, depth=None): if (ord(code) & 0xf0) == 0x90: length = (ord(code) & ~0xf0) elif code == b'\xdc': @@ -879,9 +884,12 @@ def _unpack_array(code, fp, options): raise Exception("logic error, not array: 0x{:02x}".format(ord(code))) if options.get('use_tuple'): - return tuple((_unpack(fp, options) for i in xrange(length))) + return tuple((_unpack(fp, options) if depth is None + else _unpack_depth(fp, options, depth) + for i in xrange(length))) - return [_unpack(fp, options) for i in xrange(length)] + return [_unpack(fp, options) if depth is None + else _unpack_depth(fp, options, depth) for i in xrange(length)] def _deep_list_to_tuple(obj): @@ -890,7 +898,7 @@ def _deep_list_to_tuple(obj): return obj -def _unpack_map(code, fp, options): +def _unpack_map(code, fp, options, depth=None): if (ord(code) & 0xf0) == 0x80: length = (ord(code) & ~0xf0) elif code == b'\xde': @@ -903,7 +911,8 @@ def _unpack_map(code, fp, options): d = {} if not options.get('use_ordered_dict') else collections.OrderedDict() for _ in xrange(length): # Unpack key - k = _unpack(fp, options) + k = _unpack(fp, options) if depth is None \ + else _unpack_depth(fp, options, depth) if isinstance(k, list): # Attempt to convert list into a hashable tuple @@ -916,7 +925,8 @@ def _unpack_map(code, fp, options): "encountered duplicate key: \"{:s}\" ({:s})".format(str(k), str(type(k)))) # Unpack value - v = _unpack(fp, options) + v = _unpack(fp, options) if depth is None \ + else _unpack_depth(fp, options, depth) try: d[k] = v @@ -930,6 +940,33 @@ def _unpack(fp, options): code = _read_except(fp, 1) return _unpack_dispatch_table[code](code, fp, options) + +def _unpack_depth(fp, options, depth): + code = _read_except(fp, 1) + return _unpack_depth_dispatch_table[code](code, fp, options, depth) + + +def _limit_depth(unpack_container): + """Wrap a depth-aware container unpacker to enforce options['max_depth']. + + depth is passed as a local argument and incremented on entry. + """ + @functools.wraps(unpack_container) + def wrapper(code, fp, options, depth): + depth += 1 + if depth > options['max_depth']: + raise MaximumDepthException("maximum nesting depth exceeded") + return unpack_container(code, fp, options, depth) + return wrapper + + +def _ignore_depth(unpacker): + """Adapt a scalar (non-container) unpacker to the depth dispatch signature.""" + @functools.wraps(unpacker) + def wrapper(code, fp, options, depth): + return unpacker(code, fp, options) + return wrapper + ######################################## @@ -951,6 +988,10 @@ def _unpack2(fp, **options): allow_invalid_utf8 (bool): unpack invalid strings into instances of :class:`InvalidString`, for access to the bytes (default False) + max_depth (int): maximum nesting depth of arrays/maps to unpack; when + exceeded, :class:`MaximumDepthException` is raised. If + None (default), depth is bounded only by the + interpreter's recursion limit. Returns: Python object @@ -969,13 +1010,21 @@ def _unpack2(fp, **options): The serialized map cannot be deserialized into a Python dictionary. DuplicateKeyException(UnpackException): Duplicate key encountered during map unpacking. + MaximumDepthException(UnpackException): + Maximum nesting depth exceeded while unpacking a deeply-nested + object, or the configured limit was exceeded (when ``max_depth`` is + set). Example: >>> f = open('test.bin', 'rb') >>> umsgpack.unpackb(f) {u'compact': True, u'schema': 0} """ - return _unpack(fp, options) + try: + return _unpack(fp, options) if options.get('max_depth') is None \ + else _unpack_depth(fp, options, 0) + except RecursionError: + raise MaximumDepthException("maximum nesting depth exceeded") def _unpack3(fp, **options): @@ -996,6 +1045,10 @@ def _unpack3(fp, **options): allow_invalid_utf8 (bool): unpack invalid strings into instances of :class:`InvalidString`, for access to the bytes (default False) + max_depth (int): maximum nesting depth of arrays/maps to unpack; when + exceeded, :class:`MaximumDepthException` is raised. If + None (default), depth is bounded only by the + interpreter's recursion limit. Returns: Python object @@ -1014,13 +1067,21 @@ def _unpack3(fp, **options): The serialized map cannot be deserialized into a Python dictionary. DuplicateKeyException(UnpackException): Duplicate key encountered during map unpacking. + MaximumDepthException(UnpackException): + Maximum nesting depth exceeded while unpacking a deeply-nested + object, or the configured limit was exceeded (when ``max_depth`` is + set). Example: >>> f = open('test.bin', 'rb') >>> umsgpack.unpackb(f) {'compact': True, 'schema': 0} """ - return _unpack(fp, options) + try: + return _unpack(fp, options) if options.get('max_depth') is None \ + else _unpack_depth(fp, options, 0) + except RecursionError: + raise MaximumDepthException("maximum nesting depth exceeded") # For Python 2, expects a str object @@ -1042,6 +1103,10 @@ def _unpackb2(s, **options): allow_invalid_utf8 (bool): unpack invalid strings into instances of :class:`InvalidString`, for access to the bytes (default False) + max_depth (int): maximum nesting depth of arrays/maps to unpack; when + exceeded, :class:`MaximumDepthException` is raised. If + None (default), depth is bounded only by the + interpreter's recursion limit. Returns: Python object @@ -1062,6 +1127,10 @@ def _unpackb2(s, **options): The serialized map cannot be deserialized into a Python dictionary. DuplicateKeyException(UnpackException): Duplicate key encountered during map unpacking. + MaximumDepthException(UnpackException): + Maximum nesting depth exceeded while unpacking a deeply-nested + object, or the configured limit was exceeded (when ``max_depth`` is + set). Example: >>> umsgpack.unpackb(b'\\x82\\xa7compact\\xc3\\xa6schema\\x00') @@ -1069,7 +1138,11 @@ def _unpackb2(s, **options): """ if not isinstance(s, (str, bytearray)): raise TypeError("packed data must be type 'str' or 'bytearray'") - return _unpack(io.BytesIO(s), options) + try: + return _unpack(io.BytesIO(s), options) if options.get('max_depth') is None \ + else _unpack_depth(io.BytesIO(s), options, 0) + except RecursionError: + raise MaximumDepthException("maximum nesting depth exceeded") # For Python 3, expects a bytes object @@ -1091,6 +1164,10 @@ def _unpackb3(s, **options): allow_invalid_utf8 (bool): unpack invalid strings into instances of :class:`InvalidString`, for access to the bytes (default False) + max_depth (int): maximum nesting depth of arrays/maps to unpack; when + exceeded, :class:`MaximumDepthException` is raised. If + None (default), depth is bounded only by the + interpreter's recursion limit. Returns: Python object @@ -1111,6 +1188,10 @@ def _unpackb3(s, **options): The serialized map cannot be deserialized into a Python dictionary. DuplicateKeyException(UnpackException): Duplicate key encountered during map unpacking. + MaximumDepthException(UnpackException): + Maximum nesting depth exceeded while unpacking a deeply-nested + object, or the configured limit was exceeded (when ``max_depth`` is + set). Example: >>> umsgpack.unpackb(b'\\x82\\xa7compact\\xc3\\xa6schema\\x00') @@ -1118,7 +1199,11 @@ def _unpackb3(s, **options): """ if not isinstance(s, (bytes, bytearray)): raise TypeError("packed data must be type 'bytes' or 'bytearray'") - return _unpack(io.BytesIO(s), options) + try: + return _unpack(io.BytesIO(s), options) if options.get('max_depth') is None \ + else _unpack_depth(io.BytesIO(s), options, 0) + except RecursionError: + raise MaximumDepthException("maximum nesting depth exceeded") ############################################################################# # Module Initialization @@ -1244,5 +1329,17 @@ def dst(self, dt): for code in range(0xe0, 0xff + 1): _unpack_dispatch_table[struct.pack("B", code)] = _unpack_integer + # Parallel dispatch table for the optional max_depth path. The array/map + # entries are the same container functions as the default table, wrapped by + # _limit_depth; scalar unpackers are adapted to accept (and ignore) the + # depth argument so _unpack() can dispatch uniformly with no per-code switch. + global _unpack_depth_dispatch_table + _unpack_depth_dispatch_table = {} + for code_bytes, unpacker in _unpack_dispatch_table.items(): + if unpacker is _unpack_array or unpacker is _unpack_map: + _unpack_depth_dispatch_table[code_bytes] = _limit_depth(unpacker) + else: + _unpack_depth_dispatch_table[code_bytes] = _ignore_depth(unpacker) + __init()