Fix custom class name lookup for Unity 6

CMakeLists.txt

@@ -28,6 +28,7 @@ add_compile_definitions(MOD_ID="${MOD_ID}")
 add_compile_definitions(VERSION_LONG=1)
 add_compile_definitions(NEED_UNSAFE_CSHARP)
 add_compile_definitions(UNITY_6)
+add_compile_definitions(UNITY_6_3)
 add_compile_definitions(NO_VERBOSE_LOGS)
 
 # Add any extra definitions from the build script

shared/register.hpp

@@ -7,7 +7,9 @@
 
 #include <list>
 #include <mutex>
+#include <optional>
 #include <shared_mutex>
+#include <string_view>
 
 template <>
 struct std::hash<std::pair<std::string, std::string>> {
@@ -54,6 +56,8 @@ namespace custom_types {
        public:
         CUSTOM_TYPES_EXPORT static std::unordered_map<std::pair<std::string, std::string>, Il2CppClass*> classMapping;
         CUSTOM_TYPES_EXPORT static std::vector<Il2CppClass*> classes;
+        // nullopt delegates to Unity; an engaged result (including nullptr) belongs to a custom image.
+        CUSTOM_TYPES_EXPORT static std::optional<Il2CppClass*> FindClass(Il2CppImage const* image, char const* namespaze, char const* name);
         CUSTOM_TYPES_EXPORT static std::vector<TypeRegistration*> const& getTypes() { return registeredTypes; }
         /// @brief Automatically registers all pending types.
         /// To add a type to be registered, see: AddType

src/register.cpp

@@ -8,6 +8,7 @@
 #include "beatsaber-hook/shared/hooking.hpp"
 
 #include <filesystem>
+#include <limits>
 
 #ifdef CT_USE_GCDESCRIPTOR_DEBUG
 #include "liveness.hpp"
@@ -155,22 +156,17 @@ MAKE_HOOK(GetScriptingClass, (nullptr), Il2CppClass*, void* thisptr, char* assem
     return ret;
 }
 
-} // end anonymous namespace
+#ifdef UNITY_6_3
+MAKE_HOOK(Class_FromName, (nullptr), Il2CppClass*, Il2CppImage const* image, char const* namespaze, char const* name) {
+    // Unity can inline the metadata-handle conversion here, bypassing our handle hook.
+    if (auto klass = custom_types::Register::FindClass(image, namespaze, name)) {
+        return *klass;
+    }
+    return Class_FromName(image, namespaze, name);
+}
+#endif
 
-// NOTE THAT THIS HOOK DOES NOT PERMIT TYPES OF IDENTICAL NAMESPACE AND NAME BUT
-// IN DIFFERENT IMAGES! This could be worked around if we also have image be a
-// part of the key, but as it stands, that is not necessary.
-// MAKE_HOOK(Class_FromName, nullptr, Il2CppClass*, Il2CppImage* image, const
-// char* namespaze, const char* name) {
-//     pair = std::make_pair(std::string(namespaze), std::string(name)); auto
-//     itr = custom_types::Register::classMapping.find(pair); if (itr !=
-//     custom_types::Register::classMapping.end()) {
-//         #ifndef NO_VERBOSE_LOGS
-//         custom_types::logger.debug("Returning custom class from: {}::{} lookup: {}",
-//         namespaze, name, fmt::ptr(itr->second)); #endif return itr->second;
-//     }
-//     return Class_FromName(image, namespaze, name);
-// }
+} // end anonymous namespace
 
 namespace custom_types {
     std::unordered_map<std::string, Il2CppAssembly*> Register::assembs;
@@ -186,6 +182,36 @@ namespace custom_types {
     std::vector<TypeRegistration*> Register::toRegister;
     std::vector<TypeRegistration*> Register::registeredTypes;
 
+    std::optional<Il2CppClass*> Register::FindClass(Il2CppImage const* image, char const* namespaze, char const* name) {
+        if (!image || !image->dynamic) {
+            return std::nullopt;
+        }
+        auto owned = images.find(image->name);
+        if (owned == images.end() || owned->second != image) {
+            return std::nullopt;
+        }
+        auto const& table = *image->nameToClassHashTable;
+        auto found = table.find(std::make_pair(namespaze, name));
+        if (found == table.end()) {
+            return nullptr;
+        }
+        // Custom handles encode a negative 32-bit type index with zero upper bits.
+        auto handle = reinterpret_cast<uintptr_t>(found->second);
+        constexpr auto maxIndex = std::numeric_limits<uint32_t>::max();
+        if (handle > maxIndex || handle <= std::numeric_limits<int32_t>::max()) {
+            return nullptr;
+        }
+        size_t index = maxIndex - handle;
+        if (index >= classes.size()) {
+            return nullptr;
+        }
+        auto* klass = classes[index];
+        if (!klass || klass->image != image) {
+            return nullptr;
+        }
+        return klass;
+    }
+
     Il2CppAssembly* Register::createAssembly(std::string_view name, Il2CppImage* img) {
         // Name is NOT copied, so should be a constant string
         // Check to see if an assembly with the given name already exists.
@@ -302,14 +328,14 @@ namespace custom_types {
             } else {
                 logger.warn("Failed to find 1st bl in il2cpp_type_get_class_or_element_class!");
             }
-            // {
-            //     // We need to do a tiny bit of xref tracing to find the bottom level
-            //     Class::FromName call
-            //     // Trace is: il2cpp_class_from_name --> b --> b --> result
-            //     INSTALL_HOOK(logger, Class_FromName,
-            //     (void*)cs::findNthB<1>(reinterpret_cast<const
-            //     uint32_t*>(i2c::functions::class_from_name)));
-            // }
+#ifdef UNITY_6_3
+            // Unity 6.3 inlines handle conversion; intercept Image::ClassFromName first.
+            auto class_from_name = cs::find_nth_b<1, false, -1, 4>(reinterpret_cast<uint32_t const*>(i2c::functions::class_from_name));
+            if (!class_from_name) {
+                SAFE_ABORT("Failed to find Class::FromName from il2cpp_class_from_name");
+            }
+            INSTALL_HOOK(logger, Class_FromName, reinterpret_cast<void*>(*class_from_name));
+#endif
 
 #ifdef CT_USE_GCDESCRIPTOR_DEBUG
             liveness::EnsureHooks();

src/test.cpp

@@ -49,10 +49,6 @@ void Il2CppNamespace::MyType::ctor() {
 void Il2CppNamespace::MyType::Start() {
     logger.debug("Called Il2CppNamespace::MyType::Start!");
     logger.debug("Return of asdf(1): {}", asdf(1));
-    // Runtime invoke it.
-    // We ARE NOT able to call GetClassFromName.
-    // This is because our class name is NOT in the nameToClassHashTable
-    // However, we ARE able to get our Il2CppClass* via the klass private static field of this type.
     auto* il2cppKlass = i2c::get_class_from_name("Il2CppNamespace", "MyType");
     logger.debug("il2cpp obtained klass: {}", fmt::ptr(il2cppKlass));
     logger.debug("klass: {}", fmt::ptr(___TypeRegistration::klass_ptr));
@@ -122,8 +118,6 @@ void Il2CppNamespace::MyCustomRandom2::ctor(double original) {
 
 DECLARE_CLASS(SmallTest, Test, "System", "Object", sizeof(Il2CppObject)) {
     DECLARE_STATIC_METHOD(SmallTest::Test*, SelfRef, int);
-    DECLARE_STATIC_FIELD(SmallTest::Test*, selfRefField);
-    DECLARE_STATIC_FIELD(Il2CppNamespace::MyType*, AnotherRef);
 };
 
 DEFINE_TYPE(SmallTest, Test);
@@ -168,6 +162,21 @@ DEFINE_TYPE(SmallTest, TestIt3);
 
 static custom_types::ClassWrapper* klassWrapper;
 
+static void testClassLookup() {
+    auto check = [](Il2CppImage const* image, char const* namespaze, char const* name, Il2CppClass* expected) {
+        if (i2c::functions::class_from_name(image, namespaze, name) != expected) {
+            SAFE_ABORT("Class lookup regression: {}::{} in {}", namespaze, name, image->name);
+        }
+    };
+    auto* custom = i2c::class_of<Il2CppNamespace::MyType*>();
+    auto* object = i2c::class_of<Il2CppObject*>();
+    check(custom->image, custom->namespaze, custom->name, custom);
+    check(custom->image, custom->namespaze, "MissingType", nullptr);
+    check(i2c::class_of<Il2CppNamespace::MyTypeDllTest*>()->image, custom->namespaze, custom->name, nullptr);
+    check(object->image, object->namespaze, object->name, object);
+    logger.info("Class lookup tests passed");
+}
+
 CUSTOM_TYPES_FUNC void setup(CModInfo* info) {
     info->id = MOD_ID;
     info->version = VERSION;
@@ -220,6 +229,7 @@ CUSTOM_TYPES_FUNC void load() {
     logger.debug("Registering types! (current size: {})", custom_types::Register::classes.size());
     custom_types::Register::AutoRegister();
     logger.debug("Registered: {} types!", custom_types::Register::classes.size());
+    testClassLookup();
     INSTALL_HOOK(logger, MainMenuViewController_DidActivate);
     logger.debug("Custom types size: {}", custom_types::Register::classes.size());
     logger.debug("Logging vtables for custom type! There are: {} vtables", custom_types::Register::classes[0]->vtable_count);