Always unload zygisk after specialize

native/jni/zygisk/api.hpp

@@ -152,6 +152,8 @@ enum Option : int {
     DLCLOSE_MODULE_LIBRARY = 1,
 };
 
+// All API functions will stop working after post[XXX]Specialize as Zygisk will be unloaded
+// from the specialized process afterwards.
 struct Api {
 
     // Connect to a root companion process and get a Unix domain socket for IPC.
@@ -274,22 +276,22 @@ void entry_impl(api_table *table, JNIEnv *env) {
 } // namespace internal
 
 inline int Api::connectCompanion() {
-    return impl->connectCompanion(impl->_this);
+    return impl->connectCompanion ? impl->connectCompanion(impl->_this) : -1;
 }
 inline void Api::setOption(Option opt) {
-    impl->setOption(impl->_this, opt);
+    if (impl->setOption) impl->setOption(impl->_this, opt);
 }
 inline void Api::hookJniNativeMethods(JNIEnv *env, const char *className, JNINativeMethod *methods, int numMethods) {
-    impl->hookJniNativeMethods(env, className, methods, numMethods);
+    if (impl->hookJniNativeMethods) impl->hookJniNativeMethods(env, className, methods, numMethods);
 }
 inline void Api::pltHookRegister(const char *regex, const char *symbol, void *newFunc, void **oldFunc) {
-    impl->pltHookRegister(regex, symbol, newFunc, oldFunc);
+    if (impl->pltHookRegister) impl->pltHookRegister(regex, symbol, newFunc, oldFunc);
 }
 inline void Api::pltHookExclude(const char *regex, const char *symbol) {
-    impl->pltHookExclude(regex, symbol);
+    if (impl->pltHookExclude) impl->pltHookExclude(regex, symbol);
 }
 inline bool Api::pltHookCommit() {
-    return impl->pltHookCommit();
+    return impl->pltHookCommit != nullptr && impl->pltHookCommit();
 }
 
 } // namespace zygisk

native/jni/zygisk/entry.cpp

@@ -13,7 +13,7 @@
 
 using namespace std;
 
-static void *self_handle = nullptr;
+void *self_handle = nullptr;
 
 static int zygisk_log(int prio, const char *fmt, va_list ap);
 
@@ -26,14 +26,6 @@ static void zygisk_logging() {
     log_cb.ex = nop_ex;
 }
 
