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

diff --git a/Lib/test/string_tests.py b/Lib/test/string_tests.py
index f76a9ebe08546648edb8c7f4c3b88fef571083e7..630618c317d8c8520ab6e1bc8ab524ece5e51cc8 100644 (file)
--- 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 7bddeaa99ea12fb4e7bae0d1b346a0aa81a37a32..77796dfd9024ed4b70ca1268112670b4b92de2a3 100644 (file)
--- 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 70890736ae7ab01ae0e376116423b26d5bad12e3..bcf2c38345b35d2452981e865ed5d60aa0dd0c13 100644 (file)
--- 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;