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

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

2 files changed, 76 insertions, 24 deletions

diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx
index 16de40f..f69be91 100644
--- a/msgpack/_unpacker.pyx
+++ b/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)
@@ -183,6 +202,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 +253,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 +302,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 +404,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 +414,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 +424,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.
diff --git a/msgpack/unpack.h b/msgpack/unpack.h
index 71142c6..b3bddf2 100644
--- a/msgpack/unpack.h
+++ b/msgpack/unpack.h
@@ -27,6 +27,7 @@ typedef struct unpack_user {
     PyObject *ext_hook;
     const char *encoding;
     const char *unicode_errors;
+    Py_ssize_t max_str_len, max_bin_len, max_array_len, max_map_len, max_ext_len;
 } unpack_user;
 
 typedef PyObject* msgpack_unpack_object;
@@ -55,15 +56,7 @@ static inline int unpack_callback_uint8(unpack_user* u, uint8_t d, msgpack_unpac
 
 static inline int unpack_callback_uint32(unpack_user* u, uint32_t d, msgpack_unpack_object* o)
 {
-    PyObject *p;
-#if UINT32_MAX > LONG_MAX
-    if (d > LONG_MAX) {
-        p = PyLong_FromUnsignedLong((unsigned long)d);
-    } else
-#endif
-    {
-        p = PyInt_FromUnsignedLong((long)d);
-    }
+    PyObject *p = PyInt_FromSize_t((size_t)d);
     if (!p)
         return -1;
     *o = p;
@@ -76,7 +69,7 @@ static inline int unpack_callback_uint64(unpack_user* u, uint64_t d, msgpack_unp
     if (d > LONG_MAX) {
         p = PyLong_FromUnsignedLongLong((unsigned PY_LONG_LONG)d);
     } else {
-        p = PyInt_FromLong((long)d);
+        p = PyInt_FromSize_t((size_t)d);
     }
     if (!p)
         return -1;
@@ -140,6 +133,10 @@ static inline int unpack_callback_false(unpack_user* u, msgpack_unpack_object* o
 
 static inline int unpack_callback_array(unpack_user* u, unsigned int n, msgpack_unpack_object* o)
 {
+    if (n > u->max_array_len) {
+        PyErr_Format(PyExc_ValueError, "%u exceeds max_array_len(%zd)", n, u->max_array_len);
+        return -1;
+    }
     PyObject *p = u->use_list ? PyList_New(n) : PyTuple_New(n);
 
     if (!p)
@@ -171,6 +168,10 @@ static inline int unpack_callback_array_end(unpack_user* u, msgpack_unpack_objec
 
 static inline int unpack_callback_map(unpack_user* u, unsigned int n, msgpack_unpack_object* o)
 {
+    if (n > u->max_map_len) {
+        PyErr_Format(PyExc_ValueError, "%u exceeds max_map_len(%zd)", n, u->max_map_len);
+        return -1;
+    }
     PyObject *p;
     if (u->has_pairs_hook) {
         p = PyList_New(n); // Or use tuple?
@@ -218,6 +219,11 @@ static inline int unpack_callback_map_end(unpack_user* u, msgpack_unpack_object*
 
 static inline int unpack_callback_raw(unpack_user* u, const char* b, const char* p, unsigned int l, msgpack_unpack_object* o)
 {
+    if (l > u->max_str_len) {
+        PyErr_Format(PyExc_ValueError, "%u exceeds max_str_len(%zd)", l, u->max_str_len);
+        return -1;
+    }
+
     PyObject *py;
     if(u->encoding) {
         py = PyUnicode_Decode(p, l, u->encoding, u->unicode_errors);
@@ -232,6 +238,11 @@ static inline int unpack_callback_raw(unpack_user* u, const char* b, const char*
 
 static inline int unpack_callback_bin(unpack_user* u, const char* b, const char* p, unsigned int l, msgpack_unpack_object* o)
 {
+    if (l > u->max_bin_len) {
+        PyErr_Format(PyExc_ValueError, "%u exceeds max_bin_len(%zd)", l, u->max_bin_len);
+        return -1;
+    }
+
     PyObject *py = PyBytes_FromStringAndSize(p, l);
     if (!py)
         return -1;
@@ -240,7 +251,7 @@ static inline int unpack_callback_bin(unpack_user* u, const char* b, const char*
 }
 
 static inline int unpack_callback_ext(unpack_user* u, const char* base, const char* pos,
-                                      unsigned int lenght, msgpack_unpack_object* o)
+                                      unsigned int length, msgpack_unpack_object* o)
 {
     PyObject *py;
     int8_t typecode = (int8_t)*pos++;
@@ -248,11 +259,15 @@ static inline int unpack_callback_ext(unpack_user* u, const char* base, const ch
         PyErr_SetString(PyExc_AssertionError, "u->ext_hook cannot be NULL");
         return -1;
     }
-    // length also includes the typecode, so the actual data is lenght-1
+    if (length > u->max_ext_len) {
+        PyErr_Format(PyExc_ValueError, "%u exceeds max_ext_len(%zd)", length, u->max_ext_len);
+        return -1;
+    }
+    // length also includes the typecode, so the actual data is length-1
 #if PY_MAJOR_VERSION == 2
-    py = PyObject_CallFunction(u->ext_hook, "(is#)", typecode, pos, lenght-1);
+    py = PyObject_CallFunction(u->ext_hook, "(is#)", typecode, pos, length-1);
 #else
-    py = PyObject_CallFunction(u->ext_hook, "(iy#)", typecode, pos, lenght-1);
+    py = PyObject_CallFunction(u->ext_hook, "(iy#)", typecode, pos, length-1);
 #endif
     if (!py)
         return -1;