TST: Restrict accepted shapes to plausible vectors

Coercing vector to float before checking size could cause memory blow-up

Author: effigies

nibabel/freesurfer/io.py

@@ -179,17 +179,20 @@ def write_morph_data(file_like, values, fnum=0):
         in binary write (`'wb'` mode, implementing the `write` method)
     values : array-like
         Surface morphometry values
+
+        Shape must be (N,), (N, 1), (1, N) or (N, 1, 1)
     fnum : int, optional
         Number of faces in the associated surface
     """
     magic_bytes = np.array([255, 255, 255], dtype=np.uint8)
 
-    array = np.asarray(values).astype('>f4').squeeze()
-    if len(array.shape) > 1:
-        raise ValueError("Multi-dimensional values not supported")
+    vector = np.asarray(values)
+    vnum = np.prod(vector.shape)
+    if vector.shape not in ((vnum,), (vnum, 1), (1, vnum), (vnum, 1, 1)):
+        raise ValueError("Invalid shape: argument values must be a vector")
 
     i4info = np.iinfo('i4')
-    if len(array) > i4info.max:
+    if vnum > i4info.max:
         raise ValueError("Too many values for morphometry file")
     if not i4info.min <= fnum <= i4info.max:
         raise ValueError("Argument fnum must be between {0} and {1}".format(
@@ -199,9 +202,9 @@ def write_morph_data(file_like, values, fnum=0):
         fobj.write(magic_bytes)
 
         # vertex count, face count (unused), vals per vertex (only 1 supported)
-        fobj.write(np.array([len(array), fnum, 1], dtype='>i4'))
+        fobj.write(np.array([vnum, fnum, 1], dtype='>i4'))
 
-        fobj.write(array)
+        fobj.write(vector.astype('>f4'))
 
 
 def read_annot(filepath, orig_ids=False):

nibabel/freesurfer/tests/test_io.py

@@ -104,7 +104,7 @@ def test_write_morph_data():
     """Test write_morph_data edge cases"""
     values = np.arange(20, dtype='>f4')
     okay_shapes = [(20,), (20, 1), (20, 1, 1), (1, 20)]
-    bad_shape = (10, 2)
+    bad_shapes = [(10, 2), (1, 1, 20, 1, 1)]
     big_num = np.iinfo('i4').max + 1
     with InTemporaryDirectory():
         for shape in okay_shapes:
@@ -113,10 +113,11 @@ def test_write_morph_data():
             assert_equal(values, read_morph_data('test.curv'))
         assert_raises(ValueError, write_morph_data, 'test.curv',
                       np.zeros(shape), big_num)
-        assert_raises(ValueError, write_morph_data, 'test.curv',
-                      values.reshape(bad_shape))
         assert_raises(ValueError, write_morph_data, 'test.curv',
                       strided_scalar((big_num,)))
+        for shape in bad_shapes:
+            assert_raises(ValueError, write_morph_data, 'test.curv',
+                          values.reshape(shape))
 
 
 @freesurfer_test