Migrate zygisk handler to Rust

native/src/core/zygisk/daemon.rs

@@ -0,0 +1,101 @@
+use crate::consts::MODULEROOT;
+use crate::daemon::{to_user_id, IpcRead, MagiskD, UnixSocketExt};
+use crate::ffi::{update_deny_flags, ZygiskRequest, ZygiskStateFlags};
+use base::libc::{O_CLOEXEC, O_CREAT, O_RDONLY};
+use base::{cstr, open_fd, Directory, FsPathBuf, LoggedResult, Utf8CStrBufArr, WriteExt};
+use std::os::fd::{AsRawFd, FromRawFd, RawFd};
+use std::os::unix::net::UnixStream;
+
+const UNMOUNT_MASK: u32 =
+    ZygiskStateFlags::ProcessOnDenyList.repr | ZygiskStateFlags::DenyListEnforced.repr;
+
+pub fn zygisk_should_load_module(flags: u32) -> bool {
+    flags & UNMOUNT_MASK != UNMOUNT_MASK && flags & ZygiskStateFlags::ProcessIsMagiskApp.repr == 0
+}
+
+impl MagiskD {
+    pub fn zygisk_handler(&self, client: i32) {
+        let mut client = unsafe { UnixStream::from_raw_fd(client) };
+        let _: LoggedResult<()> = try {
+            let code = ZygiskRequest {
+                repr: client.ipc_read_int()?,
+            };
+            match code {
+                ZygiskRequest::GetInfo => self.get_process_info(client)?,
+                ZygiskRequest::ConnectCompanion => self.connect_zygiskd(client.as_raw_fd()),
+                ZygiskRequest::GetModDir => self.get_mod_dir(client)?,
+                _ => {}
+            }
+        };
+    }
+
+    pub fn get_module_fds(&self, is_64_bit: bool) -> Vec<RawFd> {
+        if let Some(module_list) = self.module_list.get() {
+            module_list
+                .iter()
+                .map(|m| if is_64_bit { m.z64 } else { m.z32 })
+                .collect()
+        } else {
+            Vec::new()
+        }
+    }
+
+    fn get_process_info(&self, mut client: UnixStream) -> LoggedResult<()> {
+        let uid = client.ipc_read_int()?;
+        let process = client.ipc_read_string()?;
+        let is_64_bit = client.ipc_read_int()? != 0;
+        let mut flags: u32 = 0;
+        update_deny_flags(uid, &process, &mut flags);
+        if self.get_manager_uid(to_user_id(uid)) == uid {
+            flags |= ZygiskStateFlags::ProcessIsMagiskApp.repr
+        }
+        if self.uid_granted_root(uid) {
+            flags |= ZygiskStateFlags::ProcessGrantedRoot.repr
+        }
+
+        // First send flags
+        client.write_pod(&flags)?;
+
+        // Next send modules
+        if zygisk_should_load_module(flags) {
+            if let Some(module_list) = self.module_list.get() {
+                let module_fds: Vec<RawFd> = module_list
+                    .iter()
+                    .map(|m| if is_64_bit { m.z64 } else { m.z32 })
+                    .collect();
+                client.send_fds(&module_fds)?;
+            }
+        }
+
+        // If we're not in system_server, we are done
+        if uid != 1000 || process != "system_server" {
+            return Ok(());
+        }
+
+        // Read all failed modules
+        let failed_ids: Vec<i32> = client.ipc_read_vec()?;
+        if let Some(module_list) = self.module_list.get() {
+            for id in failed_ids {
+                let mut buf = Utf8CStrBufArr::default();
+                let path = FsPathBuf::new(&mut buf)
+                    .join(MODULEROOT)
+                    .join(&module_list[id as usize].name)
+                    .join("zygisk");
+                // Create the unloaded marker file
+                Directory::open(&path)?.open_fd(cstr!("unloaded"), O_CREAT | O_RDONLY, 0o644)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn get_mod_dir(&self, mut client: UnixStream) -> LoggedResult<()> {
+        let id = client.ipc_read_int()?;
+        let module = &self.module_list.get().unwrap()[id as usize];
+        let mut buf = Utf8CStrBufArr::default();
+        let dir = FsPathBuf::new(&mut buf).join(MODULEROOT).join(&module.name);
+        let fd = open_fd!(&dir, O_RDONLY | O_CLOEXEC)?;
+        client.send_fds(&[fd.as_raw_fd()])?;
+        Ok(())
+    }
+}

native/src/core/zygisk/entry.cpp

@@ -15,8 +15,6 @@ using namespace std;
 
 string native_bridge = "0";
 
-using modules_t = const rust::Vec<ModuleInfo>&;
-
 static bool is_compatible_with(uint32_t) {
     zygisk_logging();
     hook_entry();
@@ -30,56 +28,13 @@ extern "C" [[maybe_unused]] NativeBridgeCallbacks NativeBridgeItf {
     .isCompatibleWith = &is_compatible_with,
 };
 
-// The following code runs in zygote/app process
-
-static inline bool should_load_modules(uint32_t flags) {
-    return (flags & UNMOUNT_MASK) != UNMOUNT_MASK &&
-           (flags & PROCESS_IS_MAGISK_APP) != PROCESS_IS_MAGISK_APP;
-}
-
-int remote_get_info(int uid, const char *process, uint32_t *flags, vector<int> &fds) {
-    if (int fd = zygisk_request(ZygiskRequest::GET_INFO); fd >= 0) {
-        write_int(fd, uid);
-        write_string(fd, process);
-#ifdef __LP64__
-        write_int(fd, 1);
-#else
-        write_int(fd, 0);
-#endif
-        xxread(fd, flags, sizeof(*flags));
-        if (should_load_modules(*flags)) {
-            fds = recv_fds(fd);
-        }
-        return fd;
-    }
-    return -1;
-}
-
 // The following code runs in magiskd
 
-static vector<int> get_module_fds(bool is_64_bit, modules_t module_list) {
-    vector<int> fds;
-    // All fds passed to send_fds have to be valid file descriptors.
-    // To workaround this issue, send over STDOUT_FILENO as an indicator of an
-    // invalid fd as it will always be /dev/null in magiskd
-#if defined(__LP64__)
-    if (is_64_bit) {
-        std::transform(module_list.begin(), module_list.end(), std::back_inserter(fds),
-            [](const ModuleInfo &info) { return info.z64 < 0 ? STDOUT_FILENO : info.z64; });
-    } else
-#endif
-    {
-        std::transform(module_list.begin(), module_list.end(), std::back_inserter(fds),
-            [](const ModuleInfo &info) { return info.z32 < 0 ? STDOUT_FILENO : info.z32; });
-    }
-    return fds;
-}
-
 static pthread_mutex_t zygiskd_lock = PTHREAD_MUTEX_INITIALIZER;
 static int zygiskd_sockets[] = { -1, -1 };
 #define zygiskd_socket zygiskd_sockets[is_64_bit]
 
-static void connect_companion(int client, modules_t module_list) {
+void MagiskD::connect_zygiskd(int client) const noexcept {
     mutex_guard g(zygiskd_lock);
 
     bool is_64_bit = read_int(client);
@@ -112,7 +67,7 @@ static void connect_companion(int client, modules_t module_list) {
             exit(-1);
         }
         close(fds[1]);
-        vector<int> module_fds = get_module_fds(is_64_bit, module_list);
+        rust::Vec<int> module_fds = get_module_fds(is_64_bit);
         send_fds(zygiskd_socket, module_fds.data(), module_fds.size());
         // Wait for ack
         if (read_int(zygiskd_socket) != 0) {
@@ -123,89 +78,6 @@ static void connect_companion(int client, modules_t module_list) {
     send_fd(zygiskd_socket, client);
 }
 
-static void get_process_info(int client, const sock_cred *cred, modules_t module_list) {
-    int uid = read_int(client);
-    string process = read_string(client);
-    int arch = read_int(client);
-    auto &daemon = MagiskD();
-
-    uint32_t flags = 0;
-
-    if (is_deny_target(uid, process)) {
-        flags |= PROCESS_ON_DENYLIST;
-    }
-    if (daemon.get_manager(to_user_id(uid), nullptr, false) == uid) {
-        flags |= PROCESS_IS_MAGISK_APP;
-    }
-    if (denylist_enforced) {
-        flags |= DENYLIST_ENFORCING;
-    }
-    if (daemon.uid_granted_root(uid)) {
-        flags |= PROCESS_GRANTED_ROOT;
-    }
-
-    xwrite(client, &flags, sizeof(flags));
-
-    if (should_load_modules(flags)) {
-        vector<int> fds = get_module_fds(arch, module_list);
-        send_fds(client, fds.data(), fds.size());
-    }
-
-    if (uid != 1000 || process != "system_server")
-        return;
-
-    // Collect module status from system_server
-    int slots = read_int(client);
-    dynamic_bitset bits;
-    for (int i = 0; i < slots; ++i) {
-        dynamic_bitset::slot_type l = 0;
-        xxread(client, &l, sizeof(l));
-        bits.emplace_back(l);
-    }
-    for (int id = 0; id < module_list.size(); ++id) {
-        if (!as_const(bits)[id]) {
-            // Either not a zygisk module, or incompatible
-            char buf[4096];
-            ssprintf(buf, sizeof(buf), MODULEROOT "/%s/zygisk",
-                module_list[id].name.data());
-            if (int dirfd = open(buf, O_RDONLY | O_CLOEXEC); dirfd >= 0) {
-                close(xopenat(dirfd, "unloaded", O_CREAT | O_RDONLY, 0644));
-                close(dirfd);
-            }
-        }
-    }
-}
-
-static void get_moddir(int client, modules_t module_list) {
-    int id = read_int(client);
-    char buf[4096];
-    auto &m = module_list[id];
-    ssprintf(buf, sizeof(buf), MODULEROOT "/%.*s", (int) m.name.size(), m.name.data());
-    int dfd = xopen(buf, O_RDONLY | O_CLOEXEC);
-    send_fd(client, dfd);
-    close(dfd);
-}
-
-void zygisk_handler(int client, const sock_cred *cred) {
-    auto &module_list = MagiskD().module_list();
-    int code = read_int(client);
-    switch (code) {
-    case ZygiskRequest::GET_INFO:
-        get_process_info(client, cred, module_list);
-        break;
-    case ZygiskRequest::CONNECT_COMPANION:
-        connect_companion(client, module_list);
-        break;
-    case ZygiskRequest::GET_MODDIR:
-        get_moddir(client, module_list);
-        break;
-    default:
-        // Unknown code
-        break;
-    }
-    close(client);
-}
-
 void reset_zygisk(bool restore) {
     if (!zygisk_enabled) return;
     static atomic_uint zygote_start_count{1};

native/src/core/zygisk/mod.rs

@@ -0,0 +1,3 @@
+mod daemon;
+
+pub use daemon::zygisk_should_load_module;

native/src/core/zygisk/module.cpp

@@ -77,7 +77,7 @@ bool ZygiskModule::valid() const {
 }
 
 int ZygiskModule::connectCompanion() const {
-    if (int fd = zygisk_request(ZygiskRequest::CONNECT_COMPANION); fd >= 0) {
+    if (int fd = zygisk_request(+ZygiskRequest::ConnectCompanion); fd >= 0) {
 #ifdef __LP64__
         write_int(fd, 1);
 #else
@@ -90,11 +90,9 @@ int ZygiskModule::connectCompanion() const {
 }
 
 int ZygiskModule::getModuleDir() const {
-    if (int fd = zygisk_request(ZygiskRequest::GET_MODDIR); fd >= 0) {
+    if (owned_fd fd = zygisk_request(+ZygiskRequest::GetModDir); fd >= 0) {
         write_int(fd, id);
-        int dfd = recv_fd(fd);
-        close(fd);
-        return dfd;
+        return recv_fd(fd);
     }
     return -1;
 }
@@ -210,6 +208,24 @@ bool ZygiskContext::plt_hook_commit() {
 
 // -----------------------------------------------------------------
 
+int ZygiskContext::get_module_info(int uid, std::vector<int> &fds) {
+    if (int fd = zygisk_request(+ZygiskRequest::GetInfo); fd >= 0) {
+        write_int(fd, uid);
+        write_string(fd, process);
+#ifdef __LP64__
+        write_int(fd, 1);
+#else
+        write_int(fd, 0);
+#endif
+        xxread(fd, &info_flags, sizeof(info_flags));
+        if (zygisk_should_load_module(info_flags)) {
+            fds = recv_fds(fd);
+        }
+        return fd;
+    }
+    return -1;
+}
+
 void ZygiskContext::sanitize_fds() {
     zygisk_close_logd();
 
@@ -315,16 +331,17 @@ void ZygiskContext::fork_post() {
     sigmask(SIG_UNBLOCK, SIGCHLD);
 }
 
-void ZygiskContext::run_modules_pre(const vector<int> &fds) {
+void ZygiskContext::run_modules_pre(vector<int> &fds) {
     for (int i = 0; i < fds.size(); ++i) {
+        owned_fd fd = fds[i];
         struct stat s{};
-        if (fstat(fds[i], &s) != 0 || !S_ISREG(s.st_mode)) {
-            close(fds[i]);
+        if (fstat(fd, &s) != 0 || !S_ISREG(s.st_mode)) {
+            fds[i] = -1;
             continue;
         }
         android_dlextinfo info {
             .flags = ANDROID_DLEXT_USE_LIBRARY_FD,
-            .library_fd = fds[i],
+            .library_fd = fd,
         };
         if (void *h = android_dlopen_ext("/jit-cache", RTLD_LAZY, &info)) {
             if (void *e = dlsym(h, "zygisk_module_entry")) {
@@ -332,8 +349,8 @@ void ZygiskContext::run_modules_pre(const vector<int> &fds) {
             }
         } else if (flags & SERVER_FORK_AND_SPECIALIZE) {
             ZLOGW("Failed to dlopen zygisk module: %s\n", dlerror());
+            fds[i] = -1;
         }
-        close(fds[i]);
     }
 
     for (auto it = modules.begin(); it != modules.end();) {
@@ -370,19 +387,18 @@ void ZygiskContext::app_specialize_pre() {
     flags |= APP_SPECIALIZE;
 
     vector<int> module_fds;
-    int fd = remote_get_info(args.app->uid, process, &info_flags, module_fds);
+    owned_fd fd = get_module_info(args.app->uid, module_fds);
     if ((info_flags & UNMOUNT_MASK) == UNMOUNT_MASK) {
         ZLOGI("[%s] is on the denylist\n", process);
         flags |= DO_REVERT_UNMOUNT;
     } else if (fd >= 0) {
         run_modules_pre(module_fds);
     }
-    close(fd);
 }
 
 void ZygiskContext::app_specialize_post() {
     run_modules_post();
-    if (info_flags & PROCESS_IS_MAGISK_APP) {
+    if (info_flags & +ZygiskStateFlags::ProcessIsMagiskApp) {
         setenv("ZYGISK_ENABLED", "1", 1);
     }
 
@@ -392,24 +408,21 @@ void ZygiskContext::app_specialize_post() {
 
 void ZygiskContext::server_specialize_pre() {
     vector<int> module_fds;
-    int fd = remote_get_info(1000, "system_server", &info_flags, module_fds);
-    if (fd >= 0) {
+    if (owned_fd fd = get_module_info(1000, module_fds); fd >= 0) {
         if (module_fds.empty()) {
             write_int(fd, 0);
         } else {
             run_modules_pre(module_fds);
 
-            // Send the bitset of module status back to magiskd from system_server
-            dynamic_bitset bits;
-            for (const auto &m : modules)
-                bits[m.getId()] = true;
-            write_int(fd, static_cast<int>(bits.slots()));
-            for (int i = 0; i < bits.slots(); ++i) {
-                auto l = bits.get_slot(i);
-                xwrite(fd, &l, sizeof(l));
+            // Find all failed module ids and send it back to magiskd
+            vector<int> failed_ids;
+            for (int i = 0; i < module_fds.size(); ++i) {
+                if (module_fds[i] < 0) {
+                    failed_ids.push_back(i);
+                }
             }
+            write_vector(fd, failed_ids);
         }
-        close(fd);
     }
 }
 
@@ -435,6 +448,7 @@ void ZygiskContext::nativeSpecializeAppProcess_post() {
 void ZygiskContext::nativeForkSystemServer_pre() {
     ZLOGV("pre  forkSystemServer\n");
     flags |= SERVER_FORK_AND_SPECIALIZE;
+    process = "system_server";
 
     fork_pre();
     if (is_child()) {

native/src/core/zygisk/module.hpp

@@ -122,15 +122,13 @@ struct module_abi_v1 {
     void (*postServerSpecialize)(void *, const void *);
 };
 
+// Assert the flag values to be the same as the public API
+static_assert(+ZygiskStateFlags::ProcessGrantedRoot == zygisk::StateFlag::PROCESS_GRANTED_ROOT);
+static_assert(+ZygiskStateFlags::ProcessOnDenyList == zygisk::StateFlag::PROCESS_ON_DENYLIST);
+
 enum : uint32_t {
-    PROCESS_GRANTED_ROOT = zygisk::StateFlag::PROCESS_GRANTED_ROOT,
-    PROCESS_ON_DENYLIST = zygisk::StateFlag::PROCESS_ON_DENYLIST,
-
-    DENYLIST_ENFORCING = (1u << 30),
-    PROCESS_IS_MAGISK_APP = (1u << 31),
-
-    UNMOUNT_MASK = (PROCESS_ON_DENYLIST | DENYLIST_ENFORCING),
-    PRIVATE_MASK = (DENYLIST_ENFORCING | PROCESS_IS_MAGISK_APP)
+    UNMOUNT_MASK = (+ZygiskStateFlags::ProcessOnDenyList | +ZygiskStateFlags::DenyListEnforced),
+    PRIVATE_MASK = (+ZygiskStateFlags::DenyListEnforced | +ZygiskStateFlags::ProcessIsMagiskApp)
 };
 
 struct api_abi_base {
@@ -188,7 +186,6 @@ struct ZygiskModule {
     static uint32_t getFlags();
     void tryUnload() const;
     void clearApi() { memset(&api, 0, sizeof(api)); }
-    int getId() const { return id; }
 
     ZygiskModule(int id, void *handle, void *entry);
 
@@ -264,7 +261,7 @@ struct ZygiskContext {
     ZygiskContext(JNIEnv *env, void *args);
     ~ZygiskContext();
 
-    void run_modules_pre(const std::vector<int> &fds);
+    void run_modules_pre(std::vector<int> &fds);
     void run_modules_post();
     DCL_PRE_POST(fork)
     DCL_PRE_POST(app_specialize)
@@ -273,6 +270,7 @@ struct ZygiskContext {
     DCL_PRE_POST(nativeSpecializeAppProcess)
     DCL_PRE_POST(nativeForkSystemServer)
 
+    int get_module_info(int uid, std::vector<int> &fds);
     void sanitize_fds();
     bool exempt_fd(int fd);
     bool can_exempt_fd() const;

native/src/core/zygisk/zygisk.hpp

@@ -6,15 +6,6 @@
 #include <vector>
 #include <core.hpp>
 
-namespace ZygiskRequest {
-enum : int {
-    GET_INFO,
-    CONNECT_COMPANION,
-    GET_MODDIR,
-    END
-};
-}
-
 #if defined(__LP64__)
 #define ZLOGD(...) LOGD("zygisk64: " __VA_ARGS__)
 #define ZLOGE(...) LOGE("zygisk64: " __VA_ARGS__)
@@ -34,8 +25,6 @@ enum : int {
 void hook_entry();
 void hookJniNativeMethods(JNIEnv *env, const char *clz, JNINativeMethod *methods, int numMethods);
 
-int remote_get_info(int uid, const char *process, uint32_t *flags, std::vector<int> &fds);
-
 inline int zygisk_request(int req) {
     int fd = connect_daemon(+RequestCode::ZYGISK);
     if (fd < 0) return fd;