Fix dangling pointer in internals::registered_types_cpp_fast from pybind#5842

@oremanj pointed out in a comment on pybind#5842 that I missed part
of the nanobind PR I was porting in such a way that we could have
dangling pointers in internals::registered_types_cpp_fast. This PR
adds a test that reproed the bug and then fixes the test.

Author: swolchok

include/pybind11/detail/class.h

@@ -228,6 +228,13 @@ extern "C" inline void pybind11_meta_dealloc(PyObject *obj) {
                 internals.registered_types_cpp.erase(tindex);
 #if PYBIND11_INTERNALS_VERSION >= 12
                 internals.registered_types_cpp_fast.erase(tinfo->cpptype);
+                const alias_chain_entry *chain = tinfo->alias_chain.get();
+                while (chain) {
+                    auto num_erased = internals.registered_types_cpp_fast.erase(chain->value);
+                    (void) num_erased;
+                    assert(num_erased > 0);
+                    chain = chain->next.get();
+                }
 #endif
             }
             internals.registered_types_py.erase(tinfo->type);

include/pybind11/detail/internals.h

@@ -335,6 +335,22 @@ enum class holder_enum_t : uint8_t {
     custom_holder,
 };
 
+// When a type appears in multiple DSOs,
+// internals::registered_types_cpp_fast will have multiple distinct
+// keys (the type_info from each DSO) mapped to the same
+// type_info*. We need to keep track of these aliases so that we clean
+// them up when our type is deallocated. A linked list is appropriate
+// because this structure is expected to be 1) usually empty and 2)
+// when it's not empty, usually very small. See also `struct
+// nb_alias_chain` added in
+// https://github.com/wjakob/nanobind/commit/b515b1f7f2f4ecc0357818e6201c94a9f4cbfdc2
+#if PYBIND11_INTERNALS_VERSION >= 12
+struct alias_chain_entry {
+    std::unique_ptr<alias_chain_entry> next;
+    const std::type_info *value;
+};
+#endif
+
 /// Additional type information which does not fit into the PyTypeObject.
 /// Changes to this struct also require bumping `PYBIND11_INTERNALS_VERSION`.
 struct type_info {
@@ -357,6 +373,11 @@ struct type_info {
     void *get_buffer_data = nullptr;
     void *(*module_local_load)(PyObject *, const type_info *) = nullptr;
     holder_enum_t holder_enum_v = holder_enum_t::undefined;
+
+#if PYBIND11_INTERNALS_VERSION >= 12
+    std::unique_ptr<alias_chain_entry> alias_chain;
+#endif
+
     /* A simple type never occurs as a (direct or indirect) parent
      * of a class that makes use of multiple inheritance.
      * A type can be simple even if it has non-simple ancestors as long as it has no descendants.

include/pybind11/detail/type_caster_base.h

@@ -246,6 +246,14 @@ inline detail::type_info *get_global_type_info_lock_held(const std::type_info &t
     auto it = types.find(std::type_index(tp));
     if (it != types.end()) {
 #if PYBIND11_INTERNALS_VERSION >= 12
+        // We found the type in the slow map but not the fast one, so
+        // some other DSO added it (otherwise it would be in the fast
+        // map under &tp) and therefore we must be an alias. Record
+        // that in the chain.
+        std::unique_ptr<alias_chain_entry> chain(new alias_chain_entry);
+        chain->value = &tp;
+        chain->next = std::move(it->second->alias_chain);
+        it->second->alias_chain = std::move(chain);
         fast_types.emplace(&tp, it->second);
 #endif
         type_info = it->second;

tests/CMakeLists.txt

@@ -117,6 +117,7 @@ set(PYBIND11_TEST_FILES
     test_callbacks
     test_chrono
     test_class
+    test_class_cross_module_use_after_one_module_dealloc
     test_class_release_gil_before_calling_cpp_dtor
     test_class_sh_basic
     test_class_sh_disowning
@@ -239,8 +240,9 @@ list(SORT PYBIND11_PYTEST_FILES)
 # Contains the set of test files that require pybind11_cross_module_tests to be
 # built; if none of these are built (i.e. because TEST_OVERRIDE is used and
 # doesn't include them) the second module doesn't get built.
-tests_extra_targets("test_exceptions.py;test_local_bindings.py;test_stl.py;test_stl_binders.py"
-                    "pybind11_cross_module_tests")
+tests_extra_targets(
+  "test_class_cross_module_use_after_one_module_dealloc.py;test_exceptions.py;test_local_bindings.py;test_stl.py;test_stl_binders.py"
+  "pybind11_cross_module_tests")
 
 # And add additional targets for other tests.
 tests_extra_targets("test_exceptions.py" "cross_module_interleaved_error_already_set")

tests/pybind11_cross_module_tests.cpp

@@ -16,6 +16,13 @@
 #include <numeric>
 #include <utility>
 
+class CrossDSOClass {
+public:
+    virtual ~CrossDSOClass();
+};
+
+CrossDSOClass::~CrossDSOClass() = default;
+
 PYBIND11_MODULE(pybind11_cross_module_tests, m, py::mod_gil_not_used()) {
     m.doc() = "pybind11 cross-module test module";
 
@@ -146,4 +153,7 @@ PYBIND11_MODULE(pybind11_cross_module_tests, m, py::mod_gil_not_used()) {
     // which appears when this header is missing.
     m.def("missing_header_arg", [](const std::vector<float> &) {});
     m.def("missing_header_return", []() { return std::vector<float>(); });
+
+    // test_class_cross_module_use_after_one_module_dealloc
+    m.def("consume_cross_dso_class", [](CrossDSOClass) {});
 }

tests/test_class_cross_module_use_after_one_module_dealloc.cpp

@@ -0,0 +1,21 @@
+#include "pybind11_tests.h"
+
+#include <iostream>
+
+class CrossDSOClass {
+public:
+    virtual ~CrossDSOClass();
+};
+
+CrossDSOClass::~CrossDSOClass() = default;
+
+struct UnrelatedClass {};
+
+TEST_SUBMODULE(class_cross_module_use_after_one_module_dealloc, m) {
+    m.def("register_and_instantiate_cross_dso_class", [](py::module_ m) {
+        py::class_<CrossDSOClass>(m, "CrossDSOClass").def(py::init<>());
+        return CrossDSOClass();
+    });
+    m.def("register_unrelated_class",
+          [](py::module_ m) { py::class_<UnrelatedClass>(m, "UnrelatedClass"); });
+}

tests/test_class_cross_module_use_after_one_module_dealloc.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import gc
+import types
+import weakref
+
+import pytest
+
+import env  # noqa: F401
+from pybind11_tests import class_cross_module_use_after_one_module_dealloc as m
+
+
+def delattr_and_ensure_destroyed(*specs):
+    wrs = []
+    for mod, name in specs:
+        wrs.append(weakref.ref(getattr(mod, name)))
+        delattr(mod, name)
+
+    for _ in range(5):
+        gc.collect()
+        if all(wr() is None for wr in wrs):
+            break
+    else:
+        pytest.fail(
+            f"Could not delete bindings such as {next(wr for wr in wrs if wr() is not None)!r}"
+        )
+
+
+@pytest.mark.skipif("env.PYPY or env.GRAALPY")
+def test_cross_module_use_after_one_module_dealloc():
+    # This is a regression test for a bug that occurred during development of
+    # internals::registered_types_cpp_fast (see #5842). registered_types_cpp_fast maps
+    # &typeid(T) to a raw non-owning pointer to a Python metaclass. If two DSOs both
+    # look up the same global type, they will create two separate entries in
+    # registered_types_cpp_fast, which will look like:
+    # +=======================================+
+    # |&typeid(T) from DSO 1|metaclass pointer|
+    # |&typeid(T) from DSO 2|metaclass pointer|
+    # +=======================================+
+    #
+    # Then, if the metaclass is destroyed and we don't take extra steps to clean up the
+    # table thoroughly, the first row of the table will be cleaned up but the second one
+    # will contain a dangling pointer to the old metaclass instance. Further lookups
+    # from DSO 2 will then return that dangling pointer, which will cause use-after-frees.
+
+    import pybind11_cross_module_tests as cm
+
+    module_scope = types.ModuleType("module_scope")
+    instance = m.register_and_instantiate_cross_dso_class(module_scope)
+    cm.consume_cross_dso_class(instance)
+
+    del instance
+    delattr_and_ensure_destroyed((module_scope, "CrossDSOClass"))
+
+    # Make sure that CrossDSOClass gets allocated at a different address.
+    m.register_unrelated_class(module_scope)
+
+    instance = m.register_and_instantiate_cross_dso_class(module_scope)
+    cm.consume_cross_dso_class(instance)