From f84db35a11941d4d2f382fab6a7164cc925eaa13 Mon Sep 17 00:00:00 2001
From: LoveSy <shana@zju.edu.cn>
Date: Sun, 13 Dec 2020 14:49:25 +0800
Subject: [PATCH] Thread safe

---
 .../main/cpp/external/yahfa/src/HookMain.c    | 12 +++++++
 .../main/cpp/external/yahfa/src/trampoline.c  | 35 +++++++++++--------
 .../main/cpp/external/yahfa/src/trampoline.h  |  5 ++-
 3 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/edxp-core/src/main/cpp/external/yahfa/src/HookMain.c b/edxp-core/src/main/cpp/external/yahfa/src/HookMain.c
index b95b48fe..7dc54184 100644
--- a/edxp-core/src/main/cpp/external/yahfa/src/HookMain.c
+++ b/edxp-core/src/main/cpp/external/yahfa/src/HookMain.c
@@ -115,6 +115,16 @@ void *getEntryPoint(void* method) {
 }
 
 static int replaceMethod(void *fromMethod, void *toMethod, int isBackup) {
+    if (hookCount >= hookCap) {
+        LOGI("not enough capacity. Allocating...");
+        if (doInitHookCap()) {
+            LOGE("cannot hook method");
+            return 1;
+        }
+        LOGI("Allocating done");
+    }
+
+    LOGI("replace method from %p to %p", fromMethod, toMethod);
     // replace entry point
     void *newEntrypoint = NULL;
     if(isBackup) {
@@ -146,6 +156,8 @@ static int replaceMethod(void *fromMethod, void *toMethod, int isBackup) {
                pointer_size);
     }
 
+    hookCount += 1;
+
     return 0;
 }
 
diff --git a/edxp-core/src/main/cpp/external/yahfa/src/trampoline.c b/edxp-core/src/main/cpp/external/yahfa/src/trampoline.c
index da3523b6..58360b6c 100644
--- a/edxp-core/src/main/cpp/external/yahfa/src/trampoline.c
+++ b/edxp-core/src/main/cpp/external/yahfa/src/trampoline.c
@@ -16,6 +16,10 @@
 
 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 
+static _Thread_local unsigned char *trampolineCode; // place where trampolines are saved
+
+_Thread_local unsigned int hookCap = 1;
+_Thread_local unsigned int hookCount = 1;
 
 // trampoline:
 // 1. set eax/rdi/r0/x0 to the hook ArtMethod addr
@@ -123,22 +127,22 @@ unsigned char trampolineForBackup[] = {
 
 #endif
 
+// trampoline size required for each hook
+static unsigned int trampolineSize = roundUpToPtrSize(MAX(sizeof(trampoline), sizeof(trampolineForBackup)));
+
 static inline void FlushCache(void *addr, size_t size) {
     __builtin___clear_cache((char *) addr, (char *) ((uintptr_t) addr + size));
 }
 
 void *genTrampoline(void *toMethod, void *entrypoint) {
-    size_t size = entrypoint == NULL ? sizeof(trampoline) : sizeof(trampolineForBackup);
-
-    // TODO: make use of thread_local to avoid frequent memory allocate
-    void *targetAddr = doInitHookCap(size);
+    unsigned char *targetAddr = trampolineCode + trampolineSize * hookCount;
 
     if (targetAddr == NULL) return NULL;
 
     if (entrypoint != NULL) {
-        memcpy(targetAddr, trampolineForBackup, size);
+        memcpy(targetAddr, trampolineForBackup, sizeof(trampolineForBackup));
     } else {
-        memcpy(targetAddr, trampoline, size);
+        memcpy(targetAddr, trampoline, sizeof(trampoline)); // do not use trampolineSize since it's a rounded size
     }
 
     // replace with the actual ArtMethod addr
@@ -180,7 +184,7 @@ void *genTrampoline(void *toMethod, void *entrypoint) {
 #else
 #error Unsupported architecture
 #endif
-    FlushCache(targetAddr, size);
+    FlushCache(targetAddr, trampolineSize);
 
     return targetAddr;
 }
@@ -200,17 +204,20 @@ void setupTrampoline(uint8_t offset) {
 #endif
 }
 
-void *doInitHookCap(size_t size) {
-    if (size == 0) {
-        LOGE("invalid capacity: %zx", size);
-        return NULL;
+int doInitHookCap() {
+    if (hookCap > hookCount) {
+        return 0;
     }
-    unsigned char *buf = mmap(NULL, size, PROT_READ | PROT_WRITE | PROT_EXEC,
+    hookCap *= 2;
+    unsigned int allSize = trampolineSize * hookCap;
+    unsigned char *buf = mmap(NULL, allSize, PROT_READ | PROT_WRITE | PROT_EXEC,
                               MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
     if (buf == MAP_FAILED) {
         LOGE("mmap failed, errno = %s", strerror(errno));
-        return NULL;
+        return 1;
     }
-    return buf;
+    hookCount = 0;
+    trampolineCode = buf;
+    return 0;
 }
 
diff --git a/edxp-core/src/main/cpp/external/yahfa/src/trampoline.h b/edxp-core/src/main/cpp/external/yahfa/src/trampoline.h
index f5fca72a..eea4ea19 100644
--- a/edxp-core/src/main/cpp/external/yahfa/src/trampoline.h
+++ b/edxp-core/src/main/cpp/external/yahfa/src/trampoline.h
@@ -7,9 +7,12 @@
 
 extern int SDKVersion;
 
+extern _Thread_local unsigned int hookCap; // capacity for trampolines
+extern _Thread_local unsigned int hookCount; // current count of used trampolines
+
 extern unsigned char trampoline[];
 
-void* doInitHookCap(size_t cap);
+int doInitHookCap();
 void setupTrampoline(uint8_t offset);
 void *genTrampoline(void *toMethod, void *entrypoint);