summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--Lib/test/test_itertools.py24
-rw-r--r--Misc/NEWS3
-rw-r--r--Modules/itertoolsmodule.c40
3 files changed, 57 insertions, 10 deletions
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index d21d8ed..ea1f57c 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -183,6 +183,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
@@ -667,19 +680,22 @@ class TestBasicOps(unittest.TestCase):
self.assertEqual(take(20, c), list('defgabcdefgabcdefgab'))
# The first argument to setstate needs to be a tuple
- with self.assertRaises(SystemError):
+ with self.assertRaises(TypeError):
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)
+ c.__setstate__((tuple('defg'), 0))
+ take(20, c)
- # The first argument in the setstate tuple must be a list
+ # The second argument in the setstate tuple must be an int
with self.assertRaises(TypeError):
cycle('defg').__setstate__((list('abcdefg'), 'x'))
+ self.assertRaises(TypeError, cycle('').__setstate__, ())
+ self.assertRaises(TypeError, cycle('').__setstate__, ([],))
+
def test_groupby(self):
# Check whether it accepts arguments correctly
self.assertEqual([], list(groupby([])))
diff --git a/Misc/NEWS b/Misc/NEWS
index 8b102a5..831cb22 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -46,6 +46,9 @@ Core and Builtins
Library
-------
+- Issue #28322: Fixed possible crashes when unpickle itertools objects from
+ incorrect pickle data. Based on patch by John Leitch.
+
- Issue #28228: imghdr now supports pathlib.
- Issue #28226: compileall now supports pathlib.
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 62b6a0c..6bf04cb 100644
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -147,8 +147,13 @@ static PyObject *
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);
@@ -727,8 +732,13 @@ tee_setstate(teeobject *to, PyObject *state)
{
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;
@@ -971,9 +981,13 @@ cycle_setstate(cycleobject *lz, PyObject *state)
{
PyObject *saved=NULL;
int firstpass;
-
- if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &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_INCREF(saved);
Py_XSETREF(lz->saved, saved);
lz->firstpass = firstpass != 0;
@@ -1903,8 +1917,17 @@ 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);
@@ -3262,10 +3285,15 @@ permutations_setstate(permutationsobject *po, PyObject *state)
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 || PyTuple_GET_SIZE(cycles) != po->r) {