gh-112127: Fix use-after-free in atexit.unregister()

Lib/test/_test_atexit.py

@@ -135,6 +135,40 @@ def func():
         finally:
             atexit.unregister(func)
 
+    def test_eq_unregister_clear(self):
+        # Issue #112127: callback's __eq__ may call unregister or _clear
+        global cnt
+        cnt = 0
+        class Func:
+            def __init__(self, action, eq_ret_val):
+                self.action = action
+                self.eq_ret_val = eq_ret_val
+
+            def __call__(self):
+                return
+
+            def __eq__(self, o):
+                global cnt
+                cnt += 1
+                if cnt == 1:
+                    self.action(o)
+                return self.eq_ret_val(o)
+
+        for action, eq_ret_val in (
+            (lambda o: atexit.unregister(self), lambda o: NotImplemented),
+            (lambda o: atexit.unregister(self), lambda o: True),
+            (lambda o: atexit.unregister(o),    lambda o: NotImplemented),
+            (lambda o: atexit.unregister(o),    lambda o: True),
+            (lambda o: atexit._clear(),         lambda o: NotImplemented),
+            (lambda o: atexit._clear(),         lambda o: True),
+        ):
+            cnt = 0
+            f1 = Func(action, eq_ret_val)
+            f2 = Func(action, eq_ret_val)
+            atexit.register(f1)
+            atexit.register(f2)
+            atexit._run_exitfuncs()
+
 
 if __name__ == "__main__":
     unittest.main()

Misc/ACKS

@@ -881,6 +881,7 @@ Jim Jewett
 Pedro Diaz Jimenez
 Orjan Johansen
 Fredrik Johansson
+Benjamin Johnson
 Gregory K. Johnson
 Kent Johnson
 Michael Johnson

Modules/atexitmodule.c

@@ -278,11 +278,20 @@ atexit_unregister(PyObject *module, PyObject *func)
             continue;
         }
 
+        /*
+         * Increment refcounts of func and cb->func because equality check may
+         * unregister one or both
+         */
+        PyObject *cb_func = Py_NewRef(cb->func);
+        PyObject *cur_func = Py_NewRef(func);
         int eq = PyObject_RichCompareBool(cb->func, func, Py_EQ);
+        Py_DECREF(cb_func);
+        Py_DECREF(cur_func);
         if (eq < 0) {
             return NULL;
         }
-        if (eq) {
+        // Equality comparison might have already deleted this callback
+        if (eq && state->callbacks[i] != NULL) {
             atexit_delete_cb(state, i);
         }
     }