refactor jni hooks

.gitmodules

@@ -1,6 +1,3 @@
 [submodule "loader/src/external/lsplt"]
 	path = loader/src/external/lsplt
 	url = https://github.com/LSPosed/lsplt
-[submodule "loader/src/external/parallel-hashmap"]
-	path = loader/src/external/parallel-hashmap
-	url = https://github.com/greg7mdp/parallel-hashmap

loader/src/Android.mk

@@ -17,7 +17,7 @@ LOCAL_MODULE           := zygisk
 LOCAL_C_INCLUDES       := $(LOCAL_PATH)/include
 FILE_LIST              := $(filter %.cpp, $(call walk, $(LOCAL_PATH)/injector))
 LOCAL_SRC_FILES        := $(FILE_LIST:COMMON_FILE_LIST:$(LOCAL_PATH)/%=%)
-LOCAL_STATIC_LIBRARIES := cxx common liblsplt libphmap
+LOCAL_STATIC_LIBRARIES := cxx common liblsplt
 LOCAL_LDLIBS           := -llog
 include $(BUILD_SHARED_LIBRARY)
 

loader/src/external/Android.mk

@@ -12,10 +12,3 @@ LOCAL_SRC_FILES := \
     lsplt/lsplt/src/main/jni/elf_util.cc \
     lsplt/lsplt/src/main/jni/lsplt.cc
 include $(BUILD_STATIC_LIBRARY)
-
-# Header only library
-include $(CLEAR_VARS)
-LOCAL_MODULE:= libphmap
-LOCAL_CFLAGS := -Wno-unused-value
-LOCAL_EXPORT_C_INCLUDES := $(LOCAL_PATH)/parallel-hashmap
-include $(BUILD_STATIC_LIBRARY)

loader/src/injector/art_method.hpp

@@ -0,0 +1,82 @@
+#pragma once
+
+#include "jni_helper.hpp"
+
+template <typename T>
+constexpr inline auto RoundUpTo(T v, size_t size) {
+    return v + size - 1 - ((v + size - 1) & (size - 1));
+}
+
+inline static constexpr auto kPointerSize = sizeof(void *);
+
+namespace lsplant::art {
+
+    class ArtMethod {
+
+    public:
+        void *GetData() {
+            return *reinterpret_cast<void **>(reinterpret_cast<uintptr_t>(this) + data_offset);
+        }
+
+        static art::ArtMethod *FromReflectedMethod(JNIEnv *env, jobject method) {
+            if (art_method_field) [[likely]] {
+                return reinterpret_cast<art::ArtMethod *>(
+                        JNI_GetLongField(env, method, art_method_field));
+            } else {
+                return reinterpret_cast<art::ArtMethod *>(env->FromReflectedMethod(method));
+            }
+        }
+
+        static bool Init(JNIEnv *env) {
+            ScopedLocalRef<jclass> executable{env, nullptr};
+            executable = JNI_FindClass(env, "java/lang/reflect/Executable");
+            if (!executable) {
+                LOGE("Failed to found Executable");
+                return false;
+            }
+
+            if (art_method_field = JNI_GetFieldID(env, executable, "artMethod", "J");
+                    !art_method_field) {
+                LOGE("Failed to find artMethod field");
+                return false;
+            }
+
+            auto throwable = JNI_FindClass(env, "java/lang/Throwable");
+            if (!throwable) {
+                LOGE("Failed to found Executable");
+                return false;
+            }
+            auto clazz = JNI_FindClass(env, "java/lang/Class");
+            static_assert(std::is_same_v<decltype(clazz)::BaseType, jclass>);
+            jmethodID get_declared_constructors = JNI_GetMethodID(env, clazz, "getDeclaredConstructors",
+                                                                  "()[Ljava/lang/reflect/Constructor;");
+            const auto constructors =
+                    JNI_Cast<jobjectArray>(JNI_CallObjectMethod(env, throwable, get_declared_constructors));
+            if (constructors.size() < 2) {
+                LOGE("Throwable has less than 2 constructors");
+                return false;
+            }
+            auto &first_ctor = constructors[0];
+            auto &second_ctor = constructors[1];
+            auto *first = FromReflectedMethod(env, first_ctor.get());
+            auto *second = FromReflectedMethod(env, second_ctor.get());
+            art_method_size = reinterpret_cast<uintptr_t>(second) - reinterpret_cast<uintptr_t>(first);
+            LOGD("ArtMethod size: %zu", art_method_size);
+            if (RoundUpTo(4 * 9, kPointerSize) + kPointerSize * 3 < art_method_size) [[unlikely]] {
+                LOGW("ArtMethod size exceeds maximum assume. There may be something wrong.");
+            }
+            entry_point_offset = art_method_size - kPointerSize;
+            data_offset = entry_point_offset - kPointerSize;
+            LOGD("ArtMethod::entrypoint offset: %zu", entry_point_offset);
+            LOGD("ArtMethod::data offset: %zu", data_offset);
+            return true;
+        }
+
+    private:
+        inline static jfieldID art_method_field = nullptr;
+        inline static size_t art_method_size = 0;
+        inline static size_t entry_point_offset = 0;
+        inline static size_t data_offset = 0;
+    };
+
+}  // namespace lsplant::art

