Fix crash when debugger attached (#12)

* Fix crash when debugger attached

* 1

* 2

* 3

* 4

lsplant/src/main/jni/art/runtime/class_linker.hpp

@@ -1,8 +1,8 @@
 #pragma once
 
 #include "art/runtime/art_method.hpp"
-#include "art/runtime/obj_ptr.h"
-#include "art/thread.hpp"
+#include "art/runtime/obj_ptr.hpp"
+#include "art/runtime/thread.hpp"
 #include "common.hpp"
 
 namespace lsplant::art {
@@ -37,10 +37,8 @@ private:
                            });
 
     inline static art::ArtMethod *MayGetBackup(art::ArtMethod *method) {
-        std::shared_lock lk(hooked_methods_lock_);
-        if (auto found = hooked_methods_.find(method); found != hooked_methods_.end())
-            [[unlikely]] {
-            method = found->second.second;
+        if (auto backup = IsHooked(method); backup) [[unlikely]] {
+            method = backup;
             LOGV("propagate native method: %s", method->PrettyMethod(true).data());
         }
         return method;
@@ -102,9 +100,7 @@ private:
     static void FixTrampoline(const std::list<std::tuple<art::ArtMethod *, void *>> &methods) {
         std::shared_lock lk(hooked_methods_lock_);
         for (const auto &[art_method, old_trampoline] : methods) {
-            if (auto found = hooked_methods_.find(art_method); found != hooked_methods_.end())
-                [[likely]] {
-                auto &backup_method = found->second.second;
+            if (auto backup_method = IsHooked(art_method); backup_method) [[likely]] {
                 if (auto new_trampoline = art_method->GetEntryPoint();
                     new_trampoline != old_trampoline) [[unlikely]] {
                     LOGV("propagate entrypoint for %s", backup_method->PrettyMethod(true).data());

lsplant/src/main/jni/art/runtime/gc/scoped_gc_critical_section.hpp

@@ -1,6 +1,6 @@
 #pragma once
 
-#include "art/thread.hpp"
+#include "art/runtime/thread.hpp"
 #include "collector_type.hpp"
 #include "common.hpp"
 #include "gc_cause.hpp"

lsplant/src/main/jni/art/runtime/instrumentation.hpp

@@ -0,0 +1,46 @@
+#pragma once
+
+#include "art_method.hpp"
+#include "common.hpp"
+
+namespace lsplant::art {
+
+class Instrumentation {
+    inline static ArtMethod *MaybeUseBackupMethod(ArtMethod *art_method, const void *quick_code) {
+        if (auto backup = IsHooked(art_method); backup && art_method->GetEntryPoint() != quick_code)
+            [[unlikely]] {
+            LOGD("Propagate update method code %p for hooked method %s to its backup", quick_code,
+                 art_method->PrettyMethod().c_str());
+            return backup;
+        }
+        return art_method;
+    }
+
+    CREATE_MEM_HOOK_STUB_ENTRY(
+        "_ZN3art15instrumentation15Instrumentation40UpdateMethodsCodeToInterpreterEntryPointEPNS_9ArtMethodE",
+        void, UpdateMethodsCodeToInterpreterEntryPoint,
+        (Instrumentation * thiz, ArtMethod *art_method),
+        { backup(thiz, MaybeUseBackupMethod(art_method, nullptr)); });
+
+    CREATE_MEM_HOOK_STUB_ENTRY(
+            "_ZN3art15instrumentation15Instrumentation21InitializeMethodsCodeEPNS_9ArtMethodEPKv",
+            void, InitializeMethodsCode,
+            (Instrumentation * thiz, ArtMethod *art_method, const void* quick_code),
+            { backup(thiz, MaybeUseBackupMethod(art_method, quick_code), quick_code); });
+
+public:
+    static bool Init(JNIEnv *env, const HookHandler &handler) {
+        if (!IsJavaDebuggable(env)) [[likely]] {
+            return true;
+        }
+        int sdk_int = GetAndroidApiLevel();
+        if (sdk_int >= __ANDROID_API_P__) [[likely]] {
+            if (!HookSyms(handler, InitializeMethodsCode, UpdateMethodsCodeToInterpreterEntryPoint)) {
+                return false;
+            }
+        }
+        return true;
+    }
+};
+
+}  // namespace lsplant::art

lsplant/src/main/jni/art/runtime/jit/jit_code_cache.hpp

@@ -17,8 +17,10 @@ class JitCodeCache {
 
     CREATE_MEM_HOOK_STUB_ENTRY("_ZN3art3jit12JitCodeCache19GarbageCollectCacheEPNS_6ThreadE", void,
                                GarbageCollectCache, (JitCodeCache * thiz, Thread *self), {
-                                   LOGD("Before jit cache gc, moving hooked methods");
-                                   for (auto [target, backup] : GetJitMovements()) {
+                                   auto movements = GetJitMovements();
+                                   LOGD("Before jit cache gc, moving %zu hooked methods",
+                                        movements.size());
+                                   for (auto [target, backup] : movements) {
                                        MoveObsoleteMethod(thiz, target, backup);
                                    }
                                    backup(thiz, self);

lsplant/src/main/jni/art/runtime/jni/jni_id_manager.h

@@ -0,0 +1,33 @@
+#pragma once
+
+#include "art/runtime/art_method.hpp"
+#include "art/runtime/reflective_handle.hpp"
+#include "common.hpp"
+
+namespace lsplant::art::jni {
+
+class JniIdManager {
+private:
+    CREATE_MEM_HOOK_STUB_ENTRY(
+        "_ZN3art3jni12JniIdManager15EncodeGenericIdINS_9ArtMethodEEEmNS_16ReflectiveHandleIT_EE",
+        uintptr_t, EncodeGenericId, (JniIdManager * thiz, ReflectiveHandle<ArtMethod> method), {
+            if (auto target = IsBackup(method.Get()); target) {
+                LOGD("get generic id for %s", method.Get()->PrettyMethod().c_str());
+                method.Set(target);
+            }
+            return backup(thiz, method);
+        });
+
+public:
+    static bool Init(JNIEnv *env, const HookHandler &handler) {
+        int sdk_int = GetAndroidApiLevel();
+        if (sdk_int >= __ANDROID_API_R__) {
+            if (IsJavaDebuggable(env) && !HookSyms(handler, EncodeGenericId)) {
+                LOGW("Failed to hook EncodeGenericId, attaching debugger may crash the process");
+            }
+        }
+        return true;
+    }
+};
+
+}  // namespace lsplant::art::jni

lsplant/src/main/jni/art/runtime/obj_ptr.h

@@ -1,12 +0,0 @@
-#pragma once
-
-template <typename MirrorType>
-class ObjPtr {
-public:
-    inline MirrorType* operator->() const { return Ptr(); }
-    inline MirrorType* Ptr() const { return reference_; }
-    inline operator MirrorType*() const { return Ptr(); }
-
-private:
-    MirrorType* reference_;
-};

lsplant/src/main/jni/art/runtime/obj_ptr.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+namespace lsplant::art {
+
+template <typename MirrorType>
+class ObjPtr {
+public:
+    inline MirrorType *operator->() const { return Ptr(); }
+
+    inline MirrorType *Ptr() const { return reference_; }
+
+    inline operator MirrorType *() const { return Ptr(); }
+
+private:
+    MirrorType *reference_;
+};
+
+}  // namespace lsplant::art

lsplant/src/main/jni/art/runtime/reflective_handle.hpp

@@ -0,0 +1,23 @@
+#pragma once
+
+#include "reflective_reference.hpp"
+
+namespace lsplant::art {
+
+class ArtMethod;
+class ValueObject {};
+
+template <typename T>
+class ReflectiveHandle : public ValueObject {
+public:
+    static_assert(std::is_same_v<T, ArtMethod>, "Expected ArtField or ArtMethod");
+
+    T *Get() { return reference_->Ptr(); }
+
+    void Set(T *val) { reference_->Assign(val); }
+
+protected:
+    ReflectiveReference<T> *reference_;
+};
+
+}  // namespace lsplant::art

lsplant/src/main/jni/art/runtime/reflective_reference.hpp

@@ -0,0 +1,17 @@
+#pragma once
+
+namespace lsplant::art {
+template <class ReflectiveType>
+class ReflectiveReference {
+public:
+    static_assert(std::is_same_v<ReflectiveType, ArtMethod>, "Unknown type!");
+
+    ReflectiveType *Ptr() { return val_; }
+
+    void Assign(ReflectiveType *r) { val_ = r; }
+
+private:
+    ReflectiveType *val_;
+};
+
+}  // namespace lsplant::art

lsplant/src/main/jni/common.hpp

@@ -53,6 +53,37 @@ inline auto GetAndroidApiLevel() {
     return kApiLevel;
 }
 
+inline auto IsJavaDebuggable(JNIEnv *env) {
+    static auto kDebuggable = [&env]() {
+        auto runtime_class = JNI_FindClass(env, "dalvik/system/VMRuntime");
+        if (!runtime_class) {
+            LOGE("Failed to find VMRuntime");
+            return false;
+        }
+        auto get_runtime_method = JNI_GetStaticMethodID(env, runtime_class, "getRuntime",
+                                                        "()Ldalvik/system/VMRuntime;");
+        if (!get_runtime_method) {
+            LOGE("Failed to find VMRuntime.getRuntime()");
+            return false;
+        }
+        auto is_debuggable_method =
+                JNI_GetMethodID(env, runtime_class, "isJavaDebuggable", "()Z");
+        if (!is_debuggable_method) {
+            LOGE("Failed to find VMRuntime.isJavaDebuggable()");
+            return false;
+        }
+        auto runtime = JNI_CallStaticObjectMethod(env, runtime_class, get_runtime_method);
+        if (!runtime) {
+            LOGE("Failed to get VMRuntime");
+            return false;
+        }
+        bool is_debuggable = JNI_CallBooleanMethod(env, runtime, is_debuggable_method);
+        LOGD("java runtime debuggable %s", is_debuggable ? "true" : "false");
+        return is_debuggable;
+    }();
+    return kDebuggable;
+}
+
 inline static constexpr auto kPointerSize = sizeof(void *);
 
 namespace art {
@@ -78,9 +109,22 @@ inline std::unordered_map<const art::dex::ClassDef *, std::unordered_set<art::Ar
 inline std::shared_mutex hooked_classes_lock_;
 }  // namespace
 
-inline bool IsHooked(art::ArtMethod *art_method) {
+inline art::ArtMethod *IsHooked(art::ArtMethod *art_method, bool including_backup = false) {
     std::shared_lock lk(hooked_methods_lock_);
-    return hooked_methods_.contains(art_method);
+    if (auto it = hooked_methods_.find(art_method);
+        it != hooked_methods_.end() && (!including_backup || it->second.first)) {
+        return it->second.second;
+    }
+    return nullptr;
+}
+
+inline art::ArtMethod *IsBackup(art::ArtMethod *art_method) {
+    std::shared_lock lk(hooked_methods_lock_);
+    if (auto it = hooked_methods_.find(art_method);
+        it != hooked_methods_.end() && !it->second.first) {
+        return it->second.second;
+    }
+    return nullptr;
 }
 
 inline std::list<std::pair<art::ArtMethod *, art::ArtMethod *>> GetJitMovements() {
@@ -93,6 +137,7 @@ inline void RecordHooked(art::ArtMethod *target, const art::dex::ClassDef *class
     {
         std::unique_lock lk(hooked_methods_lock_);
         hooked_methods_[target] = {reflected_backup, backup};
+        hooked_methods_[backup] = {nullptr, target};
     }
     {
         std::unique_lock lk(hooked_classes_lock_);

lsplant/src/main/jni/include/lsplant.hpp

@@ -97,6 +97,7 @@ struct InitInfo {
 /// call on this function with the same \p target_method does not guarantee only one will success.
 /// If you call this with different \p hooker_object on the same target_method simultaneously, the
 /// behavior is undefined.
+/// \note The behavior of getting the \ref jmethodID of the backup method is undfined.
 [[nodiscard, maybe_unused, gnu::visibility("default")]] jobject Hook(JNIEnv *env,
                                                                      jobject target_method,
                                                                      jobject hooker_object,

lsplant/src/main/jni/lsplant.cc

@@ -12,9 +12,11 @@
 #include "art/runtime/class_linker.hpp"
 #include "art/runtime/dex_file.hpp"
 #include "art/runtime/gc/scoped_gc_critical_section.hpp"
+#include "art/runtime/instrumentation.hpp"
 #include "art/runtime/jit/jit_code_cache.hpp"
-#include "art/thread.hpp"
-#include "art/thread_list.hpp"
+#include "art/runtime/jni/jni_id_manager.h"
+#include "art/runtime/thread.hpp"
+#include "art/runtime/thread_list.hpp"
 #include "common.hpp"
 #include "dex_builder.h"
 #include "utils/jni_helper.hpp"
@@ -29,9 +31,11 @@ namespace lsplant {
 using art::ArtMethod;
 using art::ClassLinker;
 using art::DexFile;
+using art::Instrumentation;
 using art::Thread;
 using art::gc::ScopedGCCriticalSection;
 using art::jit::JitCodeCache;
+using art::jni::JniIdManager;
 using art::mirror::Class;
 using art::thread_list::ScopedSuspendAll;
 
@@ -246,6 +250,14 @@ bool InitNative(JNIEnv *env, const HookHandler &handler) {
         LOGE("Failed to init dex file");
         return false;
     }
+    if (!Instrumentation::Init(env, handler)) {
+        LOGE("Failed to init instrumentation");
+        return false;
+    }
+    if (!JniIdManager::Init(env, handler)) {
+        LOGE("Failed to init jni id manager");
+        return false;
+    }
     return true;
 }
 
@@ -533,7 +545,7 @@ using ::lsplant::IsHooked;
     auto *target = ArtMethod::FromReflectedMethod(env, target_method);
     bool is_static = target->IsStatic();
 
-    if (IsHooked(target)) {
+    if (IsHooked(target, true)) {
         LOGW("Skip duplicate hook");
         return nullptr;
     }
@@ -600,9 +612,14 @@ using ::lsplant::IsHooked;
     art::ArtMethod *backup = nullptr;
     {
         std::unique_lock lk(hooked_methods_lock_);
-        if (auto it = hooked_methods_.find(target); it != hooked_methods_.end()) {
+        if (auto it = hooked_methods_.find(target); it != hooked_methods_.end()) [[likely]] {
             std::tie(reflected_backup, backup) = it->second;
+            if (reflected_backup == nullptr) {
+                LOGE("Unable to unhook a method that is not hooked");
+                return false;
+            }
             hooked_methods_.erase(it);
+            hooked_methods_.erase(it->second.second);
         }
     }
     {
@@ -615,10 +632,6 @@ using ::lsplant::IsHooked;
             }
         }
     }
-    if (reflected_backup == nullptr) {
-        LOGE("Unable to unhook a method that is not hooked");
-        return false;
-    }
     env->DeleteGlobalRef(reflected_backup);
     return DoUnHook(target, backup);
 }
@@ -629,8 +642,7 @@ using ::lsplant::IsHooked;
         return false;
     }
     auto *art_method = ArtMethod::FromReflectedMethod(env, method);
-    std::shared_lock lk(hooked_methods_lock_);
-    return hooked_methods_.contains(art_method);
+    return IsHooked(art_method);
 }
 
 [[maybe_unused]] bool Deoptimize(JNIEnv *env, jobject method) {
@@ -639,12 +651,8 @@ using ::lsplant::IsHooked;
         return false;
     }
     auto *art_method = ArtMethod::FromReflectedMethod(env, method);
-    if (IsHooked(art_method)) {
-        std::shared_lock lk(hooked_methods_lock_);
-        auto it = hooked_methods_.find(art_method);
-        if (it != hooked_methods_.end()) {
-            art_method = it->second.second;
-        }
+    if (auto *backup = IsHooked(art_method); backup) {
+        art_method = backup;
     }
     if (!art_method) {
         return false;
@@ -674,7 +682,7 @@ using ::lsplant::IsHooked;
     uint8_t access_flags = JNI_GetIntField(env, target, class_access_flags);
     constexpr static uint32_t kAccFinal = 0x0010;
     JNI_SetIntField(env, target, class_access_flags, static_cast<jint>(access_flags & ~kAccFinal));
-    for (auto &constructor : constructors) {
+    for (const auto &constructor : constructors) {
         auto *method = ArtMethod::FromReflectedMethod(env, constructor.get());
         if (method && !method->IsPublic() && !method->IsProtected()) method->SetProtected();
         if (method && method->IsFinal()) method->SetNonFinal();