[core] Update Zygisk API

core/src/main/cpp/main/api/zygisk_main.cpp

@@ -27,8 +27,6 @@
 #include "config.h"
 
 namespace lspd {
-    int *allowUnload = nullptr;
-
     namespace {
         ssize_t xsendmsg(int sockfd, const struct msghdr *msg, int flags) {
             int sent = sendmsg(sockfd, msg, flags);
@@ -167,13 +165,19 @@ namespace lspd {
             if (fd < 0) return;
             xwrite(fd, &val, sizeof(val));
         }
+
+        int allow_unload = 0;
     }
 
+    int *allowUnload = &allow_unload;
+
     class ZygiskModule : public zygisk::ModuleBase {
         JNIEnv *env_;
+        zygisk::Api *api_;
 
         void onLoad(zygisk::Api *api, JNIEnv *env) override {
             env_ = env;
+            api_ = api;
             Context::GetInstance()->Init();
 
             auto companion = api->connectCompanion();
@@ -201,6 +205,7 @@ namespace lspd {
         void postAppSpecialize(const zygisk::AppSpecializeArgs *args) override {
             Context::GetInstance()->OnNativeForkAndSpecializePost(env_, args->nice_name,
                                                                   args->app_data_dir);
+            if (*allowUnload) api_->setOption(zygisk::DLCLOSE_MODULE_LIBRARY);
         }
 
         void preServerSpecialize(zygisk::ServerSpecializeArgs *args) override {

core/src/main/cpp/main/src/jni/zygisk.h

@@ -24,16 +24,18 @@ class ExampleModule : public zygisk::ModuleBase {
 public:
     void onLoad(zygisk::Api *api, JNIEnv *env) override {
         this->api = api;
+        this->env = env;
     }
     void preAppSpecialize(zygisk::AppSpecializeArgs *args) override {
         JNINativeMethod methods[] = {
             { "logger_entry_max_payload_native", "()I", (void*) my_logger_entry_max },
         };
-        api->hookJniNativeMethods("android/util/Log", methods, 1);
+        api->hookJniNativeMethods(env, "android/util/Log", methods, 1);
         *(void **) &orig_logger_entry_max = methods[0].fnPtr;
     }
 private:
     zygisk::Api *api;
+    JNIEnv *env;
 };
 REGISTER_ZYGISK_MODULE(ExampleModule)
 REGISTER_ZYGISK_COMPANION(example_handler)
@@ -76,7 +78,7 @@ namespace zygisk {
         // See preAppSpecialize(args) for more info.
         virtual void preServerSpecialize(ServerSpecializeArgs *args) {}
 
-        // This function is called after the app process is specialized.
+        // This function is called after the system server process is specialized.
         // At this point, the process runs with the privilege of system_server.
         virtual void postServerSpecialize(const ServerSpecializeArgs *args) {}
     };
@@ -120,6 +122,27 @@ namespace zygisk {
         template <class T> void entry_impl(api_table *, JNIEnv *);
     }
 
+// These values are used in Api::setOption(Option)
+    enum Option : int {
+        // Force Magisk's denylist unmount routines to run on this process.
+        //
+        // Setting this option only makes sense in preAppSpecialize.
+        // The actual unmounting happens during app process specialization.
+        //
+        // Processes added to Magisk's denylist will have all Magisk and its modules' files unmounted
+        // from its mount namespace. In addition, all Zygisk code will be unloaded from memory, which
+        // also implies that no Zygisk modules (including yours) are loaded.
+        //
+        // However, if for any reason your module still wants the unmount part of the denylist
+        // operation to be enabled EVEN IF THE PROCESS IS NOT ON THE DENYLIST, set this option.
+        FORCE_DENYLIST_UNMOUNT = 0,
+
+        // When this option is set, your module's library will be dlclose-ed after post[XXX]Specialize.
+        // Be aware that after dlclose-ing your module, all of your code will be unmapped.
+        // YOU SHOULD NOT ENABLE THIS OPTION AFTER HOOKING ANY FUNCTION IN THE PROCESS.
+        DLCLOSE_MODULE_LIBRARY = 1,
+    };
+
     struct Api {
 
         // Connect to a root companion process and get a Unix domain socket for IPC.
@@ -132,30 +155,17 @@ namespace zygisk {
         // Another good use case for a companion process is that if you want to share some resources
         // across multiple processes, hold the resources in the companion process and pass it over.
         //
-        // When this function is called, in the companion process, a socket pair will be created,
+        // The root companion process is ABI aware; that is, when calling this function from a 32-bit
-        // your module's onCompanionRequest(int) callback will receive one socket, and the other
+        // process, you will be connected to a 32-bit companion process, and vice versa for 64-bit.
-        // socket will be returned.
         //
-        // Returns a file descriptor to a socket that is connected to the socket passed to
+        // Returns a file descriptor to a socket that is connected to the socket passed to your
-        // your module's onCompanionRequest(int). Returns -1 if the connection attempt failed.
+        // module's companion request handler. Returns -1 if the connection attempt failed.
         int connectCompanion();
 
-        // Force Magisk's denylist unmount routines to run on this process.
+        // Set various options for your module.
-        //
+        // Please note that this function accepts one single option at a time.
-        // This API only works in preAppSpecialize.
+        // Check zygisk::Option for the full list of options available.
-        //
+        void setOption(Option opt);
-        // Processes added to Magisk's denylist will have all Magisk and its modules' files unmounted
-        // from its mount namespace. In addition, all Zygisk code will be unloaded from memory, which
-        // also implies that no Zygisk modules (including yours) are loaded.
-        //
-        // However, if for any reason your module still wants the unmount part of the denylist
-        // operation to be enabled EVEN IF THE PROCESS IS NOT ON THE DENYLIST, call this function.
-        // No code will be unloaded from memory (including your module) because there is no way to
-        // guarantee no crashes will occur.
-        //
-        // The unmounting does not happen immediately after the function is called. It is actually
-        // done during app process specialization.
-        void forceDenyListUnmount();
 
         // Hook JNI native methods for a class
         //
@@ -239,7 +249,7 @@ void zygisk_companion_entry(int client) { func(client); }
 
             // Zygisk functions
             int  (*connectCompanion)(void * /* _this */);
-            void (*forceDenyListUnmount)(void * /* _this */);
+            void (*setOption)(void * /* _this */, Option);
         };
 
         template <class T>
@@ -257,8 +267,8 @@ void zygisk_companion_entry(int client) { func(client); }
     int Api::connectCompanion() {
         return impl->connectCompanion(impl->_this);
     }
-    void Api::forceDenyListUnmount() {
+    void Api::setOption(Option opt) {
-        impl->forceDenyListUnmount(impl->_this);
+        impl->setOption(impl->_this, opt);
     }
     void Api::hookJniNativeMethods(JNIEnv *env, const char *className, JNINativeMethod *methods, int numMethods) {
         impl->hookJniNativeMethods(env, className, methods, numMethods);