run companion in standalone process

loader/src/ptracer/main.cpp

@@ -18,15 +18,8 @@ int main(int argc, char **argv) {
         return 0;
     } else if (argc >= 3 && argv[1] == "trace"sv) {
         if (argc >= 4 && argv[3] == "--restart"sv) {
-            constexpr auto companion = "./bin/zygisk-cp" LP_SELECT("32", "64");
             zygiskd::Init(getenv(MAGIC_PATH_ENV));
             zygiskd::ZygoteRestart();
-            if (fork_dont_care() == 0) {
-                LOGI("creating new zygisk companion");
-                execl(companion, basename(companion), nullptr);
-                PLOGE("failed to exec zygisk companion");
-                exit(1);
-            }
         }
         auto pid = strtol(argv[2], 0, 0);
         if (!trace_zygote(pid)) {

module/src/service.sh

@@ -24,3 +24,7 @@ if [ "$(which magisk)" ]; then
     fi
   done
 fi
+
+[ "$DEBUG" = true ] && export RUST_BACKTRACE=1
+unshare -m sh -c "bin/zygisk-cp64 &"
+unshare -m sh -c "bin/zygisk-cp32 &"

zygiskd/.cargo/config.toml

@@ -1,2 +1,3 @@
 [build]
 target-dir = "build/intermediates/rust"
+target = "aarch64-linux-android"

zygiskd/Cargo.toml

@@ -23,7 +23,14 @@ proc-maps = "0.3"
 rustix = { version = "0.38", features = [ "fs", "process", "mount", "net", "thread" ] }
 tokio = { version = "1.28", features = ["full"] }
 
+[profile.dev]
+# strip = true
+# split-debuginfo = "packed"
+panic = "abort"
+
 [profile.release]
 strip = true
+split-debuginfo = "packed"
+panic = "abort"
 opt-level = "z"
 lto = true

zygiskd/src/companion.rs

@@ -0,0 +1,72 @@
+use std::ffi::c_void;
+use std::os::fd::{AsRawFd, FromRawFd, RawFd};
+use std::os::unix::net::UnixStream;
+use std::thread;
+use anyhow::Result;
+use passfd::FdPassingExt;
+use rustix::fs::fstat;
+use tokio::io::AsyncWriteExt;
+use crate::utils::{check_unix_socket, UnixStreamExt};
+use crate::dl;
+
+type ZygiskCompanionEntryFn = unsafe extern "C" fn(i32);
+
+pub fn entry(fd: i32) {
+    log::info!("companion entry fd={}", fd);
+    let mut stream = unsafe { UnixStream::from_raw_fd(fd) };
+    let name = stream.read_string().expect("read name");
+    let library = stream.recv_fd().expect("receive library fd");
+    let entry = load_module(library).expect("load module");
+    unsafe { libc::close(library) };
+
+    let entry = match entry {
+        Some(entry) => {
+            log::debug!("Companion process created for `{name}`");
+            stream.write_u8(1).expect("reply 1");
+            entry
+        }
+        None => {
+            log::debug!("No companion entry for `{name}`");
+            stream.write_u8(0).expect("reply 0");
+            return ();
+        }
+    };
+
+    loop {
+        if !check_unix_socket(&stream, true) {
+            log::info!("Something bad happened in zygiskd, terminate companion");
+            std::process::exit(0);
+        }
+        let fd = stream.recv_fd().expect("recv fd");
+        log::trace!("New companion request from module `{name}` fd=`{fd}`");
+        let mut stream = unsafe { UnixStream::from_raw_fd(fd) };
+        stream.write_u8(1).expect("reply success");
+        thread::spawn(move || {
+            let st0 = fstat(&stream).expect("failed to stat stream");
+            unsafe { entry(stream.as_raw_fd()); }
+            // Only close client if it is the same file so we don't
+            // accidentally close a re-used file descriptor.
+            // This check is required because the module companion
+            // handler could've closed the file descriptor already.
+            if let Ok(st1) = fstat(&stream) {
+                if st0.st_dev != st1.st_dev || st0.st_ino != st1.st_ino {
+                    std::mem::forget(stream);
+                }
+            }
+        });
+    }
+}
+
+fn load_module(fd: RawFd) -> Result<Option<ZygiskCompanionEntryFn>> {
+    unsafe {
+        let path = format!("/proc/self/fd/{fd}");
+        let handle = dl::dlopen(&path, libc::RTLD_NOW)?;
+        let symbol = std::ffi::CString::new("zygisk_companion_entry")?;
+        let entry = libc::dlsym(handle, symbol.as_ptr());
+        if entry.is_null() {
+            return Ok(None);
+        }
+        let fnptr = std::mem::transmute::<*mut c_void, ZygiskCompanionEntryFn>(entry);
+        Ok(Some(fnptr))
+    }
+}

zygiskd/src/main.rs

@@ -3,6 +3,7 @@ mod dl;
 mod root_impl;
 mod utils;
 mod zygiskd;
+mod companion;
 
 use std::future::Future;
 use anyhow::Result;
@@ -15,19 +16,17 @@ fn init_android_logger(tag: &str) {
     );
 }
 
-fn async_start<F: Future>(future: F) -> F::Output {
-    let async_runtime = tokio::runtime::Runtime::new().unwrap();
-    async_runtime.block_on(future)
-}
-
-fn start(name: &str) -> Result<()> {
-    utils::switch_mount_namespace(1)?;
-    root_impl::setup();
-    match name.trim_start_matches("zygisk-") {
-        lp_select!("cp32", "cp64") => zygiskd::main()?,
-        _ => println!("Available command: cp[32|64]"),
+fn start() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() == 3 && args[1] == "companion" {
+        let fd: i32 = args[2].parse().unwrap();
+        companion::entry(fd);
+        return;
     }
-    Ok(())
+
+    utils::switch_mount_namespace(1).expect("switch mnt ns");
+    root_impl::setup();
+    zygiskd::main().expect("zygiskd main");
 }
 
 fn main() {
@@ -35,7 +34,5 @@ fn main() {
     let nice_name = process.split('/').last().unwrap();
     init_android_logger(nice_name);
 
-    if let Err(e) = start(nice_name) {
-        log::error!("Crashed: {}\n{}", e, e.backtrace());
-    }
+    start();
 }

zygiskd/src/utils.rs

@@ -1,7 +1,7 @@
 use anyhow::Result;
 use std::{fs, io::{Read, Write}, os::unix::net::UnixStream};
 use std::ffi::{c_char, c_void, CStr, CString};
-use std::os::fd::AsFd;
+use std::os::fd::{AsFd, AsRawFd};
 use std::os::unix::net::UnixListener;
 use std::process::Command;
 use std::sync::OnceLock;
@@ -195,6 +195,22 @@ pub fn unix_listener_from_path(path: &str) -> Result<UnixListener> {
     Ok(UnixListener::from(socket))
 }
 
+pub fn check_unix_socket(stream: &UnixStream, block: bool) -> bool {
+    unsafe {
+        let mut pfd = libc::pollfd {
+            fd: stream.as_raw_fd(),
+            events: libc::POLLIN,
+            revents: 0,
+        };
+        let timeout = if block { -1 } else { 0 };
+        libc::poll(&mut pfd, 1, timeout);
+        if pfd.revents != 0 && pfd.revents & libc::POLLIN == 0 {
+            return false;
+        }
+    }
+    return true;
+}
+
 extern "C" {
     fn __android_log_print(prio: i32, tag: *const c_char, fmt: *const c_char, ...) -> i32;
     fn __system_property_get(name: *const c_char, value: *mut c_char) -> u32;

zygiskd/src/zygiskd.rs

@@ -1,29 +1,29 @@
-use std::ffi::c_void;
 use crate::constants::{DaemonSocketAction, ProcessFlags};
-use crate::utils::UnixStreamExt;
+use crate::utils::{check_unix_socket, UnixStreamExt};
 use crate::{constants, dl, lp_select, root_impl, utils};
 use anyhow::{bail, Result};
 use passfd::FdPassingExt;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use std::thread;
 use std::fs;
-use std::os::fd::OwnedFd;
+use std::io::Error;
+use std::os::fd::{OwnedFd, RawFd};
 use std::os::unix::{
     net::{UnixListener, UnixStream},
     prelude::AsRawFd,
 };
 use std::path::PathBuf;
-use std::process::exit;
+use std::process::{Command, exit};
 use log::info;
-use rustix::fs::fstat;
-use rustix::process::{set_parent_process_death_signal, Signal};
+use std::os::unix::process::CommandExt;
+use bitflags::Flags;
 
 type ZygiskCompanionEntryFn = unsafe extern "C" fn(i32);
 
 struct Module {
     name: String,
     lib_fd: OwnedFd,
-    entry: Option<ZygiskCompanionEntryFn>,
+    companion: Mutex<Option<Option<UnixStream>>>,
 }
 
 struct Context {
@@ -32,7 +32,6 @@ struct Context {
 
 pub fn main() -> Result<()> {
     log::info!("Start zygisk companion");
-    set_parent_process_death_signal(Some(Signal::Kill))?;
 
     let arch = get_arch()?;
     log::debug!("Daemon architecture: {arch}");
@@ -60,8 +59,11 @@ pub fn main() -> Result<()> {
                 // Do nothing
             }
             DaemonSocketAction::ZygoteRestart => {
-                info!("Zygote restarted, exit");
-                exit(0);
+                info!("Zygote restarted, clean up companions");
+                for module in &context.modules {
+                    let mut companion = module.companion.lock().unwrap();
+                    companion.take();
+                }
             }
             _ => {
                 thread::spawn(move || {
@@ -112,8 +114,8 @@ fn load_modules(arch: &str) -> Result<Vec<Module>> {
                 continue;
             }
         };
-        let entry = resolve_module(&so_path.to_string_lossy())?;
-        let module = Module { name, lib_fd, entry };
+        let companion = Mutex::new(None);
+        let module = Module { name, lib_fd, companion };
         modules.push(module);
     }
 
@@ -153,17 +155,43 @@ fn create_daemon_socket() -> Result<UnixListener> {
     Ok(listener)
 }
 
-fn resolve_module(path: &str) -> Result<Option<ZygiskCompanionEntryFn>> {
+fn spawn_companion(name: &str, lib_fd: RawFd) -> Result<Option<UnixStream>> {
+    let (mut daemon, companion) = UnixStream::pair()?;
+
+    let process = std::env::args().next().unwrap();
+    let nice_name = process.split('/').last().unwrap();
+
     unsafe {
-        let handle = dl::dlopen(path, libc::RTLD_NOW)?;
-        let symbol = std::ffi::CString::new("zygisk_companion_entry")?;
-        let entry = libc::dlsym(handle, symbol.as_ptr());
-        if entry.is_null() {
-            return Ok(None);
+        let pid = libc::fork();
+        if pid < 0 {
+            bail!(Error::last_os_error());
+        } else if pid > 0 {
+            drop(companion);
+            let mut status: libc::c_int = 0;
+            libc::waitpid(pid, &mut status, 0);
+            if libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0 {
+                daemon.write_string(name)?;
+                daemon.send_fd(lib_fd)?;
+                return match daemon.read_u8()? {
+                    0 => Ok(None),
+                    1 => Ok(Some(daemon)),
+                    _ => bail!("Invalid companion response"),
+                }
+            } else {
+                bail!("exited with status {}", status);
+            }
+        } else {
+            // Remove FD_CLOEXEC flag
+            unsafe { libc::fcntl(companion.as_raw_fd() as libc::c_int, libc::F_SETFD, 0i32); };
         }
-        let fnptr = std::mem::transmute::<*mut c_void, ZygiskCompanionEntryFn>(entry);
-        Ok(Some(fnptr))
     }
+
+    Command::new(&process)
+        .arg0(format!("{}-{}", nice_name, name))
+        .arg("companion")
+        .arg(format!("{}", companion.as_raw_fd()))
+        .spawn()?;
+    exit(0)
 }
 
 fn handle_daemon_action(action: DaemonSocketAction, mut stream: UnixStream, context: &Context) -> Result<()> {
@@ -207,26 +235,41 @@ fn handle_daemon_action(action: DaemonSocketAction, mut stream: UnixStream, cont
         DaemonSocketAction::RequestCompanionSocket => {
             let index = stream.read_usize()?;
             let module = &context.modules[index];
-            match module.entry {
-                None => {
-                    stream.write_u8(0)?;
-                    return Ok(());
-                }
-                Some(companion) => {
-                    stream.write_u8(1)?;
-                    let st0 = fstat(&stream)?;
-                    unsafe { companion(stream.as_raw_fd()); }
-                    // Only close client if it is the same file so we don't
-                    // accidentally close a re-used file descriptor.
-                    // This check is required because the module companion
-                    // handler could've closed the file descriptor already.
-                    if let Ok(st1) = fstat(&stream) {
-                        if st0.st_dev != st1.st_dev || st0.st_ino != st1.st_ino {
-                            std::mem::forget(stream);
-                        }
-                    }
+            let mut companion = module.companion.lock().unwrap();
+            if let Some(Some(sock)) = companion.as_ref() {
+                if !check_unix_socket(sock, false) {
+                    log::error!("Poll companion for module `{}` crashed", module.name);
+                    companion.take();
                 }
             }
+            if companion.is_none() {
+                match spawn_companion(&module.name, module.lib_fd.as_raw_fd()) {
+                    Ok(c) => {
+                        if c.is_some() {
+                            log::trace!("  Spawned companion for `{}`", module.name);
+                        } else {
+                            log::trace!("  No companion spawned for `{}` because it has not entry", module.name);
+                        }
+                        *companion = Some(c);
+                    },
+                    Err(e) => {
+                        log::warn!("  Failed to spawn companion for `{}`: {}", module.name, e);
+                    }
+                };
+            }
+            match companion.as_ref() {
+                Some(Some(sock)) => {
+                    if let Err(_) = sock.send_fd(stream.as_raw_fd()) {
+                        log::error!("Failed to send companion fd socket of module `{}`", module.name);
+                        stream.write_u8(0)?;
+                    }
+                    // Ok: Send by companion
+                }
+                _ => {
+                    stream.write_u8(0)?;
+                }
+            }
+            log::info!("companion done");
         }
         DaemonSocketAction::GetModuleDir => {
             let index = stream.read_usize()?;