[pypy-commit] pypy py3.5: Fix segfault in itertools.product.__setstate__

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: py3.5
Changeset: r87667:c03d2d2eff2f
Date: 2016-10-09 17:30 +0200
http://bitbucket.org/pypy/pypy/changeset/c03d2d2eff2f/

Log:	Fix segfault in itertools.product.__setstate__ (CPython issue 25021)

diff --git a/pypy/module/itertools/interp_itertools.py b/pypy/module/itertools/interp_itertools.py
--- a/pypy/module/itertools/interp_itertools.py
+++ b/pypy/module/itertools/interp_itertools.py
@@ -1216,6 +1216,7 @@
         #
         for gear in self.gears:
             if len(gear) == 0:
+                self.indices = None
                 self.lst = None
                 self.stopped = True
                 break
@@ -1303,12 +1304,16 @@
         for i, gear in enumerate(self.gears):
             w_index = indices_w[i]
             index = space.int_w(w_index)
+            gear_size = len(gear)
+            if self.indices is None or gear_size == 0:
+                self.stopped = True
+                return
             if index < 0:
                 index = 0
-            if index > gear_count - 1:
-                index = gear_count - 1
+            if index > gear_size - 1:
+                index = gear_size - 1
             self.indices[i] = index
-            lst.append(gear[self.indices[i]])
+            lst.append(gear[index])
         self.lst = lst
 
 def W_Product__new__(space, w_subtype, __args__):
diff --git a/pypy/module/itertools/test/test_itertools.py b/pypy/module/itertools/test/test_itertools.py
--- a/pypy/module/itertools/test/test_itertools.py
+++ b/pypy/module/itertools/test/test_itertools.py
@@ -820,6 +820,20 @@
         assert list(product('ab', '', 'ef')) == []
         assert list(product('', 'cd', 'ef')) == []
 
+    def test_product_setstate(self):
+        # test that indices are properly clamped to the length of the tuples
+        from itertools import product
+        p = product((1, 2),(3,))
+        # will access tuple element 1 if not clamped
+        p.__setstate__((0, 0x1000))
+        assert next(p) == (2, 3)
+        # test that empty tuple in the list will result in an
+        # immediate StopIteration
+        p = product((1, 2), (), (3,))
+        # will access tuple element 1 if not clamped
+        p.__setstate__((0, 0, 0x1000))
+        raises(StopIteration, next, p)
+
     def test_permutations(self):
         from itertools import permutations
         assert list(permutations('AB')) == [('A', 'B'), ('B', 'A')]