loader/src/injector/gen_jni_hooks.py

@@ -0,0 +1,270 @@
+#!/usr/bin/env python3
+
+primitives = ['jint', 'jboolean', 'jlong']
+
+class JType:
+    def __init__(self, cpp, jni):
+        self.cpp = cpp
+        self.jni = jni
+
+
+class JArray(JType):
+    def __init__(self, type):
+        if type.cpp in primitives:
+            name = type.cpp + 'Array'
+        else:
+            name = 'jobjectArray'
+        super().__init__(name, '[' + type.jni)
+
+
+class Argument:
+    def __init__(self, name, type, set_arg = False):
+        self.name = name
+        self.type = type
+        self.set_arg = set_arg
+
+    def cpp(self):
+        return f'{self.type.cpp} {self.name}'
+
+# Args we don't care, give it an auto generated name
+class Anon(Argument):
+    cnt = 0
+    def __init__(self, type):
+        super().__init__(f'_{Anon.cnt}', type)
+        Anon.cnt += 1
+
+class Return:
+    def __init__(self, value, type):
+        self.value = value
+        self.type = type
+
+class Method:
+    def __init__(self, name, ret, args):
+        self.name = name
+        self.ret = ret
+        self.args = args
+
+    def cpp(self):
+        return ', '.join(map(lambda x: x.cpp(), self.args))
+
+    def name_list(self):
+        return ', '.join(map(lambda x: x.name, self.args))
+
+    def jni(self):
+        args = ''.join(map(lambda x: x.type.jni, self.args))
+        return f'({args}){self.ret.type.jni}'
+
+    def body(self):
+        return ''
+
+class JNIHook(Method):
+    def __init__(self, ver, ret, args):
+        name = f'{self.base_name()}_{ver}'
+        super().__init__(name, ret, args)
+
+    def base_name(self):
+        return ''
+
+    def orig_method(self):
+        return f'reinterpret_cast<decltype(&{self.name})>({self.base_name()}_orig)'
+
+def ind(i):
+    return '\n' + '    ' * i
+
+# Common types
+jint = JType('jint', 'I')
+jintArray = JArray(jint)
+jstring = JType('jstring', 'Ljava/lang/String;')
+jboolean = JType('jboolean', 'Z')
+jlong = JType('jlong', 'J')
+void = JType('void', 'V')
+
+class ForkAndSpec(JNIHook):
+    def __init__(self, ver, args):
+        super().__init__(ver, Return('ctx.pid', jint), args)
+
+    def base_name(self):
+        return 'nativeForkAndSpecialize'
+
+    def init_args(self):
+        return 'AppSpecializeArgs_v3 args(uid, gid, gids, runtime_flags, rlimits, mount_external, se_info, nice_name, instruction_set, app_data_dir);'
+
+    def body(self):
+        decl = ''
+        decl += ind(1) + self.init_args()
+        for a in self.args:
+            if a.set_arg:
+                decl += ind(1) + f'args.{a.name} = &{a.name};'
+        decl += ind(1) + 'HookContext ctx(env, &args);'
+        decl += ind(1) + f'ctx.{self.base_name()}_pre();'
+        decl += ind(1) + self.orig_method() + '('
+        decl += ind(2) + f'env, clazz, {self.name_list()}'
+        decl += ind(1) + ');'
+        decl += ind(1) + f'ctx.{self.base_name()}_post();'
+        return decl
+
+class SpecApp(ForkAndSpec):
+    def __init__(self, ver, args):
+        super().__init__(ver, args)
+        self.ret = Return('', void)
+
+    def base_name(self):
+        return 'nativeSpecializeAppProcess'
+
+class ForkServer(ForkAndSpec):
+    def base_name(self):
+        return 'nativeForkSystemServer'
+
+    def init_args(self):
+        return 'ServerSpecializeArgs_v1 args(uid, gid, gids, runtime_flags, permitted_capabilities, effective_capabilities);'
+
+# Common args
+uid = Argument('uid', jint)
+gid = Argument('gid', jint)
+gids = Argument('gids', jintArray)
+runtime_flags = Argument('runtime_flags', jint)
+rlimits = Argument('rlimits', JArray(jintArray))
+mount_external = Argument('mount_external', jint)
+se_info = Argument('se_info', jstring)
+nice_name = Argument('nice_name', jstring)
+fds_to_close = Argument('fds_to_close', jintArray)
+instruction_set = Argument('instruction_set', jstring)
+app_data_dir = Argument('app_data_dir', jstring)
+
+# o
+fds_to_ignore = Argument('fds_to_ignore', jintArray, True)
+
+# p
+is_child_zygote = Argument('is_child_zygote', jboolean, True)
+
+# q_alt
+is_top_app = Argument('is_top_app', jboolean, True)
+
+# r
+pkg_data_info_list = Argument('pkg_data_info_list', JArray(jstring), True)
+whitelisted_data_info_list = Argument('whitelisted_data_info_list', JArray(jstring), True)
+mount_data_dirs = Argument('mount_data_dirs', jboolean, True)
+mount_storage_dirs = Argument('mount_storage_dirs', jboolean, True)
+
+# server
+permitted_capabilities = Argument('permitted_capabilities', jlong)
+effective_capabilities = Argument('effective_capabilities', jlong)
+
+# Method definitions
+fas_l = ForkAndSpec('l', [uid, gid, gids, runtime_flags, rlimits, mount_external,
+    se_info, nice_name, fds_to_close, instruction_set, app_data_dir])
+
+fas_o = ForkAndSpec('o', [uid, gid, gids, runtime_flags, rlimits, mount_external,
+    se_info, nice_name, fds_to_close, fds_to_ignore, instruction_set, app_data_dir])
+
+fas_p = ForkAndSpec('p', [uid, gid, gids, runtime_flags, rlimits, mount_external, se_info,
+    nice_name, fds_to_close, fds_to_ignore, is_child_zygote, instruction_set, app_data_dir])
+
+fas_q_alt = ForkAndSpec('q_alt', [uid, gid, gids, runtime_flags, rlimits, mount_external, se_info,
+    nice_name, fds_to_close, fds_to_ignore, is_child_zygote, instruction_set, app_data_dir, is_top_app])
+
+fas_r = ForkAndSpec('r', [uid, gid, gids, runtime_flags, rlimits, mount_external, se_info,
+    nice_name, fds_to_close, fds_to_ignore, is_child_zygote, instruction_set, app_data_dir, is_top_app,
+    pkg_data_info_list, whitelisted_data_info_list, mount_data_dirs, mount_storage_dirs])
+
+fas_samsung_m = ForkAndSpec('samsung_m', [uid, gid, gids, runtime_flags, rlimits, mount_external,
+    se_info, Anon(jint), Anon(jint), nice_name, fds_to_close, instruction_set, app_data_dir])
+
+fas_samsung_n = ForkAndSpec('samsung_n', [uid, gid, gids, runtime_flags, rlimits, mount_external,
+    se_info, Anon(jint), Anon(jint), nice_name, fds_to_close, instruction_set, app_data_dir, Anon(jint)])
+
+fas_samsung_o = ForkAndSpec('samsung_o', [uid, gid, gids, runtime_flags, rlimits, mount_external,
+    se_info, Anon(jint), Anon(jint), nice_name, fds_to_close, fds_to_ignore, instruction_set, app_data_dir])
+
+fas_samsung_p = ForkAndSpec('samsung_p', [uid, gid, gids, runtime_flags, rlimits, mount_external,
+    se_info, Anon(jint), Anon(jint), nice_name, fds_to_close, fds_to_ignore, is_child_zygote,
+    instruction_set, app_data_dir])
+
+spec_q = SpecApp('q', [uid, gid, gids, runtime_flags, rlimits, mount_external, se_info,
+    nice_name, is_child_zygote, instruction_set, app_data_dir])
+
+spec_q_alt = SpecApp('q_alt', [uid, gid, gids, runtime_flags, rlimits, mount_external, se_info,
+    nice_name, is_child_zygote, instruction_set, app_data_dir, is_top_app])
+
+spec_r = SpecApp('r', [uid, gid, gids, runtime_flags, rlimits, mount_external, se_info, nice_name,
+    is_child_zygote, instruction_set, app_data_dir, is_top_app, pkg_data_info_list,
+    whitelisted_data_info_list, mount_data_dirs, mount_storage_dirs])
+
+spec_samsung_q = SpecApp('samsung_q', [uid, gid, gids, runtime_flags, rlimits, mount_external,
+    se_info, Anon(jint), Anon(jint), nice_name, is_child_zygote, instruction_set, app_data_dir])
+
+server_l = ForkServer('l', [uid, gid, gids, runtime_flags, rlimits,
+    permitted_capabilities, effective_capabilities])
+
+server_samsung_q = ForkServer('samsung_q', [uid, gid, gids, runtime_flags, Anon(jint), Anon(jint), rlimits,
+    permitted_capabilities, effective_capabilities])
+
+hook_map = {}
+
+def gen_jni_def(clz, methods):
+    if clz not in hook_map:
+        hook_map[clz] = []
+
+    decl = ''
+    for m in methods:
+        decl += ind(0) + f'[[clang::no_stack_protector]] {m.ret.type.cpp} {m.name}(JNIEnv *env, jclass clazz, {m.cpp()}) {{'
+        decl += m.body()
+        if m.ret.value:
+            decl += ind(1) + f'return {m.ret.value};'
+        decl += ind(0) + '}'
+
+    decl += ind(0) + f'std::array {m.base_name()}_methods {{'
+    for m in methods:
+        decl += ind(1) + 'JNINativeMethod {'
+        decl += ind(2) + f'"{m.base_name()}",'
+        decl += ind(2) + f'"{m.jni()}",'
+        decl += ind(2) + f'(void *) &{m.name}'
+        decl += ind(1) + '},'
+    decl += ind(0) + '};'
+    decl = ind(0) + f'void *{m.base_name()}_orig = nullptr;' + decl
+    decl += ind(0)
+
+    hook_map[clz].append(m.base_name())
+
+    return decl
+
+def gen_jni_hook():
+    decl = ''
+    decl += ind(0) + 'static void do_hook_zygote(JNIEnv *env) {'
+    decl += ind(1) + 'vector<JNINativeMethod> hooks;'
+    decl += ind(1) + 'const char *clz;'
+    for clz, methods in hook_map.items():
+        decl += ind(1) + f'clz = "{clz}";'
+        for m in methods:
+            decl += ind(1) + f'hookJniNativeMethods(env, clz, {m}_methods.data(), {m}_methods.size());'
+            decl += ind(1) + f'for (auto &method : {m}_methods) {{'
+            decl += ind(2) + f'if (method.fnPtr) {{'
+            decl += ind(3) + f'{m}_orig = method.fnPtr;'
+            decl += ind(3) + f'hooks.emplace_back(method);'
+            decl += ind(3) + f'break;'
+            decl += ind(2) + f'}}'
+            decl += ind(1) + f'}}'
+        decl += ind(1) + f'jni_hook_list->emplace(clz, std::move(hooks));'
+
+    decl += ind(0) + '}'
+    return decl
+
+with open('jni_hooks.hpp', 'w') as f:
+    f.write('// Generated by gen_jni_hooks.py\n')
+    f.write('\nnamespace {\n')
+
+    zygote = 'com/android/internal/os/Zygote'
+
+    methods = [fas_l, fas_o, fas_p, fas_q_alt, fas_r, fas_samsung_m, fas_samsung_n, fas_samsung_o, fas_samsung_p]
+    f.write(gen_jni_def(zygote, methods))
+
+    methods = [spec_q, spec_q_alt, spec_r, spec_samsung_q]
+    f.write(gen_jni_def(zygote, methods))
+
+    methods = [server_l, server_samsung_q]
+    f.write(gen_jni_def(zygote, methods))
+
+    f.write('\n} // namespace\n')
+
+    f.write(gen_jni_hook())
+    f.write('\n')

