Merge pull request #25426 from ngoldbaum/set-invalid-state

BUG: avoid seg fault from OOB access in RandomState.set_state()
numpy · Dec 19, 2023 · 0032ede · 0032ede
2 parents b9a0c9e + 952fac6
commit 0032ede
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 9 deletions.
diff --git a/numpy/random/mtrand.pyx b/numpy/random/mtrand.pyx
@@ -368,15 +368,16 @@ cdef class RandomState:
         else:
             if not isinstance(state, (tuple, list)):
                 raise TypeError('state must be a dict or a tuple.')
-            if state[0] != 'MT19937':
-                raise ValueError('set_state can only be used with legacy MT19937'
-                                 'state instances.')
-            st = {'bit_generator': state[0],
-                  'state': {'key': state[1], 'pos': state[2]}}
-            if len(state) > 3:
-                st['has_gauss'] = state[3]
-                st['gauss'] = state[4]
-                value = st
+            with cython.boundscheck(True):
+                if state[0] != 'MT19937':
+                    raise ValueError('set_state can only be used with legacy '
+                                     'MT19937 state instances.')
+                st = {'bit_generator': state[0],
+                      'state': {'key': state[1], 'pos': state[2]}}
+                if len(state) > 3:
+                    st['has_gauss'] = state[3]
+                    st['gauss'] = state[4]
+                    value = st
 
         self._aug_state.gauss = st.get('gauss', 0.0)
         self._aug_state.has_gauss = st.get('has_gauss', 0)

diff --git a/numpy/random/tests/test_random.py b/numpy/random/tests/test_random.py
@@ -147,6 +147,11 @@ def test_negative_binomial(self):
         # arguments without truncation.
         self.prng.negative_binomial(0.5, 0.5)
 
+    def test_set_invalid_state(self):
+        # gh-25402
+        with pytest.raises(IndexError):
+            self.prng.set_state(())
+
 
 class TestRandint: