Add max_<type>_len option to unpacker. (fixes #97).

Fix build error on 32bit environment (fixes #102).

Author: methane

msgpack/_unpacker.pyx

@@ -28,6 +28,11 @@ cdef extern from "unpack.h":
         PyObject* ext_hook
         char *encoding
         char *unicode_errors
+        Py_ssize_t max_str_len
+        Py_ssize_t max_bin_len
+        Py_ssize_t max_array_len
+        Py_ssize_t max_map_len
+        Py_ssize_t max_ext_len
 
     ctypedef struct unpack_context:
         msgpack_user user
@@ -46,10 +51,18 @@ cdef extern from "unpack.h":
 cdef inline init_ctx(unpack_context *ctx,
                      object object_hook, object object_pairs_hook,
                      object list_hook, object ext_hook,
-                     bint use_list, char* encoding, char* unicode_errors):
+                     bint use_list, char* encoding, char* unicode_errors,
+                     Py_ssize_t max_str_len, Py_ssize_t max_bin_len,
+                     Py_ssize_t max_array_len, Py_ssize_t max_map_len,
+                     Py_ssize_t max_ext_len):
     unpack_init(ctx)
     ctx.user.use_list = use_list
     ctx.user.object_hook = ctx.user.list_hook = <PyObject*>NULL
+    ctx.user.max_str_len = max_str_len
+    ctx.user.max_bin_len = max_bin_len
+    ctx.user.max_array_len = max_array_len
+    ctx.user.max_map_len = max_map_len
+    ctx.user.max_ext_len = max_ext_len
 
     if object_hook is not None and object_pairs_hook is not None:
         raise TypeError("object_pairs_hook and object_hook are mutually exclusive.")
@@ -85,7 +98,12 @@ def default_read_extended_type(typecode, data):
 
 def unpackb(object packed, object object_hook=None, object list_hook=None,
             bint use_list=1, encoding=None, unicode_errors="strict",
-            object_pairs_hook=None, ext_hook=ExtType):
+            object_pairs_hook=None, ext_hook=ExtType,
+            Py_ssize_t max_str_len=2147483647, # 2**32-1
+            Py_ssize_t max_bin_len=2147483647,
+            Py_ssize_t max_array_len=2147483647,
+            Py_ssize_t max_map_len=2147483647,
+            Py_ssize_t max_ext_len=2147483647):
     """
     Unpack packed_bytes to object. Returns an unpacked object.
 
@@ -115,7 +133,8 @@ def unpackb(object packed, object object_hook=None, object list_hook=None,
         cerr = PyBytes_AsString(unicode_errors)
 
     init_ctx(&ctx, object_hook, object_pairs_hook, list_hook, ext_hook,
-             use_list, cenc, cerr)
+             use_list, cenc, cerr,
+             max_str_len, max_bin_len, max_array_len, max_map_len, max_ext_len)
     ret = unpack_construct(&ctx, buf, buf_len, &off)
     if ret == 1:
         obj = unpack_data(&ctx)
@@ -144,8 +163,7 @@ def unpack(object stream, object object_hook=None, object list_hook=None,
 
 
 cdef class Unpacker(object):
-    """
-    Streaming unpacker.
+    """Streaming unpacker.
 
     arguments:
 
@@ -183,6 +201,19 @@ cdef class Unpacker(object):
         Raises `BufferFull` exception when it is insufficient.
         You shoud set this parameter when unpacking data from untrasted source.
 
+    :param int max_str_len:
+        Limits max length of str. (default: 2**31-1)
+
+    :param int max_bin_len:
+        Limits max length of bin. (default: 2**31-1)
+
+    :param int max_array_len:
+        Limits max length of array. (default: 2**31-1)
+
+    :param int max_map_len:
+        Limits max length of map. (default: 2**31-1)
+
+
     example of streaming deserialize from file-like object::
 
         unpacker = Unpacker(file_like)
@@ -221,7 +252,12 @@ cdef class Unpacker(object):
     def __init__(self, file_like=None, Py_ssize_t read_size=0, bint use_list=1,
                  object object_hook=None, object object_pairs_hook=None, object list_hook=None,
                  str encoding=None, str unicode_errors='strict', int max_buffer_size=0,
-                 object ext_hook=ExtType):
+                 object ext_hook=ExtType,
+                 Py_ssize_t max_str_len=2147483647, # 2**32-1
+                 Py_ssize_t max_bin_len=2147483647,
+                 Py_ssize_t max_array_len=2147483647,
+                 Py_ssize_t max_map_len=2147483647,
+                 Py_ssize_t max_ext_len=2147483647):
         cdef char *cenc=NULL,
         cdef char *cerr=NULL
 
@@ -265,7 +301,9 @@ cdef class Unpacker(object):
             cerr = PyBytes_AsString(self.unicode_errors)
 
         init_ctx(&self.ctx, object_hook, object_pairs_hook, list_hook,
-                 ext_hook, use_list, cenc, cerr)
+                 ext_hook, use_list, cenc, cerr,
+                 max_str_len, max_bin_len, max_array_len,
+                 max_map_len, max_ext_len)
 
     def feed(self, object next_bytes):
         """Append `next_bytes` to internal buffer."""
@@ -365,7 +403,7 @@ cdef class Unpacker(object):
                 raise ValueError("Unpack failed: error = %d" % (ret,))
 
     def read_bytes(self, Py_ssize_t nbytes):
-        """read a specified number of raw bytes from the stream"""
+        """Read a specified number of raw bytes from the stream"""
         cdef size_t nread
         nread = min(self.buf_tail - self.buf_head, nbytes)
         ret = PyBytes_FromStringAndSize(self.buf + self.buf_head, nread)
@@ -375,8 +413,7 @@ cdef class Unpacker(object):
         return ret
 
     def unpack(self, object write_bytes=None):
-        """
-        unpack one object
+        """Unpack one object
 
         If write_bytes is not None, it will be called with parts of the raw
         message as it is unpacked.
@@ -386,8 +423,7 @@ cdef class Unpacker(object):
         return self._unpack(unpack_construct, write_bytes)
 
     def skip(self, object write_bytes=None):
-        """
-        read and ignore one object, returning None
+        """Read and ignore one object, returning None
 
         If write_bytes is not None, it will be called with parts of the raw
         message as it is unpacked.

msgpack/fallback.py

@@ -102,62 +102,84 @@ def unpackb(packed, **kwargs):
 
 
 class Unpacker(object):
-    """
-    Streaming unpacker.
+    """Streaming unpacker.
+
+    arguments:
 
-    `file_like` is a file-like object having a `.read(n)` method.
-    When `Unpacker` is initialized with a `file_like`, `.feed()` is not
-    usable.
+    :param file_like:
+        File-like object having `.read(n)` method.
+        If specified, unpacker reads serialized data from it and :meth:`feed()` is not usable.
 
-    `read_size` is used for `file_like.read(read_size)`.
+    :param int read_size:
+        Used as `file_like.read(read_size)`. (default: `min(1024**2, max_buffer_size)`)
 
-    If `use_list` is True (default), msgpack lists are deserialized to Python
-    lists.  Otherwise they are deserialized to tuples.
+    :param bool use_list:
+        If true, unpack msgpack array to Python list.
+        Otherwise, unpack to Python tuple. (default: True)
 
-    `object_hook` is the same as in simplejson.  If it is not None, it should
-    be callable and Unpacker calls it with a dict argument after deserializing
-    a map.
+    :param callable object_hook:
+        When specified, it should be callable.
+        Unpacker calls it with a dict argument after unpacking msgpack map.
+        (See also simplejson)
 
-    `object_pairs_hook` is the same as in simplejson.  If it is not None, it
-    should be callable and Unpacker calls it with a list of key-value pairs
-    after deserializing a map.
+    :param callable object_pairs_hook:
+        When specified, it should be callable.
+        Unpacker calls it with a list of key-value pairs after unpacking msgpack map.
+        (See also simplejson)
 
-    `ext_hook` is callback for ext (User defined) type. It called with two
-    arguments: (code, bytes). default: `msgpack.ExtType`
+    :param str encoding:
+        Encoding used for decoding msgpack raw.
+        If it is None (default), msgpack raw is deserialized to Python bytes.
 
-    `encoding` is the encoding used for decoding msgpack bytes.  If it is
-    None (default), msgpack bytes are deserialized to Python bytes.
+    :param str unicode_errors:
+        Used for decoding msgpack raw with *encoding*.
+        (default: `'strict'`)
 
-    `unicode_errors` is used for decoding bytes.
+    :param int max_buffer_size:
+        Limits size of data waiting unpacked.  0 means system's INT_MAX (default).
+        Raises `BufferFull` exception when it is insufficient.
+        You shoud set this parameter when unpacking data from untrasted source.
 
-    `max_buffer_size` limits the buffer size.  0 means INT_MAX (default).
+    :param int max_str_len:
+        Limits max length of str. (default: 2**31-1)
 
-    Raises `BufferFull` exception when it is unsufficient.
+    :param int max_bin_len:
+        Limits max length of bin. (default: 2**31-1)
 
-    You should set this parameter when unpacking data from an untrustred source.
+    :param int max_array_len:
+        Limits max length of array. (default: 2**31-1)
 
-    example of streaming deserialization from file-like object::
+    :param int max_map_len:
+        Limits max length of map. (default: 2**31-1)
+
+
+    example of streaming deserialize from file-like object::
 
         unpacker = Unpacker(file_like)
         for o in unpacker:
-            do_something(o)
+            process(o)
 
-    example of streaming deserialization from socket::
+    example of streaming deserialize from socket::
 
         unpacker = Unpacker()
-        while 1:
-            buf = sock.recv(1024*2)
+        while True:
+            buf = sock.recv(1024**2)
             if not buf:
                 break
             unpacker.feed(buf)
             for o in unpacker:
-                do_something(o)
+                process(o)
     """
 
     def __init__(self, file_like=None, read_size=0, use_list=True,
                  object_hook=None, object_pairs_hook=None, list_hook=None,
                  encoding=None, unicode_errors='strict', max_buffer_size=0,
-                 ext_hook=ExtType):
+                 ext_hook=ExtType,
+                 max_str_len=2147483647, # 2**32-1
+                 max_bin_len=2147483647,
+                 max_array_len=2147483647,
+                 max_map_len=2147483647,
+                 max_ext_len=2147483647):
         if file_like is None:
             self._fb_feeding = True
         else:
