Store all native JNI methods in data structures

.gitmodules

@@ -34,6 +34,9 @@
 [submodule "zlib"]
 	path = native/jni/external/zlib
 	url = https://android.googlesource.com/platform/external/zlib
+[submodule "parallel-hashmap"]
+	path = native/jni/external/parallel-hashmap
+	url = https://github.com/greg7mdp/parallel-hashmap.git
 [submodule "termux-elf-cleaner"]
 	path = tools/termux-elf-cleaner
 	url = https://github.com/termux/termux-elf-cleaner.git

native/jni/Android.mk

@@ -5,13 +5,13 @@ LOCAL_PATH := $(call my-dir)
 ########################
 
 # Global toggle for the WIP zygote injection features
-ENABLE_INJECT := 0
+ENABLE_INJECT := 1
 
 ifdef B_MAGISK
 
 include $(CLEAR_VARS)
 LOCAL_MODULE := magisk
-LOCAL_STATIC_LIBRARIES := libnanopb libsystemproperties libutils
+LOCAL_STATIC_LIBRARIES := libnanopb libsystemproperties libutils libphmap
 LOCAL_C_INCLUDES := jni/include
 
 LOCAL_SRC_FILES := \
@@ -42,7 +42,8 @@ LOCAL_STATIC_LIBRARIES += libxhook
 LOCAL_SRC_FILES += \
     inject/entry.cpp \
     inject/utils.cpp \
-    inject/hook.cpp
+    inject/hook.cpp \
+    inject/memory.cpp
 else
 LOCAL_SRC_FILES += magiskhide/proc_monitor.cpp
 endif
@@ -146,7 +147,7 @@ ifneq (,$(wildcard jni/test.cpp))
 
 include $(CLEAR_VARS)
 LOCAL_MODULE := test
-LOCAL_STATIC_LIBRARIES := libutils
+LOCAL_STATIC_LIBRARIES := libutils libphmap
 LOCAL_C_INCLUDES := jni/include
 LOCAL_SRC_FILES := test.cpp
 include $(BUILD_EXECUTABLE)

native/jni/external/Android.mk

@@ -1,5 +1,12 @@
 LOCAL_PATH := $(call my-dir)
 
+# Header only library
+include $(CLEAR_VARS)
+LOCAL_MODULE:= libphmap
+LOCAL_C_INCLUDES := $(LOCAL_PATH)/parallel-hashmap
+LOCAL_EXPORT_C_INCLUDES := $(LOCAL_C_INCLUDES)
+include $(BUILD_STATIC_LIBRARY)
+
 # libxz.a
 include $(CLEAR_VARS)
 LOCAL_MODULE:= libxz

native/jni/external/parallel-hashmap

@@ -0,0 +1 @@
+Subproject commit 7684faf186806e2c88554a78188c18185b21f127

native/jni/inject/entry.cpp

@@ -4,7 +4,6 @@
 #include <sys/sendfile.h>
 #include <sys/prctl.h>
 #include <android/log.h>
-#include <atomic>
 
 #include <utils.hpp>
 
@@ -25,28 +24,12 @@ static void inject_logging() {
     log_cb.ex = nop_ex;
 }
 
-__attribute__((destructor))
-static void inject_cleanup() {
-    if (active_threads < 0)
-        return;
-
-    // Setup 1ms
-    timespec ts = { .tv_sec = 0, .tv_nsec = 1000000L };
-
-    // Check flag in busy loop
-    while (active_threads)
-        nanosleep(&ts, nullptr);
-
-    // Wait another 1ms to make sure all threads left our code
-    nanosleep(&ts, nullptr);
-}
-
 void self_unload() {
     LOGD("hook: Request to self unload\n");
     // If unhook failed, do not unload or else it will cause SIGSEGV
     if (!unhook_functions())
         return;
-    new_daemon_thread(reinterpret_cast<void *(*)(void *)>(&dlclose), self_handle);
+    new_daemon_thread(reinterpret_cast<thread_entry>(&dlclose), self_handle);
     active_threads--;
 }
 
@@ -86,6 +69,22 @@ static void sanitize_environ() {
     }
 }
 
