Merge pull request #3845 from mbyt/out_of_memory_access_as_strided

tacaswell · tacaswell · commit df3530de9b67 · 2014-11-26T15:34:02.000-05:00
BUG: ensure integer shape, strides for as_strided
diff --git a/lib/matplotlib/mlab.py b/lib/matplotlib/mlab.py
@@ -590,6 +590,11 @@ def stride_windows(x, n, noverlap=None, axis=0):
     if n > x.size:
         raise ValueError('n cannot be greater than the length of x')
 
+    # np.lib.stride_tricks.as_strided easily leads to memory corruption for
+    # non integer shape and strides, i.e. noverlap or n. See #3845.
+    noverlap = int(noverlap)
+    n = int(n)
+
     step = n - noverlap
     if axis == 0:
         shape = (n, (x.shape[-1]-noverlap)//step)
@@ -642,6 +647,10 @@ def stride_repeat(x, n, axis=0):
     if n < 1:
         raise ValueError('n cannot be less than 1')
 
+    # np.lib.stride_tricks.as_strided easily leads to memory corruption for
+    # non integer shape and strides, i.e. n. See #3845.
+    n = int(n)
+
     if axis == 0:
         shape = (n, x.size)
         strides = (0, x.strides[0])
diff --git a/lib/matplotlib/tests/test_mlab.py b/lib/matplotlib/tests/test_mlab.py
@@ -283,6 +283,24 @@ def test_stride_windows_n32_noverlap0_axis1_unflatten(self):
         assert_equal(y.shape, x1.shape)
         assert_array_equal(y, x1)
 
+    def test_stride_ensure_integer_type(self):
+        N = 100
+        x = np.empty(N + 20, dtype='>f4')
+        x.fill(np.NaN)
+        y = x[10:-10]
+        y.fill(0.3)
+        # previous to #3845 lead to corrupt access
+        y_strided = mlab.stride_windows(y, n=33, noverlap=0.6)
+        assert_array_equal(y_strided, 0.3)
+        # previous to #3845 lead to corrupt access
+        y_strided = mlab.stride_windows(y, n=33.3, noverlap=0)
+        assert_array_equal(y_strided, 0.3)
+        # even previous to #3845 could not find any problematic
+        # configuration however, let's be sure it's not accidentally
+        # introduced
+        y_strided = mlab.stride_repeat(y, n=33.815)
+        assert_array_equal(y_strided, 0.3)
+
 
 class csv_testcase(CleanupTestCase):
     def setUp(self):
