From da6c78d03fed9ee64bb587644e65b27b77ea10bb Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Mon, 9 Sep 2019 11:47:14 +0300
Subject: [PATCH] [2.7] bpo-34410: Fix a crash in the tee iterator when
 re-enter it. (GH-15625)

RuntimeError is now raised in this case..
(cherry picked from commit 526a01467b3277f9fcf7f91e66c23321caa1245d)

Co-authored-by: Serhiy Storchaka <storchaka@gmail.com>
---
 Doc/library/itertools.rst                     |  4 ++
 Lib/test/test_itertools.py                    | 41 +++++++++++++++++++
 .../2019-08-31-01-52-59.bpo-34410.7KbWZQ.rst  |  2 +
 Modules/itertoolsmodule.c                     |  9 ++++
 4 files changed, 56 insertions(+)
 create mode 100644 Misc/NEWS.d/next/Library/2019-08-31-01-52-59.bpo-34410.7KbWZQ.rst

diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst
index 17303dd01cef06..18477f1299031c 100644
--- a/Doc/library/itertools.rst
+++ b/Doc/library/itertools.rst
@@ -659,6 +659,10 @@ loops that truncate the stream.
    used anywhere else; otherwise, the *iterable* could get advanced without
    the tee objects being informed.
 
+   ``tee`` iterators are not threadsafe. A :exc:`RuntimeError` may be
+   raised when using simultaneously iterators returned by the same :func:`tee`
+   call, even if the original *iterable* is threadsafe.
+
    This itertool may require significant auxiliary storage (depending on how
    much temporary data needs to be stored). In general, if one iterator uses
    most or all of the data before another iterator starts, it is faster to use
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index 04279792065edf..2cdcbb21efcdf6 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -10,6 +10,10 @@
 import copy
 import pickle
 from functools import reduce
+try:
+    import threading
+except ImportError:
+    threading = None
 maxsize = test_support.MAX_Py_ssize_t
 minsize = -maxsize-1
 
@@ -984,6 +988,43 @@ def test_tee_del_backward(self):
             del forward, backward
             raise
 
+    def test_tee_reenter(self):
+        class I:
+            first = True
+            def __iter__(self):
+                return self
+            def next(self):
+                first = self.first
+                self.first = False
+                if first:
+                    return next(b)
+
+        a, b = tee(I())
+        with self.assertRaisesRegexp(RuntimeError, "tee"):
+            next(a)
+
+    @unittest.skipUnless(threading, 'Threading required for this test.')
+    def test_tee_concurrent(self):
+        start = threading.Event()
+        finish = threading.Event()
+        class I:
+            def __iter__(self):
+                return self
+            def next(self):
+                start.set()
+                finish.wait()
+
+        a, b = tee(I())
+        thread = threading.Thread(target=next, args=[a])
+        thread.start()
+        try:
+            start.wait()
+            with self.assertRaisesRegexp(RuntimeError, "tee"):
+                next(b)
+        finally:
+            finish.set()
+            thread.join()
+
     def test_StopIteration(self):
         self.assertRaises(StopIteration, izip().next)
 
diff --git a/Misc/NEWS.d/next/Library/2019-08-31-01-52-59.bpo-34410.7KbWZQ.rst b/Misc/NEWS.d/next/Library/2019-08-31-01-52-59.bpo-34410.7KbWZQ.rst
new file mode 100644
index 00000000000000..64e778ee0913c2
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2019-08-31-01-52-59.bpo-34410.7KbWZQ.rst
@@ -0,0 +1,2 @@
+Fixed a crash in the :func:`tee` iterator when re-enter it. RuntimeError is
+now raised in this case.
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 04076fd087099c..edd21be337cf51 100644
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -314,6 +314,7 @@ typedef struct {
     PyObject_HEAD
     PyObject *it;
     int numread;
+    int running;
     PyObject *nextlink;
     PyObject *(values[LINKCELLS]);
 } teedataobject;
@@ -336,6 +337,7 @@ teedataobject_new(PyObject *it)
     if (tdo == NULL)
         return NULL;
 
+    tdo->running = 0;
     tdo->numread = 0;
     tdo->nextlink = NULL;
     Py_INCREF(it);
@@ -364,7 +366,14 @@ teedataobject_getitem(teedataobject *tdo, int i)
     else {
         /* this is the lead iterator, so fetch more data */
         assert(i == tdo->numread);
+        if (tdo->running) {
+            PyErr_SetString(PyExc_RuntimeError,
+                            "cannot re-enter the tee iterator");
+            return NULL;
+        }
+        tdo->running = 1;
         value = PyIter_Next(tdo->it);
+        tdo->running = 0;
         if (value == NULL)
             return NULL;
         tdo->numread++;