gh-148263: Fix UNWRAP ref leak with dead proxy

Lib/test/test_weakref.py

@@ -538,6 +538,58 @@ class MyObj:
         with self.assertRaises(TypeError):
             hash(weakref.proxy(obj))
 
+    def test_proxy_unref_add_refcount(self):
+        class C:
+            def __add__(self, o):
+                return NotImplemented
+
+        # create dead proxy
+        o = C()
+        dead = weakref.proxy(o)
+        del o
+        gc.collect()
+
+        # create live proxy
+        obj = C()
+        ref = weakref.ref(obj)
+        proxy = weakref.proxy(obj)
+
+        try:
+            operator.add(proxy, dead)
+        except ReferenceError:
+            pass
+
+        del proxy, obj, dead
+        gc.collect()
+
+        self.assertIsNone(ref(), "Leaked object in add operation")
+
+    def test_proxy_unref_pow_refcount(self):
+        class C:
+            def __pow__(self, o, m=None):
+                return NotImplemented
+
+        # create dead proxy
+        o = C()
+        dead = weakref.proxy(o)
+        del o
+        gc.collect()
+
+        # create live proxy
+        obj = C()
+        ref = weakref.ref(obj)
+        proxy = weakref.proxy(obj)
+
+        try:
+            pow(proxy, dead, None)
+        except ReferenceError:
+            pass
+
+        del proxy, obj, dead
+        gc.collect()
+
+        self.assertIsNone(ref(), "Leaked object in pow operation")
+
     def test_getweakrefcount(self):
         o = C()
         ref1 = weakref.ref(o)

Objects/weakrefobject.c

@@ -530,26 +530,39 @@ proxy_check_ref(PyObject *obj)
     return true;
 }
 
-
-/* If a parameter is a proxy, check that it is still "live" and wrap it,
- * replacing the original value with the raw object.  Raises ReferenceError
- * if the param is a dead proxy.
+/*
+ * Unwrap a proxy into a strong reference.
+ * - If `o` is a live proxy: replaces `o` with the underlying object
+ *   (already Py_INCREF'd by _PyWeakref_GET_REF), sets *did_incref = 1.
+ * - If `o` is a dead proxy: sets ReferenceError, sets `o` = NULL,
+ *   sets *did_incref = 0.
+ * - If `o` is not a proxy: Py_INCREF's it, sets *did_incref = 1.
+ * Returns 1 on success, 0 on dead proxy (caller must goto error).
  */
