[Python-checkins] cpython (3.5): Issue #28322: Fixed possible crashes when unpickle itertools objects from

[serhiy.storchaka]

https://hg.python.org/cpython/rev/258ebc539b2e
changeset:   104224:258ebc539b2e
branch:      3.5
parent:      104209:36d37ff6c236
user:        Serhiy Storchaka <storchaka at gmail.com>
date:        Sun Oct 02 08:34:53 2016 +0300
summary:
  Issue #28322: Fixed possible crashes when unpickle itertools objects from
incorrect pickle data.  Based on patch by John Leitch.

files:
  Lib/test/test_itertools.py |  32 ++++++++++++++++++++
  Misc/NEWS                  |   3 +
  Modules/itertoolsmodule.c  |  40 ++++++++++++++++++++++---
  3 files changed, 70 insertions(+), 5 deletions(-)


diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -184,6 +184,19 @@
         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 @@
         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([])))
diff --git a/Misc/NEWS b/Misc/NEWS
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -85,6 +85,9 @@
 Library
 -------
 
+- Issue #28322: Fixed possible crashes when unpickle itertools objects from
+  incorrect pickle data.  Based on patch by John Leitch.
+
 - Issue #1703178: Fix the ability to pass the --link-objects option to the
   distutils build_ext command.
 
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -155,8 +155,13 @@
 groupby_setstate(groupbyobject *lz, PyObject *state)
 {
     PyObject *currkey, *currvalue, *tgtkey;
-    if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey))
+    if (!PyTuple_Check(state)) {
+        PyErr_SetString(PyExc_TypeError, "state is not a tuple");
         return NULL;
+    }
+    if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
+        return NULL;
+    }
     Py_INCREF(currkey);
     Py_XSETREF(lz->currkey, currkey);
     Py_INCREF(currvalue);
@@ -736,8 +741,13 @@
 {
     teedataobject *tdo;
     int index;
-    if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index))
+    if (!PyTuple_Check(state)) {
+        PyErr_SetString(PyExc_TypeError, "state is not a tuple");
         return NULL;
+    }
+    if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
+        return NULL;
+    }
     if (index < 0 || index > LINKCELLS) {
         PyErr_SetString(PyExc_ValueError, "Index out of range");
         return NULL;
@@ -966,8 +976,13 @@
 {
     PyObject *saved=NULL;
     int firstpass;
-    if (!PyArg_ParseTuple(state, "Oi", &saved, &firstpass))
+    if (!PyTuple_Check(state)) {
+        PyErr_SetString(PyExc_TypeError, "state is not a tuple");
         return NULL;
+    }
+    if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
+        return NULL;
+    }
     Py_XINCREF(saved);
     Py_XSETREF(lz->saved, saved);
     lz->firstpass = firstpass != 0;
@@ -1891,8 +1906,18 @@
 chain_setstate(chainobject *lz, PyObject *state)
 {
     PyObject *source, *active=NULL;
-    if (! PyArg_ParseTuple(state, "O|O", &source, &active))
+
+    if (!PyTuple_Check(state)) {
+        PyErr_SetString(PyExc_TypeError, "state is not a tuple");
         return NULL;
+    }
+    if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
+        return NULL;
+    }
+    if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
+        PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
+        return NULL;
+    }
 
     Py_INCREF(source);
     Py_XSETREF(lz->source, source);
@@ -3251,10 +3276,15 @@
     PyObject *indices, *cycles, *result;
     Py_ssize_t n, i;
 
+    if (!PyTuple_Check(state)) {
+        PyErr_SetString(PyExc_TypeError, "state is not a tuple");
+        return NULL;
+    }
     if (!PyArg_ParseTuple(state, "O!O!",
                           &PyTuple_Type, &indices,
-                          &PyTuple_Type, &cycles))
+                          &PyTuple_Type, &cycles)) {
         return NULL;
+    }
 
     n = PyTuple_GET_SIZE(po->pool);
     if (PyTuple_GET_SIZE(indices) != n ||

-- 
Repository URL: https://hg.python.org/cpython