From cd76a9de4d22595c39af34f9a5036ed4285dc90b Mon Sep 17 00:00:00 2001
From: Ruikai Liu <lrk700@gmail.com>
Date: Mon, 21 Dec 2020 15:38:12 +0800
Subject: [PATCH] compute available size for trampoline in trampoline.c

---
 library/src/main/jni/HookMain.c   | 10 ------
 library/src/main/jni/trampoline.c | 56 +++++++++++++++++++------------
 library/src/main/jni/trampoline.h |  9 +----
 3 files changed, 36 insertions(+), 39 deletions(-)

diff --git a/library/src/main/jni/HookMain.c b/library/src/main/jni/HookMain.c
index 07744b0..aab4c75 100644
--- a/library/src/main/jni/HookMain.c
+++ b/library/src/main/jni/HookMain.c
@@ -119,15 +119,6 @@ static void setNonCompilable(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
@@ -177,7 +168,6 @@ static int replaceMethod(void *fromMethod, void *toMethod, int isBackup) {
         LOGI("change access flags from 0x%x to 0x%x", old_flags, access_flags);
     }
 
-    hookCount += 1;
     return 0;
 
 }
diff --git a/library/src/main/jni/trampoline.c b/library/src/main/jni/trampoline.c
index 4041af7..9164568 100644
--- a/library/src/main/jni/trampoline.c
+++ b/library/src/main/jni/trampoline.c
@@ -4,6 +4,7 @@
 #include <sys/mman.h>
 #include <string.h>
 #include <errno.h>
+#include <unistd.h>
 
 #include "common.h"
 #include "trampoline.h"
@@ -13,8 +14,10 @@
 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;
+static unsigned char *currentTrampolineOff = 0;
+static unsigned char *trampolineSpaceEnd = 0;
+
+static void *allocTrampolineSpace();
 
 // trampoline:
 // 1. set eax/rdi/r0/x0 to the hook ArtMethod addr
@@ -124,10 +127,21 @@ unsigned char trampolineForBackup[] = {
         0x89, 0x67, 0x45, 0x23
 };
 #endif
-static unsigned int trampolineSize = roundUpToPtrSize(MAX(sizeof(trampoline), sizeof(trampolineForBackup)));
 
 void *genTrampoline(void *toMethod, void *entrypoint) {
-    unsigned char *targetAddr = trampolineCode + trampolineSize * hookCount;
+    size_t trampolineSize = entrypoint != NULL ? sizeof(trampolineForBackup) : sizeof(trampoline);
+
+    // check available space for new trampoline
+    if(currentTrampolineOff+trampolineSize > trampolineSpaceEnd) {
+        currentTrampolineOff = allocTrampolineSpace();
+        if (currentTrampolineOff == NULL) {
+            return NULL;
+        } else {
+            trampolineSpaceEnd = currentTrampolineOff + TRAMPOLINE_SPACE_SIZE;
+        }
+    }
+
+    unsigned char *targetAddr = currentTrampolineOff;
 
     if(entrypoint != NULL) {
         memcpy(targetAddr, trampolineForBackup, sizeof(trampolineForBackup));
@@ -173,11 +187,18 @@ void *genTrampoline(void *toMethod, void *entrypoint) {
     else {
         memcpy(targetAddr + 16, &toMethod, pointer_size);
     }
-
 #else
 #error Unsupported architecture
 #endif
 
+    // skip 4 bytes of code_size_
+    if(entrypoint == NULL) {
+        targetAddr += 4;
+    }
+
+    // keep each trampoline aligned
+    currentTrampolineOff += roundUpToPtrSize(trampolineSize);
+
     return targetAddr;
 }
 
@@ -196,23 +217,16 @@ void setupTrampoline(uint8_t offset) {
 #endif
 }
 
-int doInitHookCap(unsigned int cap) {
-    if (cap == 0) {
-        LOGE("invalid capacity: %d", cap);
-        return 1;
-    }
-    if (hookCap) {
-        LOGW("allocating new space for trampoline code");
-    }
-    unsigned int allSize = trampolineSize * cap;
-    unsigned char *buf = mmap(NULL, allSize, PROT_READ | PROT_WRITE | PROT_EXEC,
+static void *allocTrampolineSpace() {
+    unsigned char *buf = mmap(NULL, TRAMPOLINE_SPACE_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;
-}
+    else {
+        LOGI("allocating space for trampoline code at %p", buf);
+        return buf;
+    }
+
+}
\ No newline at end of file
diff --git a/library/src/main/jni/trampoline.h b/library/src/main/jni/trampoline.h
index 267017f..1d2ce3e 100644
--- a/library/src/main/jni/trampoline.h
+++ b/library/src/main/jni/trampoline.h
@@ -7,17 +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 setupTrampoline(uint8_t offset);
 
 void *genTrampoline(void *toMethod, void *entrypoint);
 
-#define DEFAULT_CAP 100 //size of each trampoline area would be no more than 4k Bytes(one page)
+#define TRAMPOLINE_SPACE_SIZE 4096 // 4k mem page size
 
 #endif //YAHFA_TAMPOLINE_H