From f7da5f381d814bfa9fc4e68ff04e9484d8614aec Mon Sep 17 00:00:00 2001
From: LoveSy <shana@zju.edu.cn>
Date: Wed, 3 Feb 2021 02:01:30 +0800
Subject: [PATCH] [core] Lock critical section

---
 .../main/cpp/main/src/jni/pending_hooks.cpp   | 41 +++++++++++++++----
 1 file changed, 34 insertions(+), 7 deletions(-)

diff --git a/core/src/main/cpp/main/src/jni/pending_hooks.cpp b/core/src/main/cpp/main/src/jni/pending_hooks.cpp
index 1f4f21db..2101463f 100644
--- a/core/src/main/cpp/main/src/jni/pending_hooks.cpp
+++ b/core/src/main/cpp/main/src/jni/pending_hooks.cpp
@@ -2,6 +2,7 @@
 #include <nativehelper/jni_macros.h>
 #include <string>
 #include <unordered_set>
+#include <shared_mutex>
 #include "HookMain.h"
 #include "jni.h"
 #include "native_util.h"
@@ -12,32 +13,55 @@
 namespace lspd {
     namespace {
         std::unordered_set<const void *> pending_classes_;
+        std::shared_mutex pending_classes_lock_;
+
         std::unordered_set<const void *> pending_methods_;
+        std::shared_mutex pending_methods_lock_;
 
         std::unordered_set<const void *> hooked_methods_;
+        std::shared_mutex hooked_methods_lock_;
     }
 
     bool IsClassPending(void *clazz) {
+        std::shared_lock lk(pending_classes_lock_);
         return pending_classes_.count(clazz);
     }
 
-    bool IsMethodPending(void* art_method) {
-        return pending_methods_.erase(art_method) > 0;
+    bool IsMethodPending(void *art_method) {
+        bool result;
+        {
+            std::shared_lock lk(pending_methods_lock_);
+            result = pending_methods_.count(art_method);
+        }
+        if (result) {
+            std::unique_lock lk(pending_methods_lock_);
+            pending_methods_.erase(art_method);
+        }
+        return result;
     }
 
     void DonePendingHook(void *clazz) {
+        std::unique_lock lk(pending_classes_lock_);
         pending_classes_.erase(clazz);
     }
 
-    static void PendingHooks_recordPendingMethodNative(JNI_START, jobject method_ref, jclass class_ref) {
+    static void
+    PendingHooks_recordPendingMethodNative(JNI_START, jobject method_ref, jclass class_ref) {
         auto *class_ptr = art::Thread::Current().DecodeJObject(class_ref);
         auto *method = getArtMethodYahfa(env, method_ref);
         art::mirror::Class mirror_class(class_ptr);
         if (auto def = mirror_class.GetClassDef(); LIKELY(def)) {
-            LOGD("record pending: %p (%s) with %p", class_ptr, mirror_class.GetDescriptor().c_str(), method);
+            LOGD("record pending: %p (%s) with %p", class_ptr, mirror_class.GetDescriptor().c_str(),
+                 method);
             // Add it for ShouldUseInterpreterEntrypoint
-            pending_methods_.insert(method);
-            pending_classes_.insert(def);
+            {
+                std::unique_lock lk(pending_methods_lock_);
+                pending_methods_.insert(method);
+            }
+            {
+                std::unique_lock lk(pending_classes_lock_);
+                pending_classes_.insert(def);
+            }
         } else {
             LOGW("fail to record pending for : %p (%s)", class_ptr,
                  mirror_class.GetDescriptor().c_str());
@@ -45,7 +69,8 @@ namespace lspd {
     }
 
     static JNINativeMethod gMethods[] = {
-            NATIVE_METHOD(PendingHooks, recordPendingMethodNative, "(Ljava/lang/reflect/Method;Ljava/lang/Class;)V"),
+            NATIVE_METHOD(PendingHooks, recordPendingMethodNative,
+                          "(Ljava/lang/reflect/Method;Ljava/lang/Class;)V"),
     };
 
     void RegisterPendingHooks(JNIEnv *env) {
@@ -53,10 +78,12 @@ namespace lspd {
     }
 
     bool isHooked(void *art_method) {
+        std::shared_lock lk(hooked_methods_lock_);
         return hooked_methods_.count(art_method);
     }
 
     void recordHooked(void *art_method) {
+        std::unique_lock lk(hooked_methods_lock_);
         hooked_methods_.insert(art_method);
     }