Skip to content

Commit c5cc501

Browse files
kristjanvalurkristjanvalur
committed
Make the various iterators' "setstate" sliently and consistently clip the
index. This avoids the possibility of setting an iterator to an invalid state.
2 parents 25ea45d + 25dded0 commit c5cc501

File tree

8 files changed

+66
-15
lines changed

8 files changed

+66
-15
lines changed

‎Lib/test/test_range.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -381,6 +381,18 @@ def test_iterator_pickling(self):
381381
self.assertEqual(list(it), data[1:])
382382

383383
def test_exhausted_iterator_pickling(self):
384+
r = range(2**65, 2**65+2)
385+
i = iter(r)
386+
while True:
387+
r = next(i)
388+
if r == 2**65+1:
389+
break
390+
d = pickle.dumps(i)
391+
i2 = pickle.loads(d)
392+
self.assertEqual(list(i), [])
393+
self.assertEqual(list(i2), [])
394+
395+
def test_large_exhausted_iterator_pickling(self):
384396
r = range(20)
385397
i = iter(r)
386398
while True:

‎Modules/arraymodule.c‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2838,6 +2838,8 @@ arrayiter_setstate(arrayiterobject *it, PyObject *state)
28382838
return NULL;
28392839
if (index < 0)
28402840
index = 0;
2841+
else if (index > Py_SIZE(it->ao))
2842+
index = Py_SIZE(it->ao); /* iterator exhausted */
28412843
it->index = index;
28422844
Py_RETURN_NONE;
28432845
}

‎Objects/bytearrayobject.c‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3025,9 +3025,13 @@ bytearrayiter_setstate(bytesiterobject *it, PyObject *state)
30253025
Py_ssize_t index = PyLong_AsSsize_t(state);
30263026
if (index == -1 && PyErr_Occurred())
30273027
return NULL;
3028-
if (index < 0)
3029-
index = 0;
3030-
it->it_index = index;
3028+
if (it->it_seq != NULL) {
3029+
if (index < 0)
3030+
index = 0;
3031+
else if (index > PyByteArray_GET_SIZE(it->it_seq))
3032+
index = PyByteArray_GET_SIZE(it->it_seq); /* iterator exhausted */
3033+
it->it_index = index;
3034+
}
30313035
Py_RETURN_NONE;
30323036
}
30333037

‎Objects/bytesobject.c‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2937,9 +2937,13 @@ striter_setstate(striterobject *it, PyObject *state)
29372937
Py_ssize_t index = PyLong_AsSsize_t(state);
29382938
if (index == -1 && PyErr_Occurred())
29392939
return NULL;
2940-
if (index < 0)
2941-
index = 0;
2942-
it->it_index = index;
2940+
if (it->it_seq != NULL) {
2941+
if (index < 0)
2942+
index = 0;
2943+
else if (index > PyBytes_GET_SIZE(it->it_seq))
2944+
index = PyBytes_GET_SIZE(it->it_seq); /* iterator exhausted */
2945+
it->it_index = index;
2946+
}
29432947
Py_RETURN_NONE;
29442948
}
29452949

‎Objects/listobject.c‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2811,6 +2811,8 @@ listiter_setstate(listiterobject *it, PyObject *state)
28112811
if (it->it_seq != NULL) {
28122812
if (index < 0)
28132813
index = 0;
2814+
else if (index > PyList_GET_SIZE(it->it_seq))
2815+
index = PyList_GET_SIZE(it->it_seq); /* iterator exhausted */
28142816
it->it_index = index;
28152817
}
28162818
Py_RETURN_NONE;

‎Objects/rangeobject.c‎

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -807,10 +807,11 @@ rangeiter_setstate(rangeiterobject *r, PyObject *state)
807807
long index = PyLong_AsLong(state);
808808
if (index == -1 && PyErr_Occurred())
809809
return NULL;
810-
if (index < 0 || index > r->len) {
811-
PyErr_SetString(PyExc_ValueError, "index out of range");
812-
return NULL;
813-
}
810+
/* silently clip the index value */
811+
if (index < 0)
812+
index = 0;
813+
else if (index > r->len)
814+
index = r->len; /* exhausted iterator */
814815
r->index = index;
815816
Py_RETURN_NONE;
816817
}
@@ -985,6 +986,28 @@ longrangeiter_reduce(longrangeiterobject *r)
985986
static PyObject *
986987
longrangeiter_setstate(longrangeiterobject *r, PyObject *state)
987988
{
989+
int cmp;
990+
991+
/* clip the value */
992+
PyObject *zero = PyLong_FromLong(0);
993+
if (zero == NULL)
994+
return NULL;
995+
cmp = PyObject_RichCompareBool(state, zero, Py_LT);
996+
if (cmp > 0) {
997+
Py_CLEAR(r->index);
998+
r->index = zero;
999+
Py_RETURN_NONE;
1000+
}
1001+
Py_DECREF(zero);
1002+
if (cmp < 0)
1003+
return NULL;
1004+
1005+
cmp = PyObject_RichCompareBool(r->len, state, Py_LT);
1006+
if (cmp < 0)
1007+
return NULL;
1008+
if (cmp > 0)
1009+
state = r->len;
1010+
9881011
Py_CLEAR(r->index);
9891012
r->index = state;
9901013
Py_INCREF(r->index);

‎Objects/tupleobject.c‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1011,8 +1011,8 @@ tupleiter_setstate(tupleiterobject *it, PyObject *state)
10111011
if (it->it_seq != NULL) {
10121012
if (index < 0)
10131013
index = 0;
1014-
else if (it->it_seq != NULL && index > PyTuple_GET_SIZE(it->it_seq))
1015-
index = PyTuple_GET_SIZE(it->it_seq);
1014+
else if (index > PyTuple_GET_SIZE(it->it_seq))
1015+
index = PyTuple_GET_SIZE(it->it_seq); /* exhausted iterator */
10161016
it->it_index = index;
10171017
}
10181018
Py_RETURN_NONE;

‎Objects/unicodeobject.c‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15196,9 +15196,13 @@ unicodeiter_setstate(unicodeiterobject *it, PyObject *state)
1519615196
Py_ssize_t index = PyLong_AsSsize_t(state);
1519715197
if (index == -1 && PyErr_Occurred())
1519815198
return NULL;
15199-
if (index < 0)
15200-
index = 0;
15201-
it->it_index = index;
15199+
if (it->it_seq != NULL) {
15200+
if (index < 0)
15201+
index = 0;
15202+
else if (index > PyUnicode_GET_LENGTH(it->it_seq))
15203+
index = PyUnicode_GET_LENGTH(it->it_seq); /* iterator truncated */
15204+
it->it_index = index;
15205+
}
1520215206
Py_RETURN_NONE;
1520315207
}
1520415208

0 commit comments

Comments
 (0)