From 263565175a29097048308048a3923b52e56d75ba Mon Sep 17 00:00:00 2001
From: LoveSy <shana@zju.edu.cn>
Date: Wed, 16 Feb 2022 14:58:27 +0800
Subject: [PATCH] Allow hook&backup for unhooked function to be GCed

---
 library/jni/common.hpp |  4 ++--
 library/jni/lsplant.cc | 12 ++++++------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/library/jni/common.hpp b/library/jni/common.hpp
index 5909c7e..4c49bda 100644
--- a/library/jni/common.hpp
+++ b/library/jni/common.hpp
@@ -74,7 +74,7 @@ class ClassDef;
 
 namespace {
 // target, backup
-inline std::unordered_map<art::ArtMethod *, art::ArtMethod *> hooked_methods_;
+inline std::unordered_map<art::ArtMethod *, jobject> hooked_methods_;
 inline std::shared_mutex hooked_methods_lock_;
 
 inline std::list<std::pair<art::ArtMethod *, art::ArtMethod *>> jit_movements_;
@@ -107,7 +107,7 @@ inline std::list<std::pair<art::ArtMethod *, art::ArtMethod *>> GetJitMovements(
     return std::move(jit_movements_);
 }
 
-inline void RecordHooked(art::ArtMethod *target, art::ArtMethod *backup) {
+inline void RecordHooked(art::ArtMethod *target, jobject backup) {
     std::unique_lock lk(hooked_methods_lock_);
     hooked_methods_.emplace(target, backup);
 }
diff --git a/library/jni/lsplant.cc b/library/jni/lsplant.cc
index 433f0de..1101f9a 100644
--- a/library/jni/lsplant.cc
+++ b/library/jni/lsplant.cc
@@ -394,7 +394,6 @@ bool DoHook(ArtMethod *target, ArtMethod *hook, ArtMethod *backup) {
              backup, backup->GetAccessFlags(), backup->GetEntryPoint(),
              hook, hook->GetAccessFlags(), hook->GetEntryPoint());
 
-        RecordHooked(target, backup);
         return true;
     }
 }
@@ -498,8 +497,7 @@ Hook(JNIEnv *env, jmethodID target_method, jobject hooker_object, jmethodID call
         RecordPending(class_def, target, hook, backup);
         return backup_method;
     } else if (DoHook(target, hook, backup)) {
-        JNI_NewGlobalRef(env, reflected_hook);
-        JNI_NewGlobalRef(env, reflected_backup);
+        RecordHooked(target, JNI_NewGlobalRef(env, reflected_backup));
         if (!is_proxy) [[likely]] RecordJitMovement(target, backup);
         return backup_method;
     }
@@ -511,7 +509,7 @@ Hook(JNIEnv *env, jmethodID target_method, jobject hooker_object, jmethodID call
 bool UnHook(JNIEnv *env, jmethodID target_method) {
     auto reflected_target = JNI_ToReflectedMethod(env, jclass{ nullptr }, target_method, false);
     auto *target = ArtMethod::FromReflectedMethod(env, reflected_target);
-    ArtMethod *backup = nullptr;
+    jobject reflected_backup = nullptr;
     {
         std::unique_lock lk(pending_methods_lock_);
         if (auto it = pending_methods_.find(target); it != pending_methods_.end()) {
@@ -522,14 +520,16 @@ bool UnHook(JNIEnv *env, jmethodID target_method) {
     {
         std::unique_lock lk(hooked_methods_lock_);
         if (auto it = hooked_methods_.find(target);it != hooked_methods_.end()) {
-            backup = it->second;
+            reflected_backup = it->second;
             hooked_methods_.erase(it);
         }
     }
-    if (backup == nullptr) {
+    if (reflected_backup == nullptr) {
         LOGE("Unable to unhook a method that is not hooked");
         return false;
     }
+    auto *backup = ArtMethod::FromReflectedMethod(env, reflected_backup);
+    env->DeleteGlobalRef(reflected_backup);
     return DoUnHook(target, backup);
 }