Handling too large inputs gracefully

setup.py

@@ -1,10 +1,10 @@
 #!/usr/bin/env python
 
-"""Setup file for veezio backend"""
+"""Setup file for lz4 backend"""
 
 from setuptools import setup, find_packages, Extension
 
-VERSION = (0, 4, 0)
+VERSION = (0, 4, 1)
 
 setup(
     name='lz4',

src/lz4.c

@@ -327,6 +327,10 @@ inline static int LZ4_NbCommonBytes (register U32 val)
 
 int LZ4_compressBound(int isize)
 {
+	static const int max_size = 2139095024;
+	if (isize < 0 || isize > max_size) {
+		return -1;
+	}
 	return (isize + (isize/255) + 16);
 }
 

src/python-lz4.c

@@ -1,3 +1,5 @@
+#define PY_SSIZE_T_CLEAN
+
 #include "stdlib.h"
 #include "math.h"
 #include "Python.h"
@@ -30,14 +32,25 @@ compress_with(compressor compress, PyObject *self, PyObject *args)
 {
     PyObject *result;
     const char *source;
-    int source_size;
+    size_t source_size;
     char *dest;
     int dest_size;
 
     if (!PyArg_ParseTuple(args, "s#", &source, &source_size))
         return NULL;
 
-    dest_size = hdr_size + LZ4_compressBound(source_size);
+    if (source_size > INT_MAX) {
+	PyErr_SetString(PyExc_ValueError, "input too long");
+	return NULL;
+    }
+
+    dest_size = LZ4_compressBound(source_size);
+    if (dest_size < 0 || (size_t)dest_size + hdr_size > INT_MAX) {
+	PyErr_SetString(PyExc_ValueError, "input too long");
+	return NULL;
+    }
+
+    dest_size += hdr_size;
     result = PyString_FromStringAndSize(NULL, dest_size);
     if (result == NULL)
 	return NULL;
@@ -70,7 +83,7 @@ py_lz4_uncompress(PyObject *self, PyObject *args)
 {
     PyObject *result;
     const char *source;
-    int source_size;
+    size_t source_size;
     uint32_t dest_size;
 
     if (!PyArg_ParseTuple(args, "s#", &source, &source_size))

tests/test.py

@@ -2,4 +2,20 @@
 import sys
 
 DATA = open("/dev/urandom", "rb").read(128 * 1024)  # Read 128kb
-sys.exit(DATA != lz4.loads(lz4.dumps(DATA)) and 1 or 0)
+if DATA != lz4.loads(lz4.dumps(DATA)):
+    sys.exit(1)
+
+# max size
+DATA = "x" * 2139095020
+if DATA != lz4.loads(lz4.dumps(DATA)):
+    sys.exit(1)
+
+# max size + 1
+DATA = DATA + "x"
+try:
+    lz4.dumps(DATA)
+    sys.exit(1)
+except ValueError:
+    pass
+
+sys.exit(0)