@@ -185,6 +207,11 @@ def __init__(self, file_like=None, read_size=0, use_list=True,
         self._object_hook = object_hook
         self._object_pairs_hook = object_pairs_hook
         self._ext_hook = ext_hook
+        self._max_str_len = max_str_len
+        self._max_bin_len = max_bin_len
+        self._max_array_len = max_array_len
+        self._max_map_len = max_map_len
+        self._max_ext_len = max_ext_len
 
         if list_hook is not None and not callable(list_hook):
             raise TypeError('`list_hook` is not callable')
@@ -316,12 +343,18 @@ def _read_header(self, execute=EX_CONSTRUCT, write_bytes=None):
             n = b & 0b00011111
             obj = self._fb_read(n, write_bytes)
             typ = TYPE_RAW
+            if n > self._max_str_len:
+                raise ValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
         elif b & 0b11110000 == 0b10010000:
             n = b & 0b00001111
             typ = TYPE_ARRAY
+            if n > self._max_array_len:
+                raise ValueError("%s exceeds max_array_len(%s)", n, self._max_array_len)
         elif b & 0b11110000 == 0b10000000:
             n = b & 0b00001111
             typ = TYPE_MAP
+            if n > self._max_map_len:
+                raise ValueError("%s exceeds max_map_len(%s)", n, self._max_map_len)
         elif b == 0xc0:
             obj = None
         elif b == 0xc2:
@@ -331,26 +364,38 @@ def _read_header(self, execute=EX_CONSTRUCT, write_bytes=None):
         elif b == 0xc4:
             typ = TYPE_BIN
             n = struct.unpack("B", self._fb_read(1, write_bytes))[0]
+            if n > self._max_bin_len:
+                raise ValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
             obj = self._fb_read(n, write_bytes)
         elif b == 0xc5:
             typ = TYPE_BIN
             n = struct.unpack(">H", self._fb_read(2, write_bytes))[0]
+            if n > self._max_bin_len:
+                raise ValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
             obj = self._fb_read(n, write_bytes)
         elif b == 0xc6:
             typ = TYPE_BIN
             n = struct.unpack(">I", self._fb_read(4, write_bytes))[0]
+            if n > self._max_bin_len:
+                raise ValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
             obj = self._fb_read(n, write_bytes)
         elif b == 0xc7:  # ext 8
             typ = TYPE_EXT
             L, n = struct.unpack('Bb', self._fb_read(2, write_bytes))
+            if L > self._max_ext_len:
+                raise ValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
             obj = self._fb_read(L, write_bytes)
         elif b == 0xc8:  # ext 16
             typ = TYPE_EXT
             L, n = struct.unpack('>Hb', self._fb_read(3, write_bytes))
+            if L > self._max_ext_len:
+                raise ValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
             obj = self._fb_read(L, write_bytes)
         elif b == 0xc9:  # ext 32
             typ = TYPE_EXT
             L, n = struct.unpack('>Ib', self._fb_read(5, write_bytes))
+            if L > self._max_ext_len:
+                raise ValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
             obj = self._fb_read(L, write_bytes)
         elif b == 0xca:
             obj = struct.unpack(">f", self._fb_read(4, write_bytes))[0]
@@ -374,42 +419,66 @@ def _read_header(self, execute=EX_CONSTRUCT, write_bytes=None):
             obj = struct.unpack(">q", self._fb_read(8, write_bytes))[0]
         elif b == 0xd4:  # fixext 1
             typ = TYPE_EXT
+            if self._max_ext_len < 1:
+                raise ValueError("%s exceeds max_ext_len(%s)" % (1, self._max_ext_len))
             n, obj = struct.unpack('b1s', self._fb_read(2, write_bytes))
         elif b == 0xd5:  # fixext 2
             typ = TYPE_EXT
+            if self._max_ext_len < 2:
+                raise ValueError("%s exceeds max_ext_len(%s)" % (2, self._max_ext_len))
             n, obj = struct.unpack('b2s', self._fb_read(3, write_bytes))
         elif b == 0xd6:  # fixext 4
             typ = TYPE_EXT
+            if self._max_ext_len < 4:
+                raise ValueError("%s exceeds max_ext_len(%s)" % (4, self._max_ext_len))
             n, obj = struct.unpack('b4s', self._fb_read(5, write_bytes))
         elif b == 0xd7:  # fixext 8
             typ = TYPE_EXT
+            if self._max_ext_len < 8:
+                raise ValueError("%s exceeds max_ext_len(%s)" % (8, self._max_ext_len))
             n, obj = struct.unpack('b8s', self._fb_read(9, write_bytes))
         elif b == 0xd8:  # fixext 16
             typ = TYPE_EXT
+            if self._max_ext_len < 16:
+                raise ValueError("%s exceeds max_ext_len(%s)" % (16, self._max_ext_len))
             n, obj = struct.unpack('b16s', self._fb_read(17, write_bytes))
         elif b == 0xd9:
             typ = TYPE_RAW
             n = struct.unpack("B", self._fb_read(1, write_bytes))[0]
+            if n > self._max_str_len:
+                raise ValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
             obj = self._fb_read(n, write_bytes)
         elif b == 0xda:
             typ = TYPE_RAW
             n = struct.unpack(">H", self._fb_read(2, write_bytes))[0]
+            if n > self._max_str_len:
+                raise ValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
             obj = self._fb_read(n, write_bytes)
         elif b == 0xdb:
             typ = TYPE_RAW
             n = struct.unpack(">I", self._fb_read(4, write_bytes))[0]
+            if n > self._max_str_len:
+                raise ValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
             obj = self._fb_read(n, write_bytes)
         elif b == 0xdc:
             n = struct.unpack(">H", self._fb_read(2, write_bytes))[0]
+            if n > self._max_array_len:
+                raise ValueError("%s exceeds max_array_len(%s)", n, self._max_array_len)
             typ = TYPE_ARRAY
         elif b == 0xdd:
             n = struct.unpack(">I", self._fb_read(4, write_bytes))[0]
+            if n > self._max_array_len:
+                raise ValueError("%s exceeds max_array_len(%s)", n, self._max_array_len)
             typ = TYPE_ARRAY
         elif b == 0xde:
             n = struct.unpack(">H", self._fb_read(2, write_bytes))[0]
+            if n > self._max_map_len:
+                raise ValueError("%s exceeds max_map_len(%s)", n, self._max_map_len)
             typ = TYPE_MAP
         elif b == 0xdf:
             n = struct.unpack(">I", self._fb_read(4, write_bytes))[0]
+            if n > self._max_map_len:
+                raise ValueError("%s exceeds max_map_len(%s)", n, self._max_map_len)
             typ = TYPE_MAP
         else:
             raise UnpackValueError("Unknown header: 0x%x" % b)