Migrate connect_daemon to Rust

native/src/Android.mk

@@ -17,9 +17,9 @@ LOCAL_STATIC_LIBRARIES := \
 LOCAL_SRC_FILES := \
     core/applets.cpp \
     core/magisk.cpp \
-    core/daemon.cpp \
     core/scripting.cpp \
     core/sqlite.cpp \
+    core/utils.cpp \
     core/core-rs.cpp \
     core/resetprop/sys.cpp \
     core/su/su.cpp \

native/src/base/base.cpp

@@ -158,10 +158,10 @@ void init_argv0(int argc, char **argv) {
     name_len = (argv[argc - 1] - argv[0]) + strlen(argv[argc - 1]) + 1;
 }
 
-void set_nice_name(const char *name) {
+void set_nice_name(Utf8CStr name) {
     memset(argv0, 0, name_len);
-    strscpy(argv0, name, name_len);
-    prctl(PR_SET_NAME, name);
+    strscpy(argv0, name.c_str(), name_len);
+    prctl(PR_SET_NAME, name.c_str());
 }
 
 template<typename T, int base>

native/src/base/lib.rs

@@ -10,7 +10,7 @@ pub use cstr::{
 };
 use cxx_extern::*;
 pub use dir::*;
-pub use ffi::{Utf8CStrRef, fork_dont_care};
+pub use ffi::{Utf8CStrRef, fork_dont_care, set_nice_name};
 pub use files::*;
 pub use logging::*;
 pub use misc::*;
@@ -46,6 +46,7 @@ mod ffi {
 
         fn mut_u8_patch(buf: &mut [u8], from: &[u8], to: &[u8]) -> Vec<usize>;
         fn fork_dont_care() -> i32;
+        fn set_nice_name(name: Utf8CStrRef);
 
         type FnBoolStrStr;
         fn call(self: &FnBoolStrStr, key: &str, value: &str) -> bool;

native/src/base/misc.hpp

@@ -18,6 +18,8 @@ clazz(const clazz&) = delete;  \
 clazz(clazz &&o) { swap(o); }  \
 clazz& operator=(clazz &&o) { swap(o); return *this; }
 
+struct Utf8CStr;
+
 class mutex_guard {
     DISALLOW_COPY_AND_MOVE(mutex_guard)
 public:
@@ -146,7 +148,7 @@ static inline std::string rtrim(std::string &&s) {
 int fork_dont_care();
 int fork_no_orphan();
 void init_argv0(int argc, char **argv);
-void set_nice_name(const char *name);
+void set_nice_name(Utf8CStr name);
 int switch_mnt_ns(int pid);
 std::string &replace_all(std::string &str, std::string_view from, std::string_view to);
 std::vector<std::string> split(std::string_view s, std::string_view delims);
@@ -245,4 +247,4 @@ struct FnBoolStr : public CxxFnBoolStr {
     bool call(Utf8CStr s) const {
         return operator()(s);
     }
-};
+};

native/src/core/daemon.rs

@@ -19,8 +19,9 @@ use crate::zygisk::ZygiskState;
 use base::const_format::concatcp;
 use base::{
     AtomicArc, BufReadExt, FileAttr, FsPathBuilder, ReadExt, ResultExt, Utf8CStr, Utf8CStrBuf,
-    WriteExt, cstr, info, libc,
+    WriteExt, cstr, error, fork_dont_care, info, libc, set_nice_name,
 };
+use nix::unistd::getuid;
 use nix::{
     fcntl::OFlag,
     mount::MsFlags,
@@ -30,11 +31,12 @@ use nix::{
 use num_traits::AsPrimitive;
 use std::fmt::Write as _;
 use std::io::{BufReader, Write};
-use std::os::fd::{AsFd, AsRawFd, IntoRawFd};
+use std::os::fd::{AsFd, AsRawFd, IntoRawFd, RawFd};
 use std::os::unix::net::{UCred, UnixListener, UnixStream};
 use std::process::{Command, exit};
 use std::sync::atomic::AtomicBool;
 use std::sync::{Mutex, OnceLock};
+use std::time::Duration;
 
 // Global magiskd singleton
 pub static MAGISKD: OnceLock<MagiskD> = OnceLock::new();
@@ -150,7 +152,7 @@ impl MagiskD {
         }
     }
 
-    pub fn reboot(&self) {
+    fn reboot(&self) {
         if self.is_recovery {
             Command::new("/system/bin/reboot").arg("recovery").status()
         } else {
@@ -260,7 +262,22 @@ impl MagiskD {
     }
 }
 
-pub fn daemon_entry() {
+fn switch_cgroup(cgroup: &str, pid: i32) {
+    let mut buf = cstr::buf::new::<64>()
+        .join_path(cgroup)
+        .join_path("cgroup.procs");
+    if !buf.exists() {
+        return;
+    }
+    if let Ok(mut file) = buf.open(OFlag::O_WRONLY | OFlag::O_APPEND | OFlag::O_CLOEXEC) {
+        buf.clear();
+        write!(buf, "{pid}").ok();
+        file.write_all(buf.as_bytes()).log_ok();
+    }
+}
+
+fn daemon_entry() {
+    set_nice_name(cstr!("magiskd"));
     android_logging();
 
     // Block all signals
@@ -405,16 +422,63 @@ pub fn daemon_entry() {
     }
 }
 
-fn switch_cgroup(cgroup: &str, pid: i32) {
-    let mut buf = cstr::buf::new::<64>()
-        .join_path(cgroup)
-        .join_path("cgroup.procs");
-    if !buf.exists() {
-        return;
+pub fn connect_daemon(code: RequestCode, create: bool) -> RawFd {
+    let sock_path = cstr::buf::new::<64>()
+        .join_path(get_magisk_tmp())
+        .join_path(MAIN_SOCKET);
+
+    fn send_request(code: RequestCode, mut socket: UnixStream) -> RawFd {
+        socket.write_pod(&code.repr).log_ok();
+        let mut res = -1;
+        socket.read_pod(&mut res).log_ok();
+        let res = RespondCode { repr: res };
+        match res {
+            RespondCode::OK => socket.into_raw_fd(),
+            RespondCode::ROOT_REQUIRED => {
+                error!("Root is required for this operation");
+                -1
+            }
+            RespondCode::ACCESS_DENIED => {
+                error!("Accessed denied");
+                -1
+            }
+            _ => {
+                error!("Daemon error");
+                -1
+            }
+        }
     }
-    if let Ok(mut file) = buf.open(OFlag::O_WRONLY | OFlag::O_APPEND | OFlag::O_CLOEXEC) {
-        buf.clear();
-        write!(buf, "{pid}").ok();
-        file.write_all(buf.as_bytes()).log_ok();
+
+    match UnixStream::connect(&sock_path) {
+        Ok(socket) => send_request(code, socket),
+        Err(e) => {
+            if !create || !getuid().is_root() {
+                error!("Cannot connect to daemon: {e}");
+                return -1;
+            }
+
+            let mut buf = cstr::buf::new::<64>();
+            if cstr!("/proc/self/exe").read_link(&mut buf).is_err()
+                || !buf.starts_with(get_magisk_tmp().as_str())
+            {
+                error!("Start daemon on magisk tmpfs");
+                return -1;
+            }
+
+            // Fork a process and run the daemon
+            if fork_dont_care() == 0 {
+                daemon_entry();
+                exit(0);
+            }
+
+            // In the client, we keep retry and connect to the socket
+            loop {
+                if let Ok(socket) = UnixStream::connect(&sock_path) {
+                    return send_request(code, socket);
+                } else {
+                    std::thread::sleep(Duration::from_millis(100));
+                }
+            }
+        }
     }
 }

native/src/core/include/core.hpp

@@ -19,6 +19,10 @@
 #define SDK_INT      (MagiskD::Get().sdk_int())
 #define APP_DATA_DIR (SDK_INT >= 24 ? "/data/user_de" : "/data/user")
 
+inline int connect_daemon(RequestCode req) {
+    return connect_daemon(req, false);
+}
+
 // Multi-call entrypoints
 int magisk_main(int argc, char *argv[]);
 int su_client_main(int argc, char *argv[]);
@@ -26,12 +30,10 @@ int zygisk_main(int argc, char *argv[]);
 
 struct ModuleInfo;
 
-// Daemon
-int connect_daemon(RequestCode req, bool create = false);
+// Utils
 const char *get_magisk_tmp();
 void unlock_blocks();
 bool check_key_combo();
-
 template<typename T> requires(std::is_trivially_copyable_v<T>)
 T read_any(int fd) {
     T val;
@@ -39,25 +41,21 @@ T read_any(int fd) {
         return -1;
     return val;
 }
-
 template<typename T> requires(std::is_trivially_copyable_v<T>)
 void write_any(int fd, T val) {
     if (fd < 0) return;
     xwrite(fd, &val, sizeof(val));
 }
-
-static inline int read_int(int fd) { return read_any<int>(fd); }
-static inline void write_int(int fd, int val) { write_any(fd, val); }
+inline int read_int(int fd) { return read_any<int>(fd); }
+inline void write_int(int fd, int val) { write_any(fd, val); }
 std::string read_string(int fd);
 bool read_string(int fd, std::string &str);
 void write_string(int fd, std::string_view str);
-
 template<typename T> requires(std::is_trivially_copyable_v<T>)
 void write_vector(int fd, const std::vector<T> &vec) {
     write_int(fd, vec.size());
     xwrite(fd, vec.data(), vec.size() * sizeof(T));
 }
-
 template<typename T> requires(std::is_trivially_copyable_v<T>)
 bool read_vector(int fd, std::vector<T> &vec) {
     int size = read_int(fd);
@@ -88,8 +86,8 @@ void update_deny_flags(int uid, rust::Str process, uint32_t &flags);
 void exec_root_shell(int client, int pid, SuRequest &req, MntNsMode mode);
 
 // Rust bindings
-static inline Utf8CStr get_magisk_tmp_rs() { return get_magisk_tmp(); }
-static inline rust::String resolve_preinit_dir_rs(Utf8CStr base_dir) {
+inline Utf8CStr get_magisk_tmp_rs() { return get_magisk_tmp(); }
+inline rust::String resolve_preinit_dir_rs(Utf8CStr base_dir) {
     return resolve_preinit_dir(base_dir.c_str());
 }
-static inline void exec_script_rs(Utf8CStr script) { exec_script(script.c_str()); }
+inline void exec_script_rs(Utf8CStr script) { exec_script(script.c_str()); }

native/src/core/lib.rs

@@ -9,7 +9,7 @@
 
 use crate::ffi::SuRequest;
 use crate::socket::Encodable;
-use daemon::{MagiskD, daemon_entry};
+use daemon::{MagiskD, connect_daemon};
 use derive::Decodable;
 use logging::{android_logging, zygisk_close_logd, zygisk_get_logd, zygisk_logging};
 use mount::{find_preinit_device, revert_unmount};
@@ -198,7 +198,7 @@ pub mod ffi {
         fn get_prop(name: Utf8CStrRef) -> String;
         unsafe fn resetprop_main(argc: i32, argv: *mut *mut c_char) -> i32;
 
-        fn daemon_entry();
+        fn connect_daemon(code: RequestCode, create: bool) -> i32;
     }
 
     // Default constructors

native/src/core/utils.cpp

@@ -1,6 +1,5 @@
 #include <csignal>
 #include <libgen.h>
-#include <sys/un.h>
 #include <sys/mount.h>
 #include <sys/sysmacros.h>
 #include <linux/input.h>
@@ -45,60 +44,6 @@ const char *get_magisk_tmp() {
     return path;
 }
 
-int connect_daemon(RequestCode req, bool create) {
-    int fd = xsocket(AF_LOCAL, SOCK_STREAM | SOCK_CLOEXEC, 0);
-    sockaddr_un addr = {.sun_family = AF_LOCAL};
-    const char *tmp = get_magisk_tmp();
-    ssprintf(addr.sun_path, sizeof(addr.sun_path), "%s/" MAIN_SOCKET, tmp);
-    if (connect(fd, reinterpret_cast<sockaddr *>(&addr), sizeof(addr))) {
-        if (!create || getuid() != AID_ROOT) {
-            PLOGE("Cannot connect daemon at %s,", addr.sun_path);
-            close(fd);
-            return -1;
-        }
-
-        char buf[64];
-        xreadlink("/proc/self/exe", buf, sizeof(buf));
-        if (tmp[0] == '\0' || !string_view(buf).starts_with(tmp)) {
-            LOGE("Start daemon on magisk tmpfs\n");
-            close(fd);
-            return -1;
-        }
-
-        if (fork_dont_care() == 0) {
-            close(fd);
-            set_nice_name("magiskd");
-            daemon_entry();
-        }
-
-        while (connect(fd, reinterpret_cast<sockaddr *>(&addr), sizeof(addr)))
-            usleep(10000);
-    }
-    write_int(fd, +req);
-    int res = read_int(fd);
-    if (res < +RespondCode::ERROR || res >= +RespondCode::END)
-        res = +RespondCode::ERROR;
-    switch (res) {
-    case +RespondCode::OK:
-        break;
-    case +RespondCode::ERROR:
-        LOGE("Daemon error\n");
-        close(fd);
-        return -1;
-    case +RespondCode::ROOT_REQUIRED:
-        LOGE("Root is required for this operation\n");
-        close(fd);
-        return -1;
-    case +RespondCode::ACCESS_DENIED:
-        LOGE("Access denied\n");
-        close(fd);
-        return -1;
-    default:
-        __builtin_unreachable();
-    }
-    return fd;
-}
-
 void unlock_blocks() {
     int fd, dev, OFF = 0;