+__attribute__((destructor))
+static void inject_cleanup_wait() {
+    if (active_threads < 0)
+        return;
+
+    // Setup 1ms
+    timespec ts = { .tv_sec = 0, .tv_nsec = 1000000L };
+
+    // Check flag in busy loop
+    while (active_threads)
+        nanosleep(&ts, nullptr);
+
+    // Wait another 1ms to make sure all threads left our code
+    nanosleep(&ts, nullptr);
+}
+
 __attribute__((constructor))
 static void inject_init() {
     inject_logging();

native/jni/inject/hook.cpp

@@ -1,53 +1,54 @@
-#include <jni.h>
-
 #include <xhook.h>
 #include <utils.hpp>
 #include <flags.hpp>
 #include <daemon.hpp>
 
 #include "inject.hpp"
+#include "memory.hpp"
 
 using namespace std;
-
-#define DCL_HOOK_FUNC(ret, func, ...) \
-    static ret (*old_##func)(__VA_ARGS__); \
-    static ret new_##func(__VA_ARGS__)
-
-#define DCL_JNI_FUNC(name) \
-    static const JNINativeMethod *name##_orig = nullptr; \
-    extern const JNINativeMethod name##_methods[]; \
-    extern const int name##_methods_num;
-
-namespace {
+using jni_hook::hash_map;
+using jni_hook::tree_map;
+using xstring = jni_hook::string;
 
 struct HookContext {
     int pid;
     bool do_hide;
 };
 
-// JNI method declarations
-DCL_JNI_FUNC(nativeForkAndSpecialize)
-DCL_JNI_FUNC(nativeSpecializeAppProcess)
-DCL_JNI_FUNC(nativeForkSystemServer)
-
-}
-
-// For some reason static vectors won't work, use pointers instead
 static vector<tuple<const char *, const char *, void **>> *xhook_list;
-static vector<JNINativeMethod> *jni_list;
+static vector<JNINativeMethod> *jni_hook_list;
+static hash_map<xstring, tree_map<xstring, tree_map<xstring, void *>>> *jni_method_map;
 
 static JavaVM *g_jvm;
 static int prev_fork_pid = -1;
 static HookContext *current_ctx;
 
+#define DCL_HOOK_FUNC(ret, func, ...) \
+    static ret (*old_##func)(__VA_ARGS__); \
+    static ret new_##func(__VA_ARGS__)
+
+#define DCL_JNI_FUNC(name) \
+    static int name##_orig_idx; \
+    static inline JNINativeMethod &name##_orig() {       \
+        return (*jni_hook_list)[name##_orig_idx];        \
+    }                      \
+    extern const JNINativeMethod name##_methods[];       \
+    extern const int name##_methods_num;
+
+namespace {
+// JNI method declarations
+DCL_JNI_FUNC(nativeForkAndSpecialize)
+DCL_JNI_FUNC(nativeSpecializeAppProcess)
+DCL_JNI_FUNC(nativeForkSystemServer)
+}
+
 #define HOOK_JNI(method) \
-if (newMethods[i].name == #method##sv) { \
-    auto orig = new JNINativeMethod(); \
-    memcpy(orig, &newMethods[i], sizeof(JNINativeMethod)); \
-    method##_orig = orig; \
-    jni_list->push_back(newMethods[i]); \
+if (hooked < 3 && methods[i].name == #method##sv) { \
+    jni_hook_list->push_back(methods[i]);              \
+    method##_orig_idx = jni_hook_list->size() - 1;     \
     for (int j = 0; j < method##_methods_num; ++j) {   \
-        if (strcmp(newMethods[i].signature, method##_methods[j].signature) == 0) { \
+        if (strcmp(methods[i].signature, method##_methods[j].signature) == 0) { \
             newMethods[i] = method##_methods[j];       \
             LOGI("hook: replaced #" #method "\n");     \
             ++hooked;    \
@@ -61,23 +62,26 @@ DCL_HOOK_FUNC(int, jniRegisterNativeMethods,
         JNIEnv *env, const char *className, const JNINativeMethod *methods, int numMethods) {
     LOGD("hook: jniRegisterNativeMethods %s", className);
 
-    unique_ptr<JNINativeMethod[]> newMethods;
-    int hooked = 0;
-
     if (g_jvm == nullptr) {
         // Save for later unhooking
         env->GetJavaVM(&g_jvm);
     }
 
+    unique_ptr<JNINativeMethod[]> newMethods;
+    int hooked = numeric_limits<int>::max();
     if (className == "com/android/internal/os/Zygote"sv) {
+        hooked = 0;
         newMethods = make_unique<JNINativeMethod[]>(numMethods);
         memcpy(newMethods.get(), methods, sizeof(JNINativeMethod) * numMethods);
-        for (int i = 0; i < numMethods && hooked < 3; ++i) {
+    }
+
+    auto &class_map = (*jni_method_map)[className];
+    for (int i = 0; i < numMethods; ++i) {
+        class_map[methods[i].name][methods[i].signature] = methods[i].fnPtr;
         HOOK_JNI(nativeForkAndSpecialize);
         HOOK_JNI(nativeSpecializeAppProcess);
         HOOK_JNI(nativeForkSystemServer);
     }
-    }
 
     return old_jniRegisterNativeMethods(env, className, newMethods.get() ?: methods, numMethods);
 }
@@ -239,7 +243,8 @@ void hook_functions() {
     xhook_enable_sigsegv_protection(0);
 #endif
     xhook_list = new remove_pointer_t<decltype(xhook_list)>();
-    jni_list = new remove_pointer_t<decltype(jni_list)>();
+    jni_hook_list = new remove_pointer_t<decltype(jni_hook_list)>();
+    jni_method_map = new remove_pointer_t<decltype(jni_method_map)>();
 
     XHOOK_REGISTER(".*\\libandroid_runtime.so$", jniRegisterNativeMethods);
     XHOOK_REGISTER(".*\\libandroid_runtime.so$", fork);
@@ -252,14 +257,19 @@ bool unhook_functions() {
     if (g_jvm->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION_1_6) != JNI_OK)
         return false;
 
+    // Do NOT call any destructors
+    operator delete(jni_method_map);
+    // Directly unmap the whole memory block
+    jni_hook::memory_block::release();
+
     // Unhook JNI methods
-    if (!jni_list->empty() && old_jniRegisterNativeMethods(env,
+    if (!jni_hook_list->empty() && old_jniRegisterNativeMethods(env,
             "com/android/internal/os/Zygote",
-            jni_list->data(), jni_list->size()) != 0) {
+            jni_hook_list->data(), jni_hook_list->size()) != 0) {
         LOGE("hook: Failed to register JNI hook\n");
         return false;
     }
-    delete jni_list;
+    delete jni_hook_list;
 
     // Unhook xhook
     for (auto &[path, sym, old_func] : *xhook_list) {

native/jni/inject/jni_hooks.hpp

@@ -1,5 +1,5 @@
 /*
- * Original code: https://github.com/RikkaApps/Riru/blob/master/riru/src/main/cpp/jni_native_method.cpp
+ * Original code from: https://github.com/RikkaApps/Riru
  * The code is modified and sublicensed to GPLv3 for incorporating into Magisk.
  *
  * Copyright (c) 2018-2021, RikkaW
@@ -38,7 +38,7 @@ static ret name(__VA_ARGS__)
 
 #define orig_fork(ver, ...) \
     reinterpret_cast<decltype(&nativeForkAndSpecialize_##ver)> \
-    (nativeForkAndSpecialize_orig->fnPtr)(__VA_ARGS__)
+    (nativeForkAndSpecialize_orig().fnPtr)(__VA_ARGS__)
 
 #define post_fork() \
     nativeForkAndSpecialize_post(&ctx, env, clazz); \
@@ -202,7 +202,7 @@ DCL_FORK_AND_SPECIALIZE(samsung_p,
 
 #define orig_spec(ver, ...) \
     reinterpret_cast<decltype(&nativeSpecializeAppProcess_##ver)> \
-    (nativeSpecializeAppProcess_orig->fnPtr)(__VA_ARGS__)
+    (nativeSpecializeAppProcess_orig().fnPtr)(__VA_ARGS__)
 
 #define post_spec() \
     nativeSpecializeAppProcess_post(&ctx, env, clazz)
@@ -301,7 +301,7 @@ DCL_SPECIALIZE_APP(samsung_q,
 
 #define orig_server(ver, ...) \
     reinterpret_cast<decltype(&nativeForkSystemServer_##ver)> \
-    (nativeForkSystemServer_orig->fnPtr)(__VA_ARGS__)
+    (nativeForkSystemServer_orig().fnPtr)(__VA_ARGS__)
 
 #define post_server() \
     nativeForkSystemServer_post(&ctx, env, clazz); \
@@ -355,8 +355,7 @@ const JNINativeMethod nativeSpecializeAppProcess_methods[] = {
     DEF_SPEC(r_dp2), DEF_SPEC(r_dp3)
 #endif
 };
-const int nativeSpecializeAppProcess_methods_num = std::size(
-    nativeSpecializeAppProcess_methods);
+const int nativeSpecializeAppProcess_methods_num = std::size(nativeSpecializeAppProcess_methods);
 
 const JNINativeMethod nativeForkSystemServer_methods[] = {
     DEF_SERVER(m), DEF_SERVER(samsung_q)

native/jni/inject/memory.cpp

@@ -0,0 +1,30 @@
+#include "memory.hpp"
+
+namespace jni_hook {
+
+// We know our minimum alignment is WORD size (size of pointer)
+static constexpr size_t ALIGN = sizeof(long);
+
+static constexpr size_t CAPACITY = (1 << 24);
+
+// No need to be thread safe as the initial mmap always happens on the main thread
+static uint8_t *_area = nullptr;
+
+static std::atomic<uint8_t *> _curr = nullptr;
+
+void *memory_block::allocate(size_t sz) {
+    if (!_area) {
+        // Memory will not actually be allocated because physical pages are mapped in on-demand
+        _area = static_cast<uint8_t *>(xmmap(
+                nullptr, CAPACITY, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));
+        _curr = _area;
+    }
+    return _curr.fetch_add(do_align(sz, ALIGN));
+}
+
+void memory_block::release() {
+    if (_area)
+        munmap(_area, CAPACITY);
+}
+
+} // namespace jni_hook

native/jni/inject/memory.hpp

@@ -0,0 +1,40 @@
+#pragma once
+
+#include <map>
+#include <parallel_hashmap/phmap.h>
+
+#include <utils.hpp>
+
+namespace jni_hook {
+
+struct memory_block {
+    static void *allocate(size_t sz);
+    static void deallocate(void *, size_t) { /* Monotonic increase */ }
+    static void release();
+};
+
+template<class T>
+using allocator = stateless_allocator<T, memory_block>;
+
+using string = std::basic_string<char, std::char_traits<char>, allocator<char>>;
+
+// Use node_hash_map since it will use less memory because we are using a monotonic allocator
+template<class K, class V>
+using hash_map = phmap::node_hash_map<K, V,
+    phmap::priv::hash_default_hash<K>,
+    phmap::priv::hash_default_eq<K>,
+    allocator<std::pair<const K, V>>
+>;
+
+template<class K, class V>
+using tree_map = std::map<K, V,
+    std::less<K>,
+    allocator<std::pair<const K, V>>
+>;
+
+} // namespace jni_hook
+
+// Provide heterogeneous lookup for jni_hook::string
+namespace phmap::priv {
+template <> struct HashEq<jni_hook::string> : StringHashEqT<char> {};
+} // namespace phmap::priv

native/jni/utils/misc.hpp

@@ -58,6 +58,21 @@ reversed_container<T> reversed(T &base) {
     return reversed_container<T>(base);
 }
 
+template<typename T, typename Impl>
+class stateless_allocator {
+public:
+    using value_type = T;
+    T *allocate(size_t num) { return static_cast<T*>(Impl::allocate(sizeof(T) * num)); }
+    void deallocate(T *ptr, size_t num) { Impl::deallocate(ptr, sizeof(T) * num); }
+    stateless_allocator()                           = default;
+    stateless_allocator(const stateless_allocator&) = default;
+    stateless_allocator(stateless_allocator&&)      = default;
+    template <typename U>
+    stateless_allocator(const stateless_allocator<U, Impl>&) {}
+    bool operator==(const stateless_allocator&) { return true; }
+    bool operator!=(const stateless_allocator&) { return false; }
+};
+
 int parse_int(const char *s);
 static inline int parse_int(const std::string &s) { return parse_int(s.data()); }
 static inline int parse_int(std::string_view s) { return parse_int(s.data()); }