Merge pull request #7456 from jaimefrio/reshape_segfault

BUG: int overflow in reshape, fixes #7455, fixes #7293

2 files changed, 52 insertions, 56 deletions

diff --git a/numpy/core/src/multiarray/shape.c b/numpy/core/src/multiarray/shape.c
index a307bed9c..2d8b6f602 100644
--- a/numpy/core/src/multiarray/shape.c
+++ b/numpy/core/src/multiarray/shape.c
@@ -17,8 +17,11 @@
 
 #include "shape.h"
 
+#include "templ_common.h" /* for npy_mul_with_overflow_intp */
+#include "common.h" /* for convert_shape_to_string */
+
 static int
-_fix_unknown_dimension(PyArray_Dims *newshape, npy_intp s_original);
+_fix_unknown_dimension(PyArray_Dims *newshape, PyArrayObject *arr);
 
 static int
 _attempt_nocopy_reshape(PyArrayObject *self, int newnd, npy_intp* newdims,
@@ -206,7 +209,7 @@ PyArray_Newshape(PyArrayObject *self, PyArray_Dims *newdims,
     /*
      * fix any -1 dimensions and check new-dimensions against old size
      */
-    if (_fix_unknown_dimension(newdims, PyArray_SIZE(self)) < 0) {
+    if (_fix_unknown_dimension(newdims, self) < 0) {
         return NULL;
     }
     /*
@@ -339,8 +342,10 @@ _putzero(char *optr, PyObject *zero, PyArray_Descr *dtype)
 /*
  * attempt to reshape an array without copying data
  *
- * This function should correctly handle all reshapes, including
- * axes of length 1. Zero strides should work but are untested.
+ * The requested newdims are not checked, but must be compatible with
+ * the size of self, which must be non-zero. Other than that this
+ * function should correctly handle all reshapes, including axes of
+ * length 1. Zero strides should work but are untested.
  *
  * If a copy is needed, returns 0
  * If no copy is needed, returns 1 and fills newstrides
@@ -361,7 +366,7 @@ _attempt_nocopy_reshape(PyArrayObject *self, int newnd, npy_intp* newdims,
     int oldnd;
     npy_intp olddims[NPY_MAXDIMS];
     npy_intp oldstrides[NPY_MAXDIMS];
-    npy_intp np, op, last_stride;
+    npy_intp last_stride;
     int oi, oj, ok, ni, nj, nk;
 
     oldnd = 0;
@@ -377,43 +382,14 @@ _attempt_nocopy_reshape(PyArrayObject *self, int newnd, npy_intp* newdims,
         }
     }
 
-    /*
-      fprintf(stderr, "_attempt_nocopy_reshape( (");
-      for (oi=0; oi<oldnd; oi++)
-      fprintf(stderr, "(%d,%d), ", olddims[oi], oldstrides[oi]);
-      fprintf(stderr, ") -> (");
-      for (ni=0; ni<newnd; ni++)
-      fprintf(stderr, "(%d,*), ", newdims[ni]);
-      fprintf(stderr, "), is_f_order=%d)\n", is_f_order);
-    */
-
-
-    np = 1;
-    for (ni = 0; ni < newnd; ni++) {
-        np *= newdims[ni];
-    }
-    op = 1;
-    for (oi = 0; oi < oldnd; oi++) {
-        op *= olddims[oi];
-    }
-    if (np != op) {
-        /* different total sizes; no hope */
-        return 0;
-    }
-
-    if (np == 0) {
-        /* the current code does not handle 0-sized arrays, so give up */
-        return 0;
-    }
-
     /* oi to oj and ni to nj give the axis ranges currently worked with */
     oi = 0;
     oj = 1;
     ni = 0;
     nj = 1;
     while (ni < newnd && oi < oldnd) {
-        np = newdims[ni];
-        op = olddims[oi];
+        npy_intp np = newdims[ni];
+        npy_intp op = olddims[oi];
 
         while (np != op) {
             if (np < op) {
@@ -475,26 +451,30 @@ _attempt_nocopy_reshape(PyArrayObject *self, int newnd, npy_intp* newdims,
         newstrides[nk] = last_stride;
     }
 
-    /*
-      fprintf(stderr, "success: _attempt_nocopy_reshape (");
-      for (oi=0; oi<oldnd; oi++)
-      fprintf(stderr, "(%d,%d), ", olddims[oi], oldstrides[oi]);
-      fprintf(stderr, ") -> (");
-      for (ni=0; ni<newnd; ni++)
-      fprintf(stderr, "(%d,%d), ", newdims[ni], newstrides[ni]);
-      fprintf(stderr, ")\n");
-    */
-
     return 1;
 }
 
+static void
+raise_reshape_size_mismatch(PyArray_Dims *newshape, PyArrayObject *arr)
+{
+    PyObject *msg = PyUString_FromFormat("cannot reshape array of size %zd "
+                                         "into shape ", PyArray_SIZE(arr));
+    PyObject *tmp = convert_shape_to_string(newshape->len, newshape->ptr, "");
+
+    PyUString_ConcatAndDel(&msg, tmp);
+    if (msg != NULL) {
+        PyErr_SetObject(PyExc_ValueError, msg);
+        Py_DECREF(msg);
+    }
+}
+
 static int
-_fix_unknown_dimension(PyArray_Dims *newshape, npy_intp s_original)
+_fix_unknown_dimension(PyArray_Dims *newshape, PyArrayObject *arr)
 {
     npy_intp *dimensions;
+    npy_intp s_original = PyArray_SIZE(arr);
     npy_intp i_unknown, s_known;
     int i, n;
-    static char msg[] = "total size of new array must be unchanged";
 
     dimensions = newshape->ptr;
     n = newshape->len;
@@ -508,26 +488,27 @@ _fix_unknown_dimension(PyArray_Dims *newshape, npy_intp s_original)
             }
             else {
                 PyErr_SetString(PyExc_ValueError,
-                                "can only specify one"  \
-                                " unknown dimension");
+                                "can only specify one unknown dimension");
                 return -1;
             }
         }
-        else {
-            s_known *= dimensions[i];
+        else if (npy_mul_with_overflow_intp(&s_known, s_known,
+                                            dimensions[i])) {
+            raise_reshape_size_mismatch(newshape, arr);
+            return -1;
         }
     }
 
     if (i_unknown >= 0) {
-        if ((s_known == 0) || (s_original % s_known != 0)) {
-            PyErr_SetString(PyExc_ValueError, msg);
+        if (s_known == 0 || s_original % s_known != 0) {
+            raise_reshape_size_mismatch(newshape, arr);
             return -1;
         }
-        dimensions[i_unknown] = s_original/s_known;
+        dimensions[i_unknown] = s_original / s_known;
     }
     else {
         if (s_original != s_known) {
-            PyErr_SetString(PyExc_ValueError, msg);
+            raise_reshape_size_mismatch(newshape, arr);
             return -1;
         }
     }
diff --git a/numpy/core/tests/test_regression.py b/numpy/core/tests/test_regression.py
index a61e64d8d..ace2c1814 100644
--- a/numpy/core/tests/test_regression.py
+++ b/numpy/core/tests/test_regression.py
@@ -2182,6 +2182,21 @@ class TestRegression(TestCase):
         a = np.ones(3, dtype=[('object', 'O'), ('int', '<i2')])
         a.sort()
 
+    def test_reshape_size_overflow(self):
+        # gh-7455
+        a = np.ones(20)[::2]
+        if np.dtype(np.intp).itemsize == 8:
+            # 64 bit. The following are the prime factors of 2**63 + 5,
+            # plus a leading 2, so when multiplied together as int64,
+            # the result overflows to a total size of 10.
+            new_shape = (2, 13, 419, 691, 823, 2977518503)
+        else:
+            # 32 bit. The following are the prime factors of 2**31 + 5,
+            # plus a leading 2, so when multiplied together as int32,
+            # the result overflows to a total size of 10.
+            new_shape = (2, 7, 7, 43826197)
+        assert_raises(ValueError, a.reshape, new_shape)
+
 
 if __name__ == "__main__":
     run_module_suite()