From 6dd10b1fb7a92134f1b8bcd0c87ad27da1ffd4b5 Mon Sep 17 00:00:00 2001
From: LoveSy <shana@zju.edu.cn>
Date: Sun, 13 Dec 2020 14:25:18 +0800
Subject: [PATCH] Flush CPU cache in YAHFA

---
 .../main/cpp/external/yahfa/src/HookMain.c    | 12 ----
 .../main/cpp/external/yahfa/src/trampoline.c  | 55 +++++++++----------
 .../main/cpp/external/yahfa/src/trampoline.h  |  7 +--
 3 files changed, 26 insertions(+), 48 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 e9f341b7..b95b48fe 100644
--- a/edxp-core/src/main/cpp/external/yahfa/src/HookMain.c
+++ b/edxp-core/src/main/cpp/external/yahfa/src/HookMain.c
@@ -115,17 +115,6 @@ void *getEntryPoint(void* method) {
 }
 
 static int replaceMethod(void *fromMethod, void *toMethod, int isBackup) {
-    if (hookCount >= hookCap) {
-        LOGI("not enough capacity. Allocating...");
-        if (doInitHookCap(DEFAULT_CAP)) {
-            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) {
@@ -157,7 +146,6 @@ 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 2652a8c8..da3523b6 100644
--- a/edxp-core/src/main/cpp/external/yahfa/src/trampoline.c
+++ b/edxp-core/src/main/cpp/external/yahfa/src/trampoline.c
@@ -16,11 +16,6 @@
 
 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 
-static unsigned char *trampolineCode; // place where trampolines are saved
-static unsigned int trampolineSize; // trampoline size required for each hook
-
-unsigned int hookCap = 0;
-unsigned int hookCount = 0;
 
 // trampoline:
 // 1. set eax/rdi/r0/x0 to the hook ArtMethod addr
@@ -127,17 +122,23 @@ unsigned char trampolineForBackup[] = {
 };
 
 #endif
-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) {
-    unsigned char *targetAddr = trampolineCode + trampolineSize * hookCount;
+    size_t size = entrypoint == NULL ? sizeof(trampoline) : sizeof(trampolineForBackup);
 
-    if(entrypoint != NULL) {
-        memcpy(targetAddr, trampolineForBackup, sizeof(trampolineForBackup));
-    }
-    else {
-        memcpy(targetAddr, trampoline,
-               sizeof(trampoline)); // do not use trampolineSize since it's a rounded size
+    // TODO: make use of thread_local to avoid frequent memory allocate
+    void *targetAddr = doInitHookCap(size);
+
+    if (targetAddr == NULL) return NULL;
+
+    if (entrypoint != NULL) {
+        memcpy(targetAddr, trampolineForBackup, size);
+    } else {
+        memcpy(targetAddr, trampoline, size);
     }
 
     // replace with the actual ArtMethod addr
@@ -169,17 +170,17 @@ void *genTrampoline(void *toMethod, void *entrypoint) {
     }
 
 #elif defined(__aarch64__)
-    if(entrypoint) {
+    if (entrypoint) {
         memcpy(targetAddr + 20, &entrypoint, pointer_size);
         memcpy(targetAddr + 12, &toMethod, pointer_size);
-    }
-    else {
+    } else {
         memcpy(targetAddr + 12, &toMethod, pointer_size);
     }
 
 #else
 #error Unsupported architecture
 #endif
+    FlushCache(targetAddr, size);
 
     return targetAddr;
 }
@@ -199,23 +200,17 @@ void setupTrampoline(uint8_t offset) {
 #endif
 }
 
-int doInitHookCap(unsigned int cap) {
-    if (cap == 0) {
-        LOGE("invalid capacity: %d", cap);
-        return 1;
+void *doInitHookCap(size_t size) {
+    if (size == 0) {
+        LOGE("invalid capacity: %zx", size);
+        return NULL;
     }
-    if (hookCap) {
-        LOGI("allocating new space for trampoline code");
-    }
-    unsigned int allSize = trampolineSize * cap;
-    unsigned char *buf = mmap(NULL, allSize, PROT_READ | PROT_WRITE | PROT_EXEC,
+    unsigned char *buf = mmap(NULL, size, PROT_READ | PROT_WRITE | PROT_EXEC,
                               MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
     if (buf == MAP_FAILED) {
         LOGE("mmap failed, errno = %s", strerror(errno));
-        return 1;
+        return NULL;
     }
-    hookCap = cap;
-    hookCount = 0;
-    trampolineCode = buf;
-    return 0;
+    return buf;
 }
+
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 2235a42d..f5fca72a 100644
--- a/edxp-core/src/main/cpp/external/yahfa/src/trampoline.h
+++ b/edxp-core/src/main/cpp/external/yahfa/src/trampoline.h
@@ -7,15 +7,10 @@
 
 extern int SDKVersion;
 
-extern unsigned int hookCap; // capacity for trampolines
-extern unsigned int hookCount; // current count of used trampolines
-
 extern unsigned char trampoline[];
 
-int doInitHookCap(unsigned int cap);
+void* doInitHookCap(size_t cap);
 void setupTrampoline(uint8_t offset);
 void *genTrampoline(void *toMethod, void *entrypoint);
 
-#define DEFAULT_CAP 1 //size of each trampoline area would be no more than 4k Bytes(one page)
-
 #endif //YAHFA_TAMPOLINE_H