Fix crash in itertools.cycle.__setstate__() caused by lack of type checking.

Will backport after the 3.6 release is done.

3 files changed, 37 insertions, 1 deletions

diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index fcd8869..53d6564 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -613,6 +613,39 @@ class TestBasicOps(unittest.TestCase):
         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             self.pickletest(proto, cycle('abc'))
 
+    def test_cycle_setstate(self):
+        # Verify both modes for restoring state
+
+        # Mode 0 is efficient.  It uses an incompletely consumed input
+        # iterator to build a cycle object and then passes in state with
+        # a list of previously consumed values.  There is no data
+        # overlap bewteen the two.
+        c = cycle('defg')
+        c.__setstate__((list('abc'), 0))
+        self.assertEqual(take(20, c), list('defgabcdefgabcdefgab'))
+
+        # Mode 1 is inefficient.  It starts with a cycle object built
+        # from an iterator over the remaining elements in a partial
+        # cycle and then passes in state with all of the previously
+        # seen values (this overlaps values included in the iterator).
+        c = cycle('defg')
+        c.__setstate__((list('abcdefg'), 1))
+        self.assertEqual(take(20, c), list('defgabcdefgabcdefgab'))
+
+        # The first argument to setstate needs to be a tuple
+        with self.assertRaises(SystemError):
+            cycle('defg').__setstate__([list('abcdefg'), 0])
+
+        # The first argument in the setstate tuple must be a list
+        with self.assertRaises(TypeError):
+            c = cycle('defg')
+            c.__setstate__((dict.fromkeys('defg'), 0))
+            take(20, c)
+
+        # The first argument in the setstate tuple must be a list
+        with self.assertRaises(TypeError):
+            cycle('defg').__setstate__((list('abcdefg'), 'x'))
+
     def test_groupby(self):
         # Check whether it accepts arguments correctly
         self.assertEqual([], list(groupby([])))
diff --git a/Misc/NEWS b/Misc/NEWS
index 574e279..a8b93e6 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -25,6 +25,9 @@ Library
 - Issue #21159: Improve message in configparser.InterpolationMissingOptionError.
   Patch from Łukasz Langa.
 
+- Fix crash in itertools.cycle.__setstate__() when the first argument wasn't
+  a list.
+
 - Issue #20059: urllib.parse raises ValueError on all invalid ports.
   Patch by Martin Panter.
 
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 349da0f..2cca6b3 100644
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -973,7 +973,7 @@ cycle_setstate(cycleobject *lz, PyObject *state)
 {
     PyObject *saved=NULL;
     int firstpass;
-    if (!PyArg_ParseTuple(state, "Oi", &saved, &firstpass))
+    if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass))
         return NULL;
     Py_CLEAR(lz->saved);
     lz->saved = saved;