From 7bce869dd03e1db4e7c828fde6e2187e0822b07d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jarno=20Sepp=C3=A4nen?= <jarno.seppanen@nokia.com>
Date: Wed, 20 Jun 2012 10:13:16 +0300
Subject: [PATCH] Bugfix: discard oversize inputs instead of truncating
 silently

---
 setup.py         |  4 ++--
 src/lz4.c        |  4 ++++
 src/python-lz4.c | 19 ++++++++++++++++---
 tests/test.py    | 18 +++++++++++++++++-
 4 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/setup.py b/setup.py
index 228e493..4df7f36 100644
--- a/setup.py
+++ b/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',
diff --git a/src/lz4.c b/src/lz4.c
index 06e2829..138e1b1 100644
--- a/src/lz4.c
+++ b/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);
 }
 
diff --git a/src/python-lz4.c b/src/python-lz4.c
index 2a9b537..e591187 100644
--- a/src/python-lz4.c
+++ b/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))
diff --git a/tests/test.py b/tests/test.py
index 4638605..f788b6f 100644
--- a/tests/test.py
+++ b/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)