-void self_unload() {
-    ZLOGD("Request to self unload\n");
-    // If unhooking failed, do not unload or else it will cause SIGSEGV
-    if (!unhook_functions())
-        return;
-    new_daemon_thread(reinterpret_cast<thread_entry>(&dlclose), self_handle);
-}
-
 static char *first_stage_path = nullptr;
 void unload_first_stage() {
     if (first_stage_path) {

native/jni/zygisk/hook.cpp

@@ -21,6 +21,8 @@ using xstring = jni_hook::string;
 //#define ZLOGV(...) ZLOGD(__VA_ARGS__)
 #define ZLOGV(...)
 
+static bool unhook_functions();
+
 namespace {
 
 enum {
@@ -28,6 +30,7 @@ enum {
     FORK_AND_SPECIALIZE,
     APP_SPECIALIZE,
     SERVER_SPECIALIZE,
+    CAN_DLCLOSE,
     FLAG_MAX
 };
 
@@ -51,6 +54,7 @@ struct HookContext {
     HookContext() : pid(-1), info{} {}
 
     static void close_fds();
+    void unload_zygisk();
 
     DCL_PRE_POST(fork)
     void run_modules_pre(const vector<int> &fds);
@@ -170,6 +174,15 @@ DCL_HOOK_FUNC(void, android_log_close) {
     old_android_log_close();
 }
 
+// Last point before process secontext changes
+DCL_HOOK_FUNC(int, selinux_android_setcontext,
+        uid_t uid, int isSystemServer, const char *seinfo, const char *pkgname) {
+    if (g_ctx) {
+        g_ctx->flags[CAN_DLCLOSE] = unhook_functions();
+    }
+    return old_selinux_android_setcontext(uid, isSystemServer, seinfo, pkgname);
+}
+
 // -----------------------------------------------------------------
 
 // The original android::AppRuntime virtual table
@@ -327,6 +340,7 @@ void HookContext::run_modules_pre(const vector<int> &fds) {
                 modules.emplace_back(i, h);
                 auto api = new ApiTable(&modules.back());
                 module_entry(api, env);
+                modules.back().table = api;
             }
         }
         close(fds[i]);
@@ -357,6 +371,22 @@ void HookContext::close_fds() {
     close(logd_fd.exchange(-1));
 }
 
+void HookContext::unload_zygisk() {
+    if (flags[CAN_DLCLOSE]) {
+        // Do NOT call the destructor
+        operator delete(jni_method_map);
+        // Directly unmap the whole memory block
+        jni_hook::memory_block::release();
+
+        // Strip out all API function pointers
+        for (auto &m : modules) {
+            memset(m.table, 0, sizeof(*m.table));
+        }
+
+        new_daemon_thread(reinterpret_cast<thread_entry>(&dlclose), self_handle);
+    }
+}
+
 // -----------------------------------------------------------------
 
 void HookContext::nativeSpecializeAppProcess_pre() {
@@ -387,19 +417,14 @@ void HookContext::nativeSpecializeAppProcess_post() {
     }
 
     env->ReleaseStringUTFChars(args->nice_name, process);
-    if (info.on_denylist) {
+    run_modules_post();
-        self_unload();
+    if (info.is_magisk_app) {
-    } else {
+        setenv("ZYGISK_ENABLED", "1", 1);
-        run_modules_post();
-        if (info.is_magisk_app) {
-            setenv("ZYGISK_ENABLED", "1", 1);
-        } else if (args->is_child_zygote && *args->is_child_zygote) {
-            // If we are in child zygote, unhook all zygisk hooks
-            // Modules still have their code loaded and can do whatever they want
-            unhook_functions();
-        }
     }
     g_ctx = nullptr;
+    if (!flags[FORK_AND_SPECIALIZE]) {
+        unload_zygisk();
+    }
 }
 
 void HookContext::nativeForkSystemServer_pre() {
@@ -457,6 +482,7 @@ void HookContext::fork_pre() {
 void HookContext::fork_post() {
     sigmask(SIG_UNBLOCK, SIGCHLD);
     g_ctx = nullptr;
+    unload_zygisk();
 }
 
 } // namespace
@@ -464,7 +490,6 @@ void HookContext::fork_post() {
 static bool hook_refresh() {
     if (xhook_refresh(0) == 0) {
         xhook_clear();
-        ZLOGI("xhook success\n");
         return true;
     } else {
         ZLOGE("xhook failed\n");
@@ -503,6 +528,7 @@ void hook_functions() {
     XHOOK_REGISTER(ANDROID_RUNTIME, fork);
     XHOOK_REGISTER(ANDROID_RUNTIME, unshare);
     XHOOK_REGISTER(ANDROID_RUNTIME, jniRegisterNativeMethods);
+    XHOOK_REGISTER(ANDROID_RUNTIME, selinux_android_setcontext);
     XHOOK_REGISTER_SYM(ANDROID_RUNTIME, "__android_log_close", android_log_close);
     hook_refresh();
 
@@ -526,7 +552,7 @@ void hook_functions() {
     }
 }
 
-bool unhook_functions() {
+static bool unhook_functions() {
     bool success = true;
 
     // Restore JNIEnv
@@ -539,11 +565,6 @@ bool unhook_functions() {
         }
     }
 
-    // Do NOT call the destructor
-    operator delete(jni_method_map);
-    // Directly unmap the whole memory block
-    jni_hook::memory_block::release();
-
     // Unhook JNI methods
     for (const auto &[clz, methods] : *jni_hook_list) {
         if (!methods.empty() && old_jniRegisterNativeMethods(

native/jni/zygisk/module.hpp

@@ -88,6 +88,7 @@ struct ZygiskModule {
     ZygiskModule(int id, void *handle) : handle(handle), id(id) {}
 
     void * const handle;
+    ApiTable *table = nullptr;
     bool unload = false;
 
 private:

native/jni/zygisk/zygisk.hpp

@@ -41,9 +41,9 @@ struct AppInfo {
     bool on_denylist;
 };
 
+extern void *self_handle;
+
 void unload_first_stage();
-void self_unload();
 void hook_functions();
-bool unhook_functions();
 std::vector<int> remote_get_info(int uid, const char *process, AppInfo *info);
 int remote_request_unmount();