needforspeed: check for overflow in replace (from Andrew Dalke)

3 files changed, 54 insertions, 15 deletions

diff --git a/Lib/test/string_tests.py b/Lib/test/string_tests.py
index f76a9ebe08..630618c317 100644
--- a/Lib/test/string_tests.py
+++ b/Lib/test/string_tests.py
@@ -555,15 +555,14 @@ class CommonTest(unittest.TestCase):
         self.checkraises(TypeError, 'hello', 'replace', 42, 'h')
         self.checkraises(TypeError, 'hello', 'replace', 'h', 42)
 
-### Commented out until the underlying libraries are fixed
-##    def test_replace_overflow(self):
-##        # Check for overflow checking on 32 bit machines
-##        if sys.maxint != 2147483647:
-##            return
-##        A2_16 = "A" * (2**16)
-##        self.checkraises(OverflowError, A2_16, "replace", "", A2_16)
-##        self.checkraises(OverflowError, A2_16, "replace", "A", A2_16)
-##        self.checkraises(OverflowError, A2_16, "replace", "AA", A2_16+A2_16)
+    def test_replace_overflow(self):
+        # Check for overflow checking on 32 bit machines
+        if sys.maxint != 2147483647:
+            return
+        A2_16 = "A" * (2**16)
+        self.checkraises(OverflowError, A2_16, "replace", "", A2_16)
+        self.checkraises(OverflowError, A2_16, "replace", "A", A2_16)
+        self.checkraises(OverflowError, A2_16, "replace", "AA", A2_16+A2_16)
 
     def test_zfill(self):
         self.checkequal('123', '123', 'zfill', 2)
diff --git a/Objects/stringobject.c b/Objects/stringobject.c
index 7bddeaa99e..77796dfd90 100644
--- a/Objects/stringobject.c
+++ b/Objects/stringobject.c
@@ -2460,6 +2460,7 @@ mymemreplace(const char *str, Py_ssize_t len,		/* input string */
 	char *out_s;
 	char *new_s;
 	Py_ssize_t nfound, offset, new_len;
+        Py_ssize_t product, delta;
 
 	if (len == 0 || (pat_len == 0 && sub_len == 0) || pat_len > len)
 		goto return_same;
@@ -2473,7 +2474,24 @@ mymemreplace(const char *str, Py_ssize_t len,		/* input string */
 	if (nfound == 0)
 		goto return_same;
 
-	new_len = len + nfound*(sub_len - pat_len);
+        delta = (sub_len - pat_len);
+        if (delta == 0) {
+            new_len = len;
+        } else {
+            product = nfound * (sub_len - pat_len);
+            if ((product / (sub_len - pat_len)) != nfound) {
+                PyErr_SetString(PyExc_OverflowError,
+                                "replace string is too long");
+                return NULL;
+            }
+            new_len = len + product;
+            if (new_len < 0) {
+                PyErr_SetString(PyExc_OverflowError,
+                                "replace string is too long");
+                return NULL;
+            }
+        }
+
 	if (new_len == 0) {
 		/* Have to allocate something for the caller to free(). */
 		out_s = (char *)PyMem_MALLOC(1);
@@ -2578,7 +2596,8 @@ string_replace(PyStringObject *self, PyObject *args)
 
 	new_s = mymemreplace(str,len,sub,sub_len,repl,repl_len,count,&out_len);
 	if (new_s == NULL) {
-		PyErr_NoMemory();
+		if (!PyErr_Occurred())
+			PyErr_NoMemory();
 		return NULL;
 	}
 	if (out_len == -1) {
diff --git a/Objects/unicodeobject.c b/Objects/unicodeobject.c
index 70890736ae..bcf2c38345 100644
--- a/Objects/unicodeobject.c
+++ b/Objects/unicodeobject.c
@@ -3866,9 +3866,11 @@ int PyUnicode_EncodeDecimal(Py_UNICODE *s,
    for some more background, see: http://effbot.org/stringlib */
 
 /* note: fastsearch may access s[n], which isn't a problem when using
-   Python's ordinary string types.  also, the count mode returns -1 if
-   there cannot possible be a match in the target string, and 0 if it
-   has actually checked for matches. */
+   Python's ordinary string types, but may cause problems if you're
+   using this code in other contexts.  also, the count mode returns -1
+   if there cannot possible be a match in the target string, and 0 if
+   it has actually checked for matches, but didn't find any.  callers
+   beware! */
 
 #define FAST_COUNT 0
 #define FAST_SEARCH 1
@@ -4862,6 +4864,7 @@ PyObject *replace(PyUnicodeObject *self,
     } else {
 
         Py_ssize_t n, i;
+        Py_ssize_t product, new_size, delta;
         Py_UNICODE *p;
 
         /* replace strings */
@@ -4870,7 +4873,25 @@ PyObject *replace(PyUnicodeObject *self,
             n = maxcount;
         if (n == 0)
             goto nothing;
-        u = _PyUnicode_New(self->length + n * (str2->length - str1->length));
+        /* new_size = self->length + n * (str2->length - str1->length)); */
+        delta = (str2->length - str1->length);
+        if (delta == 0) {
+            new_size = self->length;
+        } else {
+            product = n * (str2->length - str1->length);
+            if ((product / (str2->length - str1->length)) != n) {
+                PyErr_SetString(PyExc_OverflowError,
+                                "replace string is too long");
+                return NULL;
+            }
+            new_size = self->length + product;
+            if (new_size < 0) {
+                PyErr_SetString(PyExc_OverflowError,
+                                "replace string is too long");
+                return NULL;
+            }
+        }
+        u = _PyUnicode_New(new_size);
         if (!u)
             return NULL;
         i = 0;