loader/src/injector/hook.cpp

@@ -4,29 +4,30 @@
 #include <regex.h>
 #include <bitset>
 #include <list>
+#include <map>
+#include <array>
 
 #include <lsplt.hpp>
 
 #include <fcntl.h>
 #include <sys/prctl.h>
 #include <sys/stat.h>
+#include <sys/mman.h>
 #include <unistd.h>
 
 #include "dl.h"
 #include "daemon.h"
 #include "zygisk.hpp"
-#include "memory.hpp"
 #include "module.hpp"
 #include "files.hpp"
+#include "misc.hpp"
+
+#include "art_method.hpp"
 
 using namespace std;
-using jni_hook::hash_map;
-using jni_hook::tree_map;
-using xstring = jni_hook::string;
 
 static void hook_unloader();
 static void unhook_functions();
-static void restore_jni_env(JNIEnv *env);
 
 namespace {
 
@@ -88,11 +89,6 @@ struct HookContext {
     HookContext(JNIEnv *env, void *args) :
     env(env), args{args}, process(nullptr), pid(-1), info_flags(0),
     hook_info_lock(PTHREAD_MUTEX_INITIALIZER) {
-        static bool restored_env = false;
-        if (!restored_env) {
-            restore_jni_env(env);
-            restored_env = true;
-        }
         g_ctx = this;
     }
     ~HookContext();
@@ -123,98 +119,12 @@ struct HookContext {
 // Global variables
 vector<tuple<dev_t, ino_t, const char *, void **>> *plt_hook_list;
 map<string, vector<JNINativeMethod>, StringCmp> *jni_hook_list;
-hash_map<xstring, tree_map<xstring, tree_map<xstring, void *>>> *jni_method_map;
 bool should_unmap_zygisk = false;
 
-const JNINativeInterface *old_functions = nullptr;
-JNINativeInterface *new_functions = nullptr;
-
 } // namespace
 
-#define HOOK_JNI(method)                                                                     \
-if (methods[i].name == #method##sv) {                                                        \
-    int j = 0;                                                                               \
-    for (; j < method##_methods_num; ++j) {                                                  \
-        if (strcmp(methods[i].signature, method##_methods[j].signature) == 0) {              \
-            jni_hook_list->try_emplace(className).first->second.push_back(methods[i]);       \
-            method##_orig = methods[i].fnPtr;                                                \
-            newMethods[i] = method##_methods[j];                                             \
-            LOGI("replaced %s#" #method "\n", className);                                    \
-            --hook_cnt;                                                                      \
-            break;                                                                           \
-        }                                                                                    \
-    }                                                                                        \
-    if (j == method##_methods_num) {                                                         \
-        LOGE("unknown signature of %s#" #method ": %s\n", className, methods[i].signature);  \
-    }                                                                                        \
-    continue;                                                                                \
-}
-
-// JNI method hook definitions, auto generated
-#include "jni_hooks.hpp"
-
-#undef HOOK_JNI
-
 namespace {
 
-string get_class_name(JNIEnv *env, jclass clazz) {
-    static auto class_getName = env->GetMethodID(env->FindClass("java/lang/Class"), "getName", "()Ljava/lang/String;");
-    auto nameRef = (jstring) env->CallObjectMethod(clazz, class_getName);
-    const char *name = env->GetStringUTFChars(nameRef, nullptr);
-    string className(name);
-    env->ReleaseStringUTFChars(nameRef, name);
-    std::replace(className.begin(), className.end(), '.', '/');
-    return className;
-}
-
-jint env_RegisterNatives(
-        JNIEnv *env, jclass clazz, const JNINativeMethod *methods, jint numMethods) {
-    auto className = get_class_name(env, clazz);
-    LOGV("JNIEnv->RegisterNatives [%s]\n", className.data());
-    auto newMethods = hookAndSaveJNIMethods(className.data(), methods, numMethods);
-    return old_functions->RegisterNatives(env, clazz, newMethods.get() ?: methods, numMethods);
-}
-
-void replace_jni_methods() {
-    auto get_created_java_vms = reinterpret_cast<jint (*)(JavaVM **, jsize, jsize *)>(
-            dlsym(RTLD_DEFAULT, "JNI_GetCreatedJavaVMs"));
-    if (!get_created_java_vms) {
-        for (auto &map: lsplt::MapInfo::Scan()) {
-            if (!map.path.ends_with("/libnativehelper.so")) continue;
-            void *h = dlopen(map.path.data(), RTLD_LAZY);
-            if (!h) {
-                LOGW("cannot dlopen libnativehelper.so: %s\n", dlerror());
-                break;
-            }
-            get_created_java_vms = reinterpret_cast<decltype(get_created_java_vms)>(dlsym(h, "JNI_GetCreatedJavaVMs"));
-            dlclose(h);
-            break;
-        }
-        if (!get_created_java_vms) {
-            LOGW("JNI_GetCreatedJavaVMs not found\n");
-            return;
-        }
-    }
-    JavaVM *vm = nullptr;
-    jsize num = 0;
-    jint res = get_created_java_vms(&vm, 1, &num);
-    if (res != JNI_OK || vm == nullptr) return;
-    JNIEnv *env = nullptr;
-    res = vm->GetEnv(reinterpret_cast<void **>(&env), JNI_VERSION_1_6);
-    if (res != JNI_OK || env == nullptr) return;
-    default_new(new_functions);
-    memcpy(new_functions, env->functions, sizeof(*new_functions));
-    new_functions->RegisterNatives = &env_RegisterNatives;
-
-    // Replace the function table in JNIEnv to hook RegisterNatives
-    old_functions = env->functions;
-    env->functions = new_functions;
-
-    // Re-run register_com_android_internal_os_Zygote to hook JNI methods
-    auto register_zygote = dlsym(RTLD_DEFAULT, "_ZN7android39register_com_android_internal_os_ZygoteEP7_JNIEnv");
-    reinterpret_cast<void (*)(JNIEnv *)>(register_zygote)(env);
-}
-
 #define DCL_HOOK_FUNC(ret, func, ...) \
 ret (*old_##func)(__VA_ARGS__);       \
 ret new_##func(__VA_ARGS__)
@@ -288,40 +198,103 @@ DCL_HOOK_FUNC(int, pthread_attr_destroy, void *target) {
 
 // -----------------------------------------------------------------
 
+static bool can_hook_jni = false;
+static jint MODIFIER_NATIVE = 0;
+static jmethodID member_getModifiers = nullptr;
+
 void hookJniNativeMethods(JNIEnv *env, const char *clz, JNINativeMethod *methods, int numMethods) {
-    auto class_map = jni_method_map->find(clz);
-    if (class_map == jni_method_map->end()) {
-        for (int i = 0; i < numMethods; ++i) {
+    if (!can_hook_jni) return;
+    auto clazz = env->FindClass(clz);
+    if (clazz == nullptr) {
+        env->ExceptionClear();
+        for (int i = 0; i < numMethods; i++) {
             methods[i].fnPtr = nullptr;
         }
         return;
     }
 
     vector<JNINativeMethod> hooks;
-    for (int i = 0; i < numMethods; ++i) {
-        auto method_map = class_map->second.find(methods[i].name);
-        if (method_map != class_map->second.end()) {
-            auto it = method_map->second.find(methods[i].signature);
-            if (it != method_map->second.end()) {
-                // Copy the JNINativeMethod
-                hooks.push_back(methods[i]);
-                // Save the original function pointer
-                methods[i].fnPtr = it->second;
-                // Do not allow double hook, remove method from map
-                method_map->second.erase(it);
-                continue;
-            }
+    for (int i = 0; i < numMethods; i++) {
+        auto &nm = methods[i];
+        auto mid = env->GetMethodID(clazz, nm.name, nm.signature);
+        bool is_static = false;
+        if (mid == nullptr) {
+            env->ExceptionClear();
+            mid = env->GetStaticMethodID(clazz, nm.name, nm.signature);
+            is_static = true;
         }
-        // No matching method found, set fnPtr to null
-        methods[i].fnPtr = nullptr;
+        if (mid == nullptr) {
+            env->ExceptionClear();
+            nm.fnPtr = nullptr;
+            continue;
+        }
+        auto method = lsplant::JNI_ToReflectedMethod(env, clazz, mid, is_static);
+        auto modifier = lsplant::JNI_CallIntMethod(env, method, member_getModifiers);
+        if ((modifier & MODIFIER_NATIVE) == 0) {
+            nm.fnPtr = nullptr;
+            continue;
+        }
+        auto artMethod = lsplant::art::ArtMethod::FromReflectedMethod(env, method);
+        hooks.push_back(nm);
+        auto orig = artMethod->GetData();
+        LOGD("replaced %s %s orig %p", clz, nm.name, orig);
+        nm.fnPtr = orig;
     }
 
-    if (hooks.empty())
-        return;
-
-    old_functions->RegisterNatives(env, env->FindClass(clz), hooks.data(), static_cast<int>(hooks.size()));
+    if (hooks.empty()) return;
+    env->RegisterNatives(clazz, hooks.data(), hooks.size());
 }
 
+// JNI method hook definitions, auto generated
+#include "jni_hooks.hpp"
+
+void initialize_jni_hook() {
+    auto get_created_java_vms = reinterpret_cast<jint (*)(JavaVM **, jsize, jsize *)>(
+            dlsym(RTLD_DEFAULT, "JNI_GetCreatedJavaVMs"));
+    if (!get_created_java_vms) {
+        for (auto &map: lsplt::MapInfo::Scan()) {
+            if (!map.path.ends_with("/libnativehelper.so")) continue;
+            void *h = dlopen(map.path.data(), RTLD_LAZY);
+            if (!h) {
+                LOGW("cannot dlopen libnativehelper.so: %s\n", dlerror());
+                break;
+            }
+            get_created_java_vms = reinterpret_cast<decltype(get_created_java_vms)>(dlsym(h, "JNI_GetCreatedJavaVMs"));
+            dlclose(h);
+            break;
+        }
+        if (!get_created_java_vms) {
+            LOGW("JNI_GetCreatedJavaVMs not found\n");
+            return;
+        }
+    }
+    JavaVM *vm = nullptr;
+    jsize num = 0;
+    jint res = get_created_java_vms(&vm, 1, &num);
+    if (res != JNI_OK || vm == nullptr) return;
+    JNIEnv *env = nullptr;
+    res = vm->GetEnv(reinterpret_cast<void **>(&env), JNI_VERSION_1_6);
+    if (res != JNI_OK || env == nullptr) return;
+
+    auto classMember = lsplant::JNI_FindClass(env, "java/lang/reflect/Member");
+    if (classMember != nullptr) member_getModifiers = lsplant::JNI_GetMethodID(env, classMember, "getModifiers", "()I");
+    auto classModifier = lsplant::JNI_FindClass(env, "java/lang/reflect/Modifier");
+    if (classModifier != nullptr) {
+        auto fieldId = lsplant::JNI_GetStaticFieldID(env, classModifier, "NATIVE", "I");
+        if (fieldId != nullptr) MODIFIER_NATIVE = lsplant::JNI_GetStaticIntField(env, classModifier, fieldId);
+    }
+    if (member_getModifiers == nullptr || MODIFIER_NATIVE == 0) return;
+    if (!lsplant::art::ArtMethod::Init(env)) {
+        LOGE("failed to init ArtMethod");
+        return;
+    }
+
+    can_hook_jni = true;
+    do_hook_zygote(env);
+}
+
+// -----------------------------------------------------------------
+
 ZygiskModule::ZygiskModule(int id, void *handle, void *entry)
 : id(id), handle(handle), entry{entry}, api{}, mod{nullptr} {
     // Make sure all pointers are null
@@ -716,12 +689,6 @@ HookContext::~HookContext() {
 
 } // namespace
 
-static void restore_jni_env(JNIEnv *env) {
-    env->functions = old_functions;
-    delete new_functions;
-    new_functions = nullptr;
-}
-
 static bool hook_commit() {
     if (lsplt::CommitHook()) {
         return true;
@@ -748,7 +715,6 @@ static void hook_register(dev_t dev, ino_t inode, const char *symbol, void *new_
 void hook_functions() {
     default_new(plt_hook_list);
     default_new(jni_hook_list);
-    default_new(jni_method_map);
 
     ino_t android_runtime_inode = 0;
     dev_t android_runtime_dev = 0;
@@ -771,7 +737,7 @@ void hook_functions() {
                            [](auto &t) { return *std::get<3>(t) == nullptr;}),
             plt_hook_list->end());
 
-    replace_jni_methods();
+    initialize_jni_hook();
 }
 
 static void hook_unloader() {

loader/src/injector/jni_hooks.hpp

@@ -109,54 +109,53 @@ void *nativeForkAndSpecialize_orig = nullptr;
     ctx.nativeForkAndSpecialize_post();
     return ctx.pid;
 }
-const JNINativeMethod nativeForkAndSpecialize_methods[] = {
-    {
+std::array nativeForkAndSpecialize_methods {
+    JNINativeMethod {
         "nativeForkAndSpecialize",
         "(II[II[[IILjava/lang/String;Ljava/lang/String;[ILjava/lang/String;Ljava/lang/String;)I",
         (void *) &nativeForkAndSpecialize_l
     },
-    {
+    JNINativeMethod {
         "nativeForkAndSpecialize",
         "(II[II[[IILjava/lang/String;Ljava/lang/String;[I[ILjava/lang/String;Ljava/lang/String;)I",
         (void *) &nativeForkAndSpecialize_o
     },
-    {
+    JNINativeMethod {
         "nativeForkAndSpecialize",
         "(II[II[[IILjava/lang/String;Ljava/lang/String;[I[IZLjava/lang/String;Ljava/lang/String;)I",
         (void *) &nativeForkAndSpecialize_p
     },
-    {
+    JNINativeMethod {
         "nativeForkAndSpecialize",
         "(II[II[[IILjava/lang/String;Ljava/lang/String;[I[IZLjava/lang/String;Ljava/lang/String;Z)I",
         (void *) &nativeForkAndSpecialize_q_alt
     },
-    {
+    JNINativeMethod {
         "nativeForkAndSpecialize",
         "(II[II[[IILjava/lang/String;Ljava/lang/String;[I[IZLjava/lang/String;Ljava/lang/String;Z[Ljava/lang/String;[Ljava/lang/String;ZZ)I",
         (void *) &nativeForkAndSpecialize_r
     },
-    {
+    JNINativeMethod {
         "nativeForkAndSpecialize",
         "(II[II[[IILjava/lang/String;IILjava/lang/String;[ILjava/lang/String;Ljava/lang/String;)I",
         (void *) &nativeForkAndSpecialize_samsung_m
     },
-    {
+    JNINativeMethod {
         "nativeForkAndSpecialize",
         "(II[II[[IILjava/lang/String;IILjava/lang/String;[ILjava/lang/String;Ljava/lang/String;I)I",
         (void *) &nativeForkAndSpecialize_samsung_n
     },
-    {
+    JNINativeMethod {
         "nativeForkAndSpecialize",
         "(II[II[[IILjava/lang/String;IILjava/lang/String;[I[ILjava/lang/String;Ljava/lang/String;)I",
         (void *) &nativeForkAndSpecialize_samsung_o
     },
-    {
+    JNINativeMethod {
         "nativeForkAndSpecialize",
         "(II[II[[IILjava/lang/String;IILjava/lang/String;[I[IZLjava/lang/String;Ljava/lang/String;)I",
         (void *) &nativeForkAndSpecialize_samsung_p
     },
 };
-constexpr int nativeForkAndSpecialize_methods_num = std::size(nativeForkAndSpecialize_methods);
 
 void *nativeSpecializeAppProcess_orig = nullptr;
 [[clang::no_stack_protector]] void nativeSpecializeAppProcess_q(JNIEnv *env, jclass clazz, jint uid, jint gid, jintArray gids, jint runtime_flags, jobjectArray rlimits, jint mount_external, jstring se_info, jstring nice_name, jboolean is_child_zygote, jstring instruction_set, jstring app_data_dir) {
@@ -205,29 +204,28 @@ void *nativeSpecializeAppProcess_orig = nullptr;
     );
     ctx.nativeSpecializeAppProcess_post();
 }
-const JNINativeMethod nativeSpecializeAppProcess_methods[] = {
-    {
+std::array nativeSpecializeAppProcess_methods {
+    JNINativeMethod {
         "nativeSpecializeAppProcess",
         "(II[II[[IILjava/lang/String;Ljava/lang/String;ZLjava/lang/String;Ljava/lang/String;)V",
         (void *) &nativeSpecializeAppProcess_q
     },
-    {
+    JNINativeMethod {
         "nativeSpecializeAppProcess",
         "(II[II[[IILjava/lang/String;Ljava/lang/String;ZLjava/lang/String;Ljava/lang/String;Z)V",
         (void *) &nativeSpecializeAppProcess_q_alt
     },
-    {
+    JNINativeMethod {
         "nativeSpecializeAppProcess",
         "(II[II[[IILjava/lang/String;Ljava/lang/String;ZLjava/lang/String;Ljava/lang/String;Z[Ljava/lang/String;[Ljava/lang/String;ZZ)V",
         (void *) &nativeSpecializeAppProcess_r
     },
-    {
+    JNINativeMethod {
         "nativeSpecializeAppProcess",
         "(II[II[[IILjava/lang/String;IILjava/lang/String;ZLjava/lang/String;Ljava/lang/String;)V",
         (void *) &nativeSpecializeAppProcess_samsung_q
     },
 };
-constexpr int nativeSpecializeAppProcess_methods_num = std::size(nativeSpecializeAppProcess_methods);
 
 void *nativeForkSystemServer_orig = nullptr;
 [[clang::no_stack_protector]] jint nativeForkSystemServer_l(JNIEnv *env, jclass clazz, jint uid, jint gid, jintArray gids, jint runtime_flags, jobjectArray rlimits, jlong permitted_capabilities, jlong effective_capabilities) {
@@ -250,45 +248,48 @@ void *nativeForkSystemServer_orig = nullptr;
     ctx.nativeForkSystemServer_post();
     return ctx.pid;
 }
-const JNINativeMethod nativeForkSystemServer_methods[] = {
-    {
+std::array nativeForkSystemServer_methods {
+    JNINativeMethod {
         "nativeForkSystemServer",
         "(II[II[[IJJ)I",
         (void *) &nativeForkSystemServer_l
     },
-    {
+    JNINativeMethod {
         "nativeForkSystemServer",
         "(II[IIII[[IJJ)I",
         (void *) &nativeForkSystemServer_samsung_q
     },
 };
-constexpr int nativeForkSystemServer_methods_num = std::size(nativeForkSystemServer_methods);
-
-unique_ptr<JNINativeMethod[]> hookAndSaveJNIMethods(const char *className, const JNINativeMethod *methods, int numMethods) {
-    unique_ptr<JNINativeMethod[]> newMethods;
-    int clz_id = -1;
-    int hook_cnt = 0;
-    do {
-        if (className == "com/android/internal/os/Zygote"sv) {
-            clz_id = 0;
-            hook_cnt = 3;
-            break;
-        }
-    } while (false);
-    if (hook_cnt) {
-        newMethods = make_unique<JNINativeMethod[]>(numMethods);
-        memcpy(newMethods.get(), methods, sizeof(JNINativeMethod) * numMethods);
-    }
-    auto &class_map = (*jni_method_map)[className];
-    for (int i = 0; i < numMethods; ++i) {
-        if (hook_cnt && clz_id == 0) {
-            HOOK_JNI(nativeForkAndSpecialize)
-            HOOK_JNI(nativeSpecializeAppProcess)
-            HOOK_JNI(nativeForkSystemServer)
-        }
-        class_map[methods[i].name][methods[i].signature] = methods[i].fnPtr;
-    }
-    return newMethods;
-}
 
 } // namespace
+
+static void do_hook_zygote(JNIEnv *env) {
+    vector<JNINativeMethod> hooks;
+    const char *clz;
+    clz = "com/android/internal/os/Zygote";
+    hookJniNativeMethods(env, clz, nativeForkAndSpecialize_methods.data(), nativeForkAndSpecialize_methods.size());
+    for (auto &method : nativeForkAndSpecialize_methods) {
+        if (method.fnPtr) {
+            nativeForkAndSpecialize_orig = method.fnPtr;
+            hooks.emplace_back(method);
+            break;
+        }
+    }
+    hookJniNativeMethods(env, clz, nativeSpecializeAppProcess_methods.data(), nativeSpecializeAppProcess_methods.size());
+    for (auto &method : nativeSpecializeAppProcess_methods) {
+        if (method.fnPtr) {
+            nativeSpecializeAppProcess_orig = method.fnPtr;
+            hooks.emplace_back(method);
+            break;
+        }
+    }
+    hookJniNativeMethods(env, clz, nativeForkSystemServer_methods.data(), nativeForkSystemServer_methods.size());
+    for (auto &method : nativeForkSystemServer_methods) {
+        if (method.fnPtr) {
+            nativeForkSystemServer_orig = method.fnPtr;
+            hooks.emplace_back(method);
+            break;
+        }
+    }
+    jni_hook_list->emplace(clz, std::move(hooks));
+}

loader/src/injector/memory.cpp

@@ -1,31 +0,0 @@
-#include "memory.hpp"
-
-namespace jni_hook {
-
-// We know our minimum alignment is WORD size (size of pointer)
-static constexpr size_t ALIGN = sizeof(long);
-
-// 4MB is more than enough
-static constexpr size_t CAPACITY = (1 << 22);
-
-// 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 *>(mmap(
-                nullptr, CAPACITY, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));
-        _curr = _area;
-    }
-    return _curr.fetch_add(align_to(sz, ALIGN));
-}
-
-void memory_block::release() {
-    if (_area)
-        munmap(_area, CAPACITY);
-}
-
-} // namespace jni_hook

loader/src/injector/memory.hpp

@@ -1,44 +0,0 @@
-#pragma once
-
-#include <map>
-#include <sys/mman.h>
-
-#pragma clang diagnostic push
-#include <parallel_hashmap/phmap.h>
-#pragma clang diagnostic pop
-
-#include "misc.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