From fd9f8799a12babc1319058477037b33cf489633f Mon Sep 17 00:00:00 2001
From: 5ec1cff <ewtqyqyewtqyqy@gmail.com>
Date: Mon, 13 Nov 2023 21:51:37 +0800
Subject: [PATCH] refine tracing logic

---
 loader/src/ptracer/monitor.cpp |  2 +-
 loader/src/ptracer/ptracer.cpp | 20 ++------------------
 loader/src/ptracer/utils.cpp   | 20 +++++++++++++-------
 loader/src/ptracer/utils.hpp   |  2 +-
 4 files changed, 17 insertions(+), 27 deletions(-)

diff --git a/loader/src/ptracer/monitor.cpp b/loader/src/ptracer/monitor.cpp
index 533032a..04e7c5b 100644
--- a/loader/src/ptracer/monitor.cpp
+++ b/loader/src/ptracer/monitor.cpp
@@ -319,7 +319,7 @@ public:
                                 LOGD("stopping %d", pid);
                                 kill(pid, SIGSTOP);
                                 ptrace(PTRACE_CONT, pid, 0, 0);
-                                wait_pid(pid, &status, __WALL);
+                                waitpid(pid, &status, __WALL);
                                 if (STOPPED_WITH(SIGSTOP, 0)) {
                                     LOGD("detaching %d", pid);
                                     ptrace(PTRACE_DETACH, pid, 0, SIGSTOP);
diff --git a/loader/src/ptracer/ptracer.cpp b/loader/src/ptracer/ptracer.cpp
index 4bcbdf8..9e81776 100644
--- a/loader/src/ptracer/ptracer.cpp
+++ b/loader/src/ptracer/ptracer.cpp
@@ -68,10 +68,7 @@ bool inject_on_main(int pid, const char *lib_path) {
     if (!write_proc(pid, (uintptr_t *) addr_of_entry_addr, &break_addr, sizeof(break_addr))) return false;
     ptrace(PTRACE_CONT, pid, 0, 0);
     int status;
-    if (waitpid(pid, &status, __WALL) == -1) {
-        PLOGE("wait");
-        return false;
-    }
+    wait_for_trace(pid, &status, __WALL);
     if (WIFSTOPPED(status) && WSTOPSIG(status) == SIGSEGV) {
         if (!get_regs(pid, regs)) return false;
         if ((regs.REG_IP & ~1) != (break_addr & ~1)) {
@@ -156,19 +153,6 @@ bool inject_on_main(int pid, const char *lib_path) {
         if (!set_regs(pid, backup)) return false;
 
         return true;
-
-        /*
-        ptrace(PTRACE_CONT, pid, 0, 0);
-        waitpid(pid, &status, __WALL);
-        if (WIFSTOPPED(status)) {
-            siginfo_t siginfo;
-            ptrace(PTRACE_GETSIGINFO, pid, 0, &siginfo);
-            LOGD("process stopped by signal %d %s si_code=%d si_addr=%p", WSTOPSIG(status),
-                 strsignal(WSTOPSIG(status)), siginfo.si_code, siginfo.si_addr);
-            pause();
-        } else {
-            LOGD("other reason %d", status);
-        }*/
     } else {
         LOGE("stopped by other reason: %s", parse_status(status).c_str());
     }
@@ -179,7 +163,7 @@ bool inject_on_main(int pid, const char *lib_path) {
 
 bool trace_zygote(int pid) {
     LOGI("start tracing %d", pid);
-#define WAIT_OR_DIE if (wait_pid(pid, &status, __WALL) != pid) return false;
+#define WAIT_OR_DIE wait_for_trace(pid, &status, __WALL);
 #define CONT_OR_DIE \
     if (ptrace(PTRACE_CONT, pid, 0, 0) == -1) { \
         PLOGE("cont"); \
diff --git a/loader/src/ptracer/utils.cpp b/loader/src/ptracer/utils.cpp
index 8179d38..cc0ab4d 100644
--- a/loader/src/ptracer/utils.cpp
+++ b/loader/src/ptracer/utils.cpp
@@ -308,14 +308,12 @@ uintptr_t remote_call(int pid, struct user_regs_struct &regs, uintptr_t func_add
     }
     ptrace(PTRACE_CONT, pid, 0, 0);
     int status;
-    if (waitpid(pid, &status, __WALL) == -1) {
-        PLOGE("wait");
-    }
+    wait_for_trace(pid, &status, __WALL);
     if (!get_regs(pid, regs)) {
         LOGE("failed to get regs after call");
         return 0;
     }
-    if (WIFSTOPPED(status) && WSTOPSIG(status) == SIGSEGV) {
+    if (WSTOPSIG(status) == SIGSEGV) {
         if (regs.REG_IP != return_addr) {
             LOGE("wrong return addr %p", (void *) regs.REG_IP);
             return 0;
@@ -345,11 +343,19 @@ int fork_dont_care() {
     return pid;
 }
 
-int wait_pid(int pid, int* status, int flags) {
+void wait_for_trace(int pid, int* status, int flags) {
     while (true) {
         auto result = waitpid(pid, status, flags);
-        if (result == -1 && errno == EINTR) continue;
-        return result;
+        if (result == -1) {
+            if (errno == EINTR) continue;
+        } else {
+            PLOGE("wait %d failed", pid);
+            exit(1);
+        }
+        if (!WIFSTOPPED(*status)) {
+            LOGE("process %d not stopped for trace: %s, exit", pid, parse_status(*status).c_str());
+            exit(1);
+        }
     }
 }
 
diff --git a/loader/src/ptracer/utils.hpp b/loader/src/ptracer/utils.hpp
index 965952b..d8fc59f 100644
--- a/loader/src/ptracer/utils.hpp
+++ b/loader/src/ptracer/utils.hpp
@@ -87,7 +87,7 @@ uintptr_t remote_call(int pid, struct user_regs_struct &regs, uintptr_t func_add
 
 int fork_dont_care();
 
-int wait_pid(int pid, int* status, int flags);
+void wait_for_trace(int pid, int* status, int flags);
 
 std::string parse_status(int status);