From 2327e7b398699ea24179e0f8b3a08005030abfbc Mon Sep 17 00:00:00 2001
From: LoveSy <shana@zju.edu.cn>
Date: Sat, 3 Dec 2022 02:37:23 +0800
Subject: [PATCH] Guard backup during hook (#2263)

---
 core/src/main/jni/src/jni/hook_bridge.cpp | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/core/src/main/jni/src/jni/hook_bridge.cpp b/core/src/main/jni/src/jni/hook_bridge.cpp
index ec1cabd5..b789ff12 100644
--- a/core/src/main/jni/src/jni/hook_bridge.cpp
+++ b/core/src/main/jni/src/jni/hook_bridge.cpp
@@ -23,6 +23,7 @@
 #include <absl/container/flat_hash_map.h>
 #include <memory>
 #include <shared_mutex>
+#include <mutex>
 #include <set>
 
 using namespace lsplant;
@@ -35,6 +36,7 @@ struct HookItem {
 };
 
 std::shared_mutex hooked_lock;
+std::recursive_mutex backup_lock;
 absl::flat_hash_map<jmethodID, std::unique_ptr<HookItem>> hooked_methods;
 
 jmethodID invoke = nullptr;
@@ -83,6 +85,7 @@ LSP_DEF_NATIVE_METHOD(jboolean, HookBridge, hookMethod, jobject hookMethod,
                                                                                "([Ljava/lang/Object;)Ljava/lang/Object;"),
                                                       false);
         auto hooker_object = env->NewObject(hooker, init, hookMethod);
+        std::unique_lock lk(backup_lock);
         hook_item->backup = lsplant::Hook(env, hookMethod, hooker_object, callback_method);
         env->DeleteLocalRef(hooker_object);
     }
@@ -101,7 +104,11 @@ LSP_DEF_NATIVE_METHOD(jboolean, HookBridge, unhookMethod, jobject hookMethod, jo
         }
     }
     if (!hook_item) return JNI_FALSE;
-    auto backup = hook_item->backup;
+    jobject backup = nullptr;
+    {
+        std::unique_lock lk(backup_lock);
+        backup = hook_item->backup;
+    }
     if (!backup) return JNI_FALSE;
     JNIMonitor monitor(env, backup);
     for (auto i = hook_item->callbacks.begin(); i != hook_item->callbacks.end(); ++i) {
@@ -129,8 +136,9 @@ LSP_DEF_NATIVE_METHOD(jobject, HookBridge, invokeOriginalMethod, jobject hookMet
         }
     }
     jobject to_call = hookMethod;
-    if (hook_item && hook_item->backup) {
-        to_call = hook_item->backup;
+    if (hook_item) {
+        std::unique_lock lk(backup_lock);
+        if (hook_item->backup) to_call = hook_item->backup;
     }
     return env->CallObjectMethod(to_call, invoke, thiz, args);
 }
@@ -153,7 +161,11 @@ LSP_DEF_NATIVE_METHOD(jobjectArray, HookBridge, callbackSnapshot, jobject method
         }
     }
     if (!hook_item) return nullptr;
-    auto backup = hook_item->backup;
+    jobject backup = nullptr;
+    {
+        std::unique_lock lk(backup_lock);
+        backup = hook_item->backup;
+    }
     if (!backup) return nullptr;
     JNIMonitor monitor(env, backup);
     auto res = env->NewObjectArray((jsize) hook_item->callbacks.size(), env->FindClass("java/lang/Object"), nullptr);