Issue #28322: Fixed possible crashes when unpickle itertools objects from

incorrect pickle data.  Based on patch by John Leitch.

1 files changed, 32 insertions, 0 deletions

diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index 141791c..e054303 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -184,6 +184,19 @@ class TestBasicOps(unittest.TestCase):
         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             self.pickletest(proto, chain('abc', 'def'), compare=list('abcdef'))
 
+    def test_chain_setstate(self):
+        self.assertRaises(TypeError, chain().__setstate__, ())
+        self.assertRaises(TypeError, chain().__setstate__, [])
+        self.assertRaises(TypeError, chain().__setstate__, 0)
+        self.assertRaises(TypeError, chain().__setstate__, ([],))
+        self.assertRaises(TypeError, chain().__setstate__, (iter([]), []))
+        it = chain()
+        it.__setstate__((iter(['abc', 'def']),))
+        self.assertEqual(list(it), ['a', 'b', 'c', 'd', 'e', 'f'])
+        it = chain()
+        it.__setstate__((iter(['abc', 'def']), iter(['ghi'])))
+        self.assertEqual(list(it), ['ghi', 'a', 'b', 'c', 'd', 'e', 'f'])
+
     def test_combinations(self):
         self.assertRaises(TypeError, combinations, 'abc')       # missing r argument
         self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
@@ -631,6 +644,25 @@ class TestBasicOps(unittest.TestCase):
         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             self.pickletest(proto, cycle('abc'))
 
+    def test_cycle_setstate(self):
+        self.assertRaises(TypeError, cycle('').__setstate__, ())
+        self.assertRaises(TypeError, cycle('').__setstate__, [])
+        self.assertRaises(TypeError, cycle('').__setstate__, 0)
+        self.assertRaises(TypeError, cycle('').__setstate__, ([],))
+        self.assertRaises(TypeError, cycle('').__setstate__, ((), 0))
+        it = cycle('abc')
+        it.__setstate__((['de', 'fg'], 0))
+        self.assertEqual(list(islice(it, 15)),
+                         ['a', 'b', 'c', 'de', 'fg',
+                          'a', 'b', 'c', 'de', 'fg',
+                          'a', 'b', 'c', 'de', 'fg'])
+        it = cycle('abc')
+        it.__setstate__((['de', 'fg'], 1))
+        self.assertEqual(list(islice(it, 15)),
+                         ['a', 'b', 'c', 'de', 'fg',
+                          'de', 'fg', 'de', 'fg', 'de',
+                          'fg', 'de', 'fg', 'de', 'fg'])
+
     def test_groupby(self):
         # Check whether it accepts arguments correctly
         self.assertEqual([], list(groupby([])))