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,23 +28,14 @@ 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() {
     char *cur = environ[0];
 
     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(void *handle, const char *tmp, char *path) {
-    self_handle = handle;
-    MAGISKTMP = tmp;
-    unsetenv(INJECT_ENV_2);
-    unsetenv(SECOND_STAGE_PTR);
-
+static void second_stage_entry() {
     zygisk_logging();
     ZLOGD("inject 2nd stage\n");
-    hook_functions();
 
-    // First stage will be unloaded before the first fork
-    first_stage_path = path;
+    char path[PATH_MAX];
+    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;
-
-    // Update path to 2nd stage lib
-    *num = '2';
+    // 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.
+
+    int fd = xopen(path, O_RDONLY | O_CLOEXEC);
+    // 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);
-    remap_all(path);
-
-    // 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);
+    // Force dlopen ourselves to make ourselves dlclose-able.
+    // After this call, all global variables will be reset.
+    android_dlopen_ext(path, RTLD_LAZY, &info);
 }
 
-__attribute__((constructor))
+[[gnu::constructor]] [[maybe_unused]]
 static void zygisk_init() {
-    if (getenv(INJECT_ENV_2)) {
-        // Return function pointer to first stage
-        char buf[128];
-        snprintf(buf, sizeof(buf), "%p", &second_stage_entry);
-        setenv(SECOND_STAGE_PTR, buf, 1);
-    } else if (getenv(INJECT_ENV_1)) {
+    if (getenv(INJECT_ENV_1)) {
+        unsetenv(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]);
-        if (void *h = dlopen(buf, RTLD_LAZY)) {
+        android_dlextinfo info {
+            .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__)
-            string lib = tmp + "/" ZYGISKBIN "/zygisk.app_process64.1.so";
-#else
-            string lib = tmp + "/" ZYGISKBIN "/zygisk.app_process32.1.so";
-#endif
+            xreadlink("/proc/self/exe", buf, sizeof(buf));
             if (char *ld = getenv("LD_PRELOAD")) {
-                char env[256];
-                sprintf(env, "%s:%s", ld, lib.data());
-                setenv("LD_PRELOAD", env, 1);
+                string env = ld;
+                env += ':';
+                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);