Merge pull request #94 from msgpack/strict-input-check

Add tests for limits.

3 files changed, 92 insertions, 4 deletions

diff --git a/msgpack/_packer.pyx b/msgpack/_packer.pyx
index f261f08..82e4a63 100644
--- a/msgpack/_packer.pyx
+++ b/msgpack/_packer.pyx
@@ -135,6 +135,8 @@ cdef class Packer(object):
                 else:
                     ret = msgpack_pack_false(&self.pk)
             elif PyLong_Check(o):
+                # PyInt_Check(long) is True for Python 3.
+                # Sow we should test long before int.
                 if o > 0:
                     ullval = o
                     ret = msgpack_pack_unsigned_long_long(&self.pk, ullval)
@@ -152,8 +154,10 @@ cdef class Packer(object):
                    dval = o
                    ret = msgpack_pack_double(&self.pk, dval)
             elif PyBytes_Check(o):
-                rawval = o
                 L = len(o)
+                if L > (2**32)-1:
+                    raise ValueError("bytes is too large")
+                rawval = o
                 ret = msgpack_pack_bin(&self.pk, L)
                 if ret == 0:
                     ret = msgpack_pack_raw_body(&self.pk, rawval, L)
@@ -161,13 +165,19 @@ cdef class Packer(object):
                 if not self.encoding:
                     raise TypeError("Can't encode unicode string: no encoding is specified")
                 o = PyUnicode_AsEncodedString(o, self.encoding, self.unicode_errors)
+                L = len(o)
+                if L > (2**32)-1:
+                    raise ValueError("dict is too large")
                 rawval = o
                 ret = msgpack_pack_raw(&self.pk, len(o))
                 if ret == 0:
                     ret = msgpack_pack_raw_body(&self.pk, rawval, len(o))
             elif PyDict_CheckExact(o):
                 d = <dict>o
-                ret = msgpack_pack_map(&self.pk, len(d))
+                L = len(d)
+                if L > (2**32)-1:
+                    raise ValueError("dict is too large")
+                ret = msgpack_pack_map(&self.pk, L)
                 if ret == 0:
                     for k, v in d.iteritems():
                         ret = self._pack(k, nest_limit-1)
@@ -175,7 +185,10 @@ cdef class Packer(object):
                         ret = self._pack(v, nest_limit-1)
                         if ret != 0: break
             elif PyDict_Check(o):
-                ret = msgpack_pack_map(&self.pk, len(o))
+                L = len(o)
+                if L > (2**32)-1:
+                    raise ValueError("dict is too large")
+                ret = msgpack_pack_map(&self.pk, L)
                 if ret == 0:
                     for k, v in o.items():
                         ret = self._pack(k, nest_limit-1)
@@ -187,10 +200,15 @@ cdef class Packer(object):
                 longval = o.code
                 rawval = o.data
                 L = len(o.data)
+                if L > (2**32)-1:
+                    raise ValueError("EXT data is too large")
                 ret = msgpack_pack_ext(&self.pk, longval, L)
                 ret = msgpack_pack_raw_body(&self.pk, rawval, L)
             elif PyTuple_Check(o) or PyList_Check(o):
-                ret = msgpack_pack_array(&self.pk, len(o))
+                L = len(o)
+                if L > (2**32)-1:
+                    raise ValueError("list is too large")
+                ret = msgpack_pack_array(&self.pk, L)
                 if ret == 0:
                     for v in o:
                         ret = self._pack(v, nest_limit-1)
@@ -220,6 +238,8 @@ cdef class Packer(object):
         msgpack_pack_raw_body(&self.pk, data, len(data))
 
     def pack_array_header(self, size_t size):
+        if size > (2**32-1):
+            raise ValueError
         cdef int ret = msgpack_pack_array(&self.pk, size)
         if ret == -1:
             raise MemoryError
@@ -231,6 +251,8 @@ cdef class Packer(object):
             return buf
 
     def pack_map_header(self, size_t size):
+        if size > (2**32-1):
+            raise ValueError
         cdef int ret = msgpack_pack_map(&self.pk, size)
         if ret == -1:
             raise MemoryError
diff --git a/msgpack/fallback.py b/msgpack/fallback.py
index 55de72c..49323e6 100644
--- a/msgpack/fallback.py
+++ b/msgpack/fallback.py
@@ -633,6 +633,8 @@ class Packer(object):
         return ret
 
     def pack_array_header(self, n):
+        if n >= 2**32:
+            raise ValueError
         self._fb_pack_array_header(n)
         ret = self._buffer.getvalue()
         if self._autoreset:
@@ -642,6 +644,8 @@ class Packer(object):
         return ret
 
     def pack_map_header(self, n):
+        if n >= 2**32:
+            raise ValueError
         self._fb_pack_map_header(n)
         ret = self._buffer.getvalue()
         if self._autoreset:
diff --git a/test/test_limits.py b/test/test_limits.py
new file mode 100644
index 0000000..da8cd2b
--- /dev/null
+++ b/test/test_limits.py
@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+# coding: utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+import pytest
+
+from msgpack import packb, unpackb, Packer
+
+
+def test_integer():
+    x = -(2 ** 63)
+    assert unpackb(packb(x)) == x
+    with pytest.raises((OverflowError, ValueError)):
+        packb(x-1)
+
+    x = 2 ** 64 - 1
+    assert unpackb(packb(x)) == x
+    with pytest.raises((OverflowError, ValueError)):
+        packb(x+1)
+
+
+def test_array_header():
+    packer = Packer()
+    packer.pack_array_header(2**32-1)
+    with pytest.raises(ValueError):
+        packer.pack_array_header(2**32)
+
+
+def test_map_header():
+    packer = Packer()
+    packer.pack_map_header(2**32-1)
+    with pytest.raises(ValueError):
+        packer.pack_array_header(2**32)
+
+
+# PyPy fails following tests because of constant folding?
+# https://bugs.pypy.org/issue1721
+#@pytest.mark.skipif(True, reason="Requires very large memory.")
+#def test_binary():
+#    x = b'x' * (2**32 - 1)
+#    assert unpackb(packb(x)) == x
+#    del x
+#    x = b'x' * (2**32)
+#    with pytest.raises(ValueError):
+#        packb(x)
+#
+#
+#@pytest.mark.skipif(True, reason="Requires very large memory.")
+#def test_string():
+#    x = 'x' * (2**32 - 1)
+#    assert unpackb(packb(x)) == x
+#    x += 'y'
+#    with pytest.raises(ValueError):
+#        packb(x)
+#
+#
+#@pytest.mark.skipif(True, reason="Requires very large memory.")
+#def test_array():
+#    x = [0] * (2**32 - 1)
+#    assert unpackb(packb(x)) == x
+#    x.append(0)
+#    with pytest.raises(ValueError):
+#        packb(x)