-#define UNWRAP(o) \
-        if (PyWeakref_CheckProxy(o)) { \
-            o = _PyWeakref_GET_REF(o); \
-            if (!proxy_check_ref(o)) { \
-                return NULL; \
-            } \
-        } \
-        else { \
-            Py_INCREF(o); \
+static inline int
+_proxy_unwrap(PyObject **op, int *did_incref)
+{
+    if (PyWeakref_CheckProxy(*op)) {
+        *op = _PyWeakref_GET_REF(*op);
+        if (!proxy_check_ref(*op)) {
+            *did_incref = 0;
+            return 0;
         }
+        /* _PyWeakref_GET_REF already returned a strong ref */
+    }
+    else {
+        Py_INCREF(*op);
+    }
+    *did_incref = 1;
+    return 1;
+}
 
 #define WRAP_UNARY(method, generic) \
     static PyObject * \
     method(PyObject *proxy) { \
-        UNWRAP(proxy); \
+        int proxy_incref = 0; \
+        if (!_proxy_unwrap(&proxy, &proxy_incref)) \
+            return NULL; \
         PyObject* res = generic(proxy); \
         Py_DECREF(proxy); \
         return res; \
@@ -558,12 +571,19 @@ proxy_check_ref(PyObject *obj)
 #define WRAP_BINARY(method, generic) \
     static PyObject * \
     method(PyObject *x, PyObject *y) { \
-        UNWRAP(x); \
-        UNWRAP(y); \
-        PyObject* res = generic(x, y); \
-        Py_DECREF(x); \
-        Py_DECREF(y); \
-        return res; \
+        int x_incref = 0, y_incref = 0; \
+        if (!_proxy_unwrap(&x, &x_incref)) goto clean_up; \
+        if (!_proxy_unwrap(&y, &y_incref)) goto clean_up; \
+        { \
+            PyObject* res = generic(x, y); \
+            Py_DECREF(x); \
+            Py_DECREF(y); \
+            return res; \
+        } \
+    clean_up: \
+        if (x_incref) Py_DECREF(x); \
+        if (y_incref) Py_DECREF(y); \
+        return NULL; \
     }
 
 /* Note that the third arg needs to be checked for NULL since the tp_call
@@ -572,27 +592,36 @@ proxy_check_ref(PyObject *obj)
 #define WRAP_TERNARY(method, generic) \
     static PyObject * \
     method(PyObject *proxy, PyObject *v, PyObject *w) { \
-        UNWRAP(proxy); \
-        UNWRAP(v); \
+        int proxy_incref = 0, v_incref = 0, w_incref = 0; \
+        if (!_proxy_unwrap(&proxy, &proxy_incref)) goto clean_up; \
+        if (!_proxy_unwrap(&v, &v_incref))         goto clean_up; \
         if (w != NULL) { \
-            UNWRAP(w); \
+            if (!_proxy_unwrap(&w, &w_incref))     goto clean_up; \
         } \
-        PyObject* res = generic(proxy, v, w); \
-        Py_DECREF(proxy); \
-        Py_DECREF(v); \
-        Py_XDECREF(w); \
-        return res; \
+        { \
+            PyObject* res = generic(proxy, v, w); \
+            Py_DECREF(proxy); \
+            Py_DECREF(v); \
+            Py_XDECREF(w); \
+            return res; \
+        } \
+    clean_up: \
+        if (proxy_incref) Py_DECREF(proxy); \
+        if (v_incref)     Py_DECREF(v); \
+        if (w_incref)     Py_DECREF(w); \
+        return NULL; \
     }
 
 #define WRAP_METHOD(method, SPECIAL) \
     static PyObject * \
     method(PyObject *proxy, PyObject *Py_UNUSED(ignored)) { \
-            UNWRAP(proxy); \
-            PyObject* res = PyObject_CallMethodNoArgs(proxy, &_Py_ID(SPECIAL)); \
-            Py_DECREF(proxy); \
-            return res; \
-        }
-
+        int proxy_incref = 0; \
+        if (!_proxy_unwrap(&proxy, &proxy_incref)) \
+            return NULL; \
+        PyObject* res = PyObject_CallMethodNoArgs(proxy, &_Py_ID(SPECIAL)); \
+        Py_DECREF(proxy); \
+        return res; \
+    }
 
 /* direct slots */
 
@@ -635,12 +664,19 @@ proxy_setattr(PyObject *proxy, PyObject *name, PyObject *value)
 static PyObject *
 proxy_richcompare(PyObject *proxy, PyObject *v, int op)
 {
-    UNWRAP(proxy);
-    UNWRAP(v);
-    PyObject* res = PyObject_RichCompare(proxy, v, op);
-    Py_DECREF(proxy);
-    Py_DECREF(v);
-    return res;
+    int proxy_incref = 0, v_incref = 0;
+    if (!_proxy_unwrap(&proxy, &proxy_incref)) goto clean_up;
+    if (!_proxy_unwrap(&v, &v_incref))         goto clean_up;
+    {
+        PyObject* res = PyObject_RichCompare(proxy, v, op);
+        Py_DECREF(proxy);
+        Py_DECREF(v);
+        return res;
+    }
+clean_up:
+    if (proxy_incref) Py_DECREF(proxy);
+    if (v_incref)     Py_DECREF(v);
+    return NULL;
 }
 
 /* number slots */