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

2 files changed, 46 insertions, 6 deletions

diff --git a/Objects/stringobject.c b/Objects/stringobject.c
index 7bddeaa..77796df 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 7089073..bcf2c38 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;