bpo-44720: Don't crash when calling weakref.proxy(not_an_iterator).__next__ (GH-27316)

3 files changed, 37 insertions, 0 deletions

diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py
index dd5a781..1a5314c 100644
--- a/Lib/test/test_weakref.py
+++ b/Lib/test/test_weakref.py
@@ -411,6 +411,36 @@ class ReferencesTestCase(TestBase):
             # can be killed in the middle of the call
             "blech" in p
 
+    def test_proxy_next(self):
+        arr = [4, 5, 6]
+        def iterator_func():
+            yield from arr
+        it = iterator_func()
+
+        class IteratesWeakly:
+            def __iter__(self):
+                return weakref.proxy(it)
+
+        weak_it = IteratesWeakly()
+
+        # Calls proxy.__next__
+        self.assertEqual(list(weak_it), [4, 5, 6])
+
+    def test_proxy_bad_next(self):
+        # bpo-44720: PyIter_Next() shouldn't be called if the reference
+        # isn't an iterator.
+
+        not_an_iterator = lambda: 0
+
+        class A:
+            def __iter__(self):
+                return weakref.proxy(not_an_iterator)
+        a = A()
+
+        msg = "Weakref proxy referenced a non-iterator"
+        with self.assertRaisesRegex(TypeError, msg):
+            list(a)
+
     def test_proxy_reversed(self):
         class MyObj:
             def __len__(self):
diff --git a/Misc/NEWS.d/next/Library/2021-07-24-02-17-59.bpo-44720.shU5Qm.rst b/Misc/NEWS.d/next/Library/2021-07-24-02-17-59.bpo-44720.shU5Qm.rst
new file mode 100644
index 0000000..83694f3
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2021-07-24-02-17-59.bpo-44720.shU5Qm.rst
@@ -0,0 +1 @@
+``weakref.proxy`` objects referencing non-iterators now raise ``TypeError`` rather than dereferencing the null ``tp_iternext`` slot and crashing.
\ No newline at end of file
diff --git a/Objects/weakrefobject.c b/Objects/weakrefobject.c
index c36d239..bb56c7d 100644
--- a/Objects/weakrefobject.c
+++ b/Objects/weakrefobject.c
@@ -657,6 +657,12 @@ proxy_iternext(PyWeakReference *proxy)
         return NULL;
 
     PyObject *obj = PyWeakref_GET_OBJECT(proxy);
+    if (!PyIter_Check(obj)) {
+        PyErr_Format(PyExc_TypeError,
+            "Weakref proxy referenced a non-iterator '%.200s' object",
+            Py_TYPE(obj)->tp_name);
+        return NULL;
+    }
     Py_INCREF(obj);
     PyObject* res = PyIter_Next(obj);
     Py_DECREF(obj);