restart zygiskd when zygote restart

loader/src/common/daemon.cpp

@@ -9,6 +9,7 @@
 namespace zygiskd {
     static std::string zygisk_path;
     void Init(const char *path) {
+        LOGI("zygisk path set to %s", path);
         zygisk_path = path;
     }
 
@@ -25,8 +26,10 @@ namespace zygiskd {
         while (retry--) {
             int r = connect(fd, reinterpret_cast<struct sockaddr*>(&addr), socklen);
             if (r == 0) return fd;
-            PLOGE("Retrying to connect to zygiskd, sleep 1s");
+            if (retry) {
-            sleep(1);
+                PLOGE("Retrying to connect to zygiskd, sleep 1s");
+                sleep(1);
+            }
         }
 
         close(fd);
@@ -111,9 +114,15 @@ namespace zygiskd {
     void ZygoteRestart() {
         UniqueFd fd = Connect(1);
         if (fd == -1) {
-            PLOGE("Could not notify ZygoteRestart");
+            if (errno == ENOENT) {
+                LOGD("Could not notify ZygoteRestart (maybe it hasn't been created)");
+            } else {
+                PLOGE("Could not notify ZygoteRestart");
+            }
             return;
         }
-        socket_utils::write_u8(fd, (uint8_t) SocketAction::ZygoteRestart);
+        if (!socket_utils::write_u8(fd, (uint8_t) SocketAction::ZygoteRestart)) {
+            PLOGE("Failed to request ZygoteRestart");
+        }
     }
 }

loader/src/include/daemon.h

@@ -12,6 +12,7 @@
 #endif
 
 constexpr auto kCPSocketName = "/" LP_SELECT("cp32", "cp64") ".sock";
+constexpr const auto MAGIC_PATH_ENV = "MAGIC_PATH";
 
 class UniqueFd {
     using Fd = int;

loader/src/ptracer/main.cpp

@@ -17,6 +17,17 @@ int main(int argc, char **argv) {
         init_monitor();
         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)) {
             kill(pid, SIGKILL);

loader/src/ptracer/main.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+
 void init_monitor();
 bool trace_zygote(int pid);
 

loader/src/ptracer/monitor.cpp

@@ -256,7 +256,7 @@ public:
                     if (STOPPED_WITH(SIGTRAP, PTRACE_EVENT_FORK)) {
                         long child_pid;
                         ptrace(PTRACE_GETEVENTMSG, pid, 0, &child_pid);
-                        LOGI("forked %ld", child_pid);
+                        LOGV("forked %ld", child_pid);
                     } else if (STOPPED_WITH(SIGTRAP, PTRACE_EVENT_STOP) &&
                                tracing_state == STOPPING) {
                         if (ptrace(PTRACE_DETACH, 1, 0, 0) == -1)
@@ -283,7 +283,7 @@ public:
                 }
                 auto state = process.find(pid);
                 if (state == process.end()) {
-                    LOGI("new process %d attached", pid);
+                    LOGV("new process %d attached", pid);
                     process.emplace(pid);
                     ptrace(PTRACE_SETOPTIONS, pid, 0, PTRACE_O_TRACEEXEC);
                     ptrace(PTRACE_CONT, pid, 0, 0);
@@ -291,7 +291,7 @@ public:
                 } else {
                     if (STOPPED_WITH(SIGTRAP, PTRACE_EVENT_EXEC)) {
                         auto program = get_program(pid);
-                        LOGD("%d program %s", pid, program.c_str());
+                        LOGV("%d program %s", pid, program.c_str());
                         const char* tracer = nullptr;
                         do {
                             if (tracing_state != TRACING) {
@@ -327,7 +327,7 @@ public:
                                     auto p = fork_dont_care();
                                     if (p == 0) {
                                         execl(tracer, basename(tracer), "trace",
-                                              std::to_string(pid).c_str(), nullptr);
+                                              std::to_string(pid).c_str(), "--restart", nullptr);
                                         PLOGE("failed to exec, kill");
                                         kill(pid, SIGKILL);
                                         exit(1);
@@ -340,18 +340,17 @@ public:
                         } while (false);
 
                     } else {
-                        LOGD("process %d received unknown status %s", pid,
+                        LOGE("process %d received unknown status %s", pid,
                              parse_status(status).c_str());
                     }
                     process.erase(state);
                     if (WIFSTOPPED(status)) {
-                        LOGI("detach process %d", pid);
+                        LOGV("detach process %d", pid);
                         ptrace(PTRACE_DETACH, pid, 0, 0);
                     }
                 }
             }
         }
-        LOGD("sigchld handle done");
     }
 
     ~PtraceHandler() {

loader/src/ptracer/ptracer.cpp

@@ -180,7 +180,7 @@ bool trace_zygote(int pid) {
     }
     WAIT_OR_DIE
     if (STOPPED_WITH(SIGSTOP, PTRACE_EVENT_STOP)) {
-        std::string magic_path = getenv("MAGIC_PATH");
+        std::string magic_path = getenv(MAGIC_PATH_ENV);
         std::string lib_path = magic_path + "/lib" LP_SELECT("", "64") "/libzygisk.so";
         if (!inject_on_main(pid, lib_path.c_str(), magic_path.c_str())) {
             LOGE("failed to inject");

module/src/post-fs-data.sh

@@ -43,4 +43,5 @@ if [ -f $MODDIR/lib/libzygisk.so ];then
   chcon u:object_r:system_file:s0 $MAGIC_PATH/lib/libzygisk.so
 fi
 
+[ "$DEBUG" = true ] && export RUST_BACKTRACE=1
 unshare -m sh -c "./bin/zygisk-ptrace64 monitor &"

module/src/service.sh

@@ -24,7 +24,3 @@ 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/src/constants.rs

@@ -36,7 +36,7 @@ pub enum DaemonSocketAction {
     ReadModules,
     RequestCompanionSocket,
     GetModuleDir,
-    ZygoteRestarted,
+    ZygoteRestart,
 }
 
 // Zygisk process flags

zygiskd/src/zygiskd.rs

@@ -50,13 +50,27 @@ pub fn main() -> Result<()> {
 
     log::info!("Handle zygote connections");
     for stream in listener.incoming() {
-        let stream = stream?;
+        let mut stream = stream?;
         let context = Arc::clone(&context);
-        thread::spawn(move || {
+        let action = stream.read_u8()?;
-            if let Err(e) = handle_daemon_action(stream, &context) {
+        let action = DaemonSocketAction::try_from(action)?;
-                log::warn!("Error handling daemon action: {}\n{}", e, e.backtrace());
+        log::trace!("New daemon action {:?}", action);
+        match action {
+            DaemonSocketAction::PingHeartbeat => {
+                // Do nothing
             }
-        });
+            DaemonSocketAction::ZygoteRestart => {
+                info!("Zygote restarted, exit");
+                exit(0);
+            }
+            _ => {
+                thread::spawn(move || {
+                    if let Err(e) = handle_daemon_action(action, stream, &context) {
+                        log::warn!("Error handling daemon action: {}\n{}", e, e.backtrace());
+                    }
+                });
+            }
+        }
     }
 
     Ok(())
@@ -152,14 +166,8 @@ fn resolve_module(path: &str) -> Result<Option<ZygiskCompanionEntryFn>> {
     }
 }
 
-fn handle_daemon_action(mut stream: UnixStream, context: &Context) -> Result<()> {
+fn handle_daemon_action(action: DaemonSocketAction, mut stream: UnixStream, context: &Context) -> Result<()> {
-    let action = stream.read_u8()?;
-    let action = DaemonSocketAction::try_from(action)?;
-    log::trace!("New daemon action {:?}", action);
     match action {
-        DaemonSocketAction::PingHeartbeat => {
-            // Do nothing
-        }
         DaemonSocketAction::RequestLogcatFd => {
             loop {
                 let level = match stream.read_u8() {
@@ -227,10 +235,7 @@ fn handle_daemon_action(mut stream: UnixStream, context: &Context) -> Result<()>
             let dir = fs::File::open(dir)?;
             stream.send_fd(dir.as_raw_fd())?;
         }
-        DaemonSocketAction::ZygoteRestarted => {
+        _ => {}
-            info!("zygote restarted, exit");
-            exit(0);
-        }
     }
     Ok(())
 }