Refactor Zygisk loading

Co-authored-by: topjohnwu <topjohnwu@gmail.com>

native/jni/utils/missing.hpp

@@ -1,7 +1,10 @@
 #pragma once
 
 #include <sys/syscall.h>
+#include <linux/fcntl.h>
 #include <unistd.h>
+#include <cerrno>
+#include <cstdio>
 
 static inline int sigtimedwait(const sigset_t* set, siginfo_t* info, const timespec* timeout) {
     union {
@@ -11,3 +14,13 @@ static inline int sigtimedwait(const sigset_t* set, siginfo_t* info, const times
     s.set = *set;
     return syscall(__NR_rt_sigtimedwait, &s.set64, info, timeout, sizeof(sigset64_t));
 }
+
+static inline int fexecve(int fd, char* const* argv, char* const* envp) {
+    syscall(__NR_execveat, fd, "", argv, envp, AT_EMPTY_PATH);
+    if (errno == ENOSYS) {
+        char buf[256];
+        std::snprintf(buf, sizeof(buf), "/proc/self/fd/%d", fd);
+        execve(buf, argv, envp);
+    }
+    return -1;
+}

native/jni/zygisk/entry.cpp

@@ -2,6 +2,7 @@
 #include <dlfcn.h>
 #include <sys/prctl.h>
 #include <android/log.h>
+#include <android/dlext.h>
 
 #include <utils.hpp>
 #include <daemon.hpp>
@@ -27,15 +28,6 @@ static void zygisk_logging() {
     log_cb.ex = nop_ex;
 }
 
-static char *first_stage_path = nullptr;
-void unload_first_stage() {
-    if (first_stage_path) {
-        unmap_all(first_stage_path);
-        free(first_stage_path);
-        first_stage_path = nullptr;
-    }
-}
-
 // Make sure /proc/self/environ is sanitized
 // Filter env and reset MM_ENV_END
 static void sanitize_environ() {
@@ -43,7 +35,7 @@ static void sanitize_environ() {
 
     for (int i = 0; environ[i]; ++i) {
         // Copy all env onto the original stack
-        int len = strlen(environ[i]);
+        size_t len = strlen(environ[i]);
         memmove(cur, environ[i], len + 1);
         environ[i] = cur;
         cur += len + 1;
@@ -52,7 +44,7 @@ static void sanitize_environ() {
     prctl(PR_SET_MM, PR_SET_MM_ENV_END, cur, 0, 0);
 }
 
-__attribute__((destructor))
+[[gnu::destructor]] [[maybe_unused]]
 static void zygisk_cleanup_wait() {
     if (self_handle) {
         // Wait 10us to make sure none of our code is executing
@@ -61,71 +53,84 @@ static void zygisk_cleanup_wait() {
     }
 }
 
-#define SECOND_STAGE_PTR "ZYGISK_PTR"
+static void second_stage_entry() {
-
-static void second_stage_entry(void *handle, const char *tmp, char *path) {
-    self_handle = handle;
-    MAGISKTMP = tmp;
-    unsetenv(INJECT_ENV_2);
-    unsetenv(SECOND_STAGE_PTR);
-
     zygisk_logging();
     ZLOGD("inject 2nd stage\n");
-    hook_functions();
 
-    // First stage will be unloaded before the first fork
+    char path[PATH_MAX];
-    first_stage_path = path;
+    MAGISKTMP = getenv(MAGISKTMP_ENV);
+    int fd = parse_int(getenv(MAGISKFD_ENV));
+
+    snprintf(path, sizeof(path), "/proc/self/fd/%d", fd);
+    xreadlink(path, path, PATH_MAX);
+    android_dlextinfo info {
+        .flags = ANDROID_DLEXT_USE_LIBRARY_FD,
+        .library_fd = fd,
+    };
+    self_handle = android_dlopen_ext(path, RTLD_LAZY, &info);
+    dlclose(self_handle);
+    close(fd);
+    unsetenv(MAGISKTMP_ENV);
+    unsetenv(MAGISKFD_ENV);
+    sanitize_environ();
+    hook_functions();
 }
 
 static void first_stage_entry() {
     android_logging();
     ZLOGD("inject 1st stage\n");
 
+    char path[PATH_MAX];
+    char buf[256];
     char *ld = getenv("LD_PRELOAD");
-    char tmp[128];
-    strlcpy(tmp, getenv("MAGISKTMP"), sizeof(tmp));
-    char *path;
     if (char *c = strrchr(ld, ':')) {
         *c = '\0';
+        strlcpy(path, c + 1, sizeof(path));
         setenv("LD_PRELOAD", ld, 1);  // Restore original LD_PRELOAD
-        path = strdup(c + 1);
     } else {
         unsetenv("LD_PRELOAD");
-        path = strdup(ld);
+        strlcpy(path, ld, sizeof(path));
     }
-    unsetenv(INJECT_ENV_1);
-    unsetenv("MAGISKTMP");
-    sanitize_environ();
 
-    char *num = strrchr(path, '.') - 1;
+    // Force the linker to load the library on top of ourselves, so we do not
+    // need to unmap the 1st stage library that was loaded with LD_PRELOAD.
 
-    // Update path to 2nd stage lib
+    int fd = xopen(path, O_RDONLY | O_CLOEXEC);
-    *num = '2';
+    // Use fd here instead of path to make sure inode is the same as 2nd stage
+    snprintf(buf, sizeof(buf), "%d", fd);
+    setenv(MAGISKFD_ENV, buf, 1);
+    struct stat s{};
+    xfstat(fd, &s);
+
+    android_dlextinfo info {
+        .flags = ANDROID_DLEXT_FORCE_LOAD | ANDROID_DLEXT_USE_LIBRARY_FD,
+        .library_fd = fd,
+    };
+    auto [addr, size] = find_map_range(path, s.st_ino);
+    if (addr && size) {
+        info.flags |= ANDROID_DLEXT_RESERVED_ADDRESS;
+        info.reserved_addr = addr;
+        // The existing address is guaranteed to fit, as 1st stage and 2nd stage
+        // are exactly the same ELF (same inode). However, the linker could over
+        // estimate the required size and refuse to dlopen. Add 2 more page_sizes
+        // (one at the beginning and one at the end) as a safety measure.
+        info.reserved_size = size + 2 * 4096;
+    }
 
-    // Load second stage
     setenv(INJECT_ENV_2, "1", 1);
-    void *handle = dlopen(path, RTLD_LAZY);
+    // Force dlopen ourselves to make ourselves dlclose-able.
-    remap_all(path);
+    // After this call, all global variables will be reset.
-
+    android_dlopen_ext(path, RTLD_LAZY, &info);
-    // Revert path to 1st stage lib
-    *num = '1';
-
-    // Run second stage entry
-    char *env = getenv(SECOND_STAGE_PTR);
-    decltype(&second_stage_entry) second_stage;
-    sscanf(env, "%p", &second_stage);
-    second_stage(handle, tmp, path);
 }
 
-__attribute__((constructor))
+[[gnu::constructor]] [[maybe_unused]]
 static void zygisk_init() {
-    if (getenv(INJECT_ENV_2)) {
+    if (getenv(INJECT_ENV_1)) {
-        // Return function pointer to first stage
+        unsetenv(INJECT_ENV_1);
-        char buf[128];
-        snprintf(buf, sizeof(buf), "%p", &second_stage_entry);
-        setenv(SECOND_STAGE_PTR, buf, 1);
-    } else if (getenv(INJECT_ENV_1)) {
         first_stage_entry();
+    } else if (getenv(INJECT_ENV_2)) {
+        unsetenv(INJECT_ENV_2);
+        second_stage_entry();
     }
 }
 
@@ -167,7 +172,7 @@ static int zygisk_log(int prio, const char *fmt, va_list ap) {
     return ret;
 }
 
-static inline bool should_load_modules(int flags) {
+static inline bool should_load_modules(uint32_t flags) {
     return (flags & UNMOUNT_MASK) != UNMOUNT_MASK &&
            (flags & PROCESS_IS_MAGISK_APP) != PROCESS_IS_MAGISK_APP;
 }
@@ -304,10 +309,6 @@ static void setup_files(int client, const sock_cred *cred) {
 
     write_int(client, 0);
     send_fd(client, is_64_bit ? app_process_64 : app_process_32);
-
-    string path = MAGISKTMP + "/" ZYGISKBIN "/zygisk." + basename(buf);
-    cp_afc(buf, (path + ".1.so").data());
-    cp_afc(buf, (path + ".2.so").data());
     write_string(client, MAGISKTMP);
 }
 

native/jni/zygisk/hook.cpp

@@ -1,8 +1,10 @@
-#include <dlfcn.h>
+#include <android/dlext.h>
 #include <sys/mount.h>
-#include <xhook.h>
+#include <dlfcn.h>
 #include <bitset>
 
+#include <xhook.h>
+
 #include <utils.hpp>
 #include <flags.h>
 #include <daemon.hpp>
@@ -147,7 +149,6 @@ DCL_HOOK_FUNC(int, jniRegisterNativeMethods,
 // Skip actual fork and return cached result if applicable
 // Also unload first stage zygisk if necessary
 DCL_HOOK_FUNC(int, fork) {
-    unload_first_stage();
     return (g_ctx && g_ctx->pid >= 0) ? g_ctx->pid : old_fork();
 }
 
@@ -360,15 +361,17 @@ uint32_t ZygiskModule::getFlags() {
 }
 
 void HookContext::run_modules_pre(const vector<int> &fds) {
-    char buf[256];
 
     // Since we directly use the pointer to elements in the vector, in order to prevent dangling
     // pointers, the vector has to be pre-allocated to ensure reallocation does not occur
     modules.reserve(fds.size());
 
     for (int i = 0; i < fds.size(); ++i) {
-        snprintf(buf, sizeof(buf), "/proc/self/fd/%d", fds[i]);
+        android_dlextinfo info {
-        if (void *h = dlopen(buf, RTLD_LAZY)) {
+            .flags = ANDROID_DLEXT_USE_LIBRARY_FD,
+            .library_fd = fds[i],
+        };
+        if (void *h = android_dlopen_ext("/jit-cache", RTLD_LAZY, &info)) {
             if (void *e = dlsym(h, "zygisk_module_entry")) {
                 modules.emplace_back(i, h, e);
             }

native/jni/zygisk/main.cpp

@@ -1,4 +1,5 @@
 #include <sys/mount.h>
+#include <android/dlext.h>
 #include <dlfcn.h>
 
 #include <magisk.hpp>
@@ -13,7 +14,7 @@ using namespace std;
 // Entrypoint for app_process overlay
 int app_process_main(int argc, char *argv[]) {
     android_logging();
-    char buf[256];
+    char buf[PATH_MAX];
 
     bool zygote = false;
     if (auto fp = open_file("/proc/self/attr/current", "r")) {
@@ -51,9 +52,8 @@ int app_process_main(int argc, char *argv[]) {
             return 1;
         close(fds[0]);
 
-        snprintf(buf, sizeof(buf), "/proc/self/fd/%d", app_proc_fd);
         fcntl(app_proc_fd, F_SETFD, FD_CLOEXEC);
-        execve(buf, argv, environ);
+        fexecve(app_proc_fd, argv, environ);
         return 1;
     }
 
@@ -70,26 +70,22 @@ int app_process_main(int argc, char *argv[]) {
                 break;
 
             string tmp = read_string(socket);
-#if defined(__LP64__)
+            xreadlink("/proc/self/exe", buf, sizeof(buf));
-            string lib = tmp + "/" ZYGISKBIN "/zygisk.app_process64.1.so";
-#else
-            string lib = tmp + "/" ZYGISKBIN "/zygisk.app_process32.1.so";
-#endif
             if (char *ld = getenv("LD_PRELOAD")) {
-                char env[256];
+                string env = ld;
-                sprintf(env, "%s:%s", ld, lib.data());
+                env += ':';
-                setenv("LD_PRELOAD", env, 1);
+                env += buf;
+                setenv("LD_PRELOAD", env.data(), 1);
             } else {
-                setenv("LD_PRELOAD", lib.data(), 1);
+                setenv("LD_PRELOAD", buf, 1);
             }
             setenv(INJECT_ENV_1, "1", 1);
-            setenv("MAGISKTMP", tmp.data(), 1);
+            setenv(MAGISKTMP_ENV, tmp.data(), 1);
 
             close(socket);
 
-            snprintf(buf, sizeof(buf), "/proc/self/fd/%d", app_proc_fd);
             fcntl(app_proc_fd, F_SETFD, FD_CLOEXEC);
-            execve(buf, argv, environ);
+            fexecve(app_proc_fd, argv, environ);
         } while (false);
 
         close(socket);

native/jni/zygisk/utils.cpp

@@ -12,9 +12,9 @@ struct map_info {
     uintptr_t end;
     uintptr_t off;
     int perms;
-    const char *path;
+    unsigned long inode;
 
-    map_info() : start(0), end(0), off(0), perms(0), path(nullptr) {}
+    map_info() : start(0), end(0), off(0), perms(0), inode(0) {}
 };
 
 } // namespace
@@ -31,8 +31,8 @@ static void parse_maps(int pid, Func fn) {
         char perm[5];
         int path_off;
 
-        if (sscanf(line, "%" PRIxPTR "-%" PRIxPTR " %4s %" PRIxPTR " %*x:%*x %*d %n%*s",
+        if (sscanf(line, "%" PRIxPTR "-%" PRIxPTR " %4s %" PRIxPTR " %*x:%*x %lu %n%*s",
-                   &info.start, &info.end, perm, &info.off, &path_off) != 4)
+                   &info.start, &info.end, perm, &info.off, &info.inode, &path_off) != 5)
             return true;
 
         // Parse permissions
@@ -43,22 +43,38 @@ static void parse_maps(int pid, Func fn) {
         if (perm[2] != '-')
             info.perms |= PROT_EXEC;
 
-        info.path = line + path_off;
+        return fn(info, line + path_off);
-
-        return fn(info);
     });
 }
 
 static vector<map_info> find_maps(const char *name) {
     vector<map_info> maps;
-    parse_maps(getpid(), [=, &maps](map_info &info) -> bool {
+    parse_maps(getpid(), [=, &maps](const map_info &info, const char *path) -> bool {
-        if (strcmp(info.path, name) == 0)
+        if (strcmp(path, name) == 0)
             maps.emplace_back(info);
         return true;
     });
     return maps;
 }
 
+std::pair<void *, size_t> find_map_range(const char *name, unsigned long inode) {
+    vector<map_info> maps = find_maps(name);
+    uintptr_t start = 0u;
+    uintptr_t end = 0u;
+    for (const auto &map : maps) {
+        if (map.inode == inode) {
+            if (start == 0) {
+                start = map.start;
+                end = map.end;
+            } else if (map.start == end) {
+                end = map.end;
+            }
+        }
+    }
+    LOGD("found map %s with start = %zx, end = %zx\n", name, start, end);
+    return make_pair(reinterpret_cast<void *>(start), end - start);
+}
+
 void unmap_all(const char *name) {
     vector<map_info> maps = find_maps(name);
     for (map_info &info : maps) {
@@ -91,10 +107,10 @@ void remap_all(const char *name) {
 
 uintptr_t get_function_off(int pid, uintptr_t addr, char *lib) {
     uintptr_t off = 0;
-    parse_maps(pid, [=, &off](map_info &info) -> bool {
+    parse_maps(pid, [=, &off](const map_info &info, const char *path) -> bool {
         if (addr >= info.start && addr < info.end) {
             if (lib)
-                strcpy(lib, info.path);
+                strcpy(lib, path);
             off = addr - info.start + info.off;
             return false;
         }
@@ -105,8 +121,8 @@ uintptr_t get_function_off(int pid, uintptr_t addr, char *lib) {
 
 uintptr_t get_function_addr(int pid, const char *lib, uintptr_t off) {
     uintptr_t addr = 0;
-    parse_maps(pid, [=, &addr](map_info &info) -> bool {
+    parse_maps(pid, [=, &addr](const map_info &info, const char *path) -> bool {
-        if (strcmp(info.path, lib) == 0 && (info.perms & PROT_EXEC)) {
+        if (strcmp(path, lib) == 0 && (info.perms & PROT_EXEC)) {
             addr = info.start - info.off + off;
             return false;
         }

native/jni/zygisk/zygisk.hpp

@@ -6,6 +6,8 @@
 
 #define INJECT_ENV_1   "MAGISK_INJ_1"
 #define INJECT_ENV_2   "MAGISK_INJ_2"
+#define MAGISKFD_ENV   "MAGISKFD"
+#define MAGISKTMP_ENV  "MAGISKTMP"
 
 enum : int {
     ZYGISK_SETUP,
@@ -25,6 +27,9 @@ enum : int {
 #define ZLOGI(...) LOGI("zygisk32: " __VA_ARGS__)
 #endif
 
+// Find the memory address + size of the pages matching name + inode
+std::pair<void*, size_t> find_map_range(const char *name, unsigned long inode);
+
 // Unmap all pages matching the name
 void unmap_all(const char *name);
 
@@ -39,7 +44,6 @@ uintptr_t get_function_addr(int pid, const char *lib, uintptr_t off);
 
 extern void *self_handle;
 
-void unload_first_stage();
 void hook_functions();
 int remote_get_info(int uid, const char *process, uint32_t *flags, std::vector<int> &fds);
 int remote_request_unmount();