From 09e18240809acad7ba12fb5bc62898b9a8592256 Mon Sep 17 00:00:00 2001
From: Levi Zim <rsworktech@outlook.com>
Date: Sun, 20 Oct 2024 21:43:30 +0800
Subject: [PATCH] refactor/ptrace: use PTRACE_GET_SYSCALL_INFO for syscall
 result

---
 src/arch.rs          |  1 -
 src/arch/aarch64.rs  |  4 ----
 src/arch/riscv64.rs  |  4 ----
 src/arch/x86_64.rs   | 11 ++---------
 src/tracer.rs        | 10 +++++-----
 src/tracer/ptrace.rs | 34 +++++++++++++++++++++++++++++++++-
 6 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/src/arch.rs b/src/arch.rs
index ab8bb69..501d4a8 100644
--- a/src/arch.rs
+++ b/src/arch.rs
@@ -24,5 +24,4 @@ pub use audit::*;
 
 pub trait RegsExt {
   fn syscall_arg(&self, idx: usize, is_32bit: bool) -> usize;
-  fn syscall_ret(&self) -> isize;
 }
diff --git a/src/arch/aarch64.rs b/src/arch/aarch64.rs
index e4250a2..3b134f4 100644
--- a/src/arch/aarch64.rs
+++ b/src/arch/aarch64.rs
@@ -23,8 +23,4 @@ impl RegsExt for Regs {
       _ => unimplemented!(),
     } as usize)
   }
-
-  fn syscall_ret(&self) -> isize {
-    self.regs[0] as isize
-  }
 }
diff --git a/src/arch/riscv64.rs b/src/arch/riscv64.rs
index 35ca7a8..5cdbe8c 100644
--- a/src/arch/riscv64.rs
+++ b/src/arch/riscv64.rs
@@ -23,8 +23,4 @@ impl RegsExt for Regs {
       _ => unimplemented!(),
     } as usize)
   }
-
-  fn syscall_ret(&self) -> isize {
-    self.a0 as isize
-  }
 }
diff --git a/src/arch/x86_64.rs b/src/arch/x86_64.rs
index 9742359..95ca2f6 100644
--- a/src/arch/x86_64.rs
+++ b/src/arch/x86_64.rs
@@ -61,7 +61,7 @@ use super::RegsExt;
 impl RegsExt for Regs {
   fn syscall_arg(&self, idx: usize, is_32bit: bool) -> usize {
     match self {
-      Regs::X86(regs) => {
+      Self::X86(regs) => {
         debug_assert!(is_32bit);
         (match idx {
           0 => regs.ebx,
@@ -73,7 +73,7 @@ impl RegsExt for Regs {
           _ => unreachable!(),
         } as usize)
       }
-      Regs::X64(regs) => {
+      Self::X64(regs) => {
         if is_32bit {
           (match idx {
             0 => regs.rbx,
@@ -98,11 +98,4 @@ impl RegsExt for Regs {
       }
     }
   }
-
-  fn syscall_ret(&self) -> isize {
-    match self {
-      Regs::X86(regs) => regs.eax as isize,
-      Regs::X64(regs) => regs.rax as isize,
-    }
-  }
 }
diff --git a/src/tracer.rs b/src/tracer.rs
index d77d9fa..9fdeeeb 100644
--- a/src/tracer.rs
+++ b/src/tracer.rs
@@ -817,17 +817,17 @@ impl Tracer {
     let mut store = self.store.write().unwrap();
     let p = store.get_current_mut(pid).unwrap();
     p.presyscall = !p.presyscall;
-    let regs = match ptrace_getregs(pid) {
-      Ok(regs) => regs,
+    let result = match ptrace::syscall_exit_result(pid) {
+      Ok(r) => r,
       Err(Errno::ESRCH) => {
-        info!("ptrace getregs failed: {pid}, ESRCH, child probably gone!");
+        info!("ptrace get_syscall_info failed: {pid}, ESRCH, child probably gone!");
         return Ok(());
       }
       e => e?,
     };
-    let result = regs.syscall_ret() as i64;
     // If exec is successful, the register value might be clobbered.
-    let exec_result = if p.is_exec_successful { 0 } else { result };
+    // TODO: would the value in ptrace_syscall_info be clobbered?
+    let exec_result = if p.is_exec_successful { 0 } else { result } as i64;
     match p.syscall {
       Syscall::Execve | Syscall::Execveat => {
         trace!("post execve(at) in exec");
diff --git a/src/tracer/ptrace.rs b/src/tracer/ptrace.rs
index 9082971..bb17739 100644
--- a/src/tracer/ptrace.rs
+++ b/src/tracer/ptrace.rs
@@ -5,7 +5,7 @@ use nix::{
   errno::Errno,
   libc::{
     ptrace_syscall_info, SYS_execve, SYS_execveat, PTRACE_GET_SYSCALL_INFO,
-    PTRACE_SYSCALL_INFO_ENTRY, PTRACE_SYSCALL_INFO_SECCOMP,
+    PTRACE_SYSCALL_INFO_ENTRY, PTRACE_SYSCALL_INFO_EXIT, PTRACE_SYSCALL_INFO_SECCOMP,
   },
   sys::{ptrace, signal::Signal},
   unistd::Pid,
@@ -81,6 +81,11 @@ impl SyscallInfo {
   }
 }
 
+/// Get [`SyscallInfo`] on ptrace syscall entry/seccomp stop
+///
+/// # Precondition
+///
+/// The caller is the tracer thread and at the syscall entry/seccomp stop.
 pub fn syscall_entry_info(pid: Pid) -> Result<SyscallInfo, Errno> {
   let mut info = MaybeUninit::<ptrace_syscall_info>::uninit();
   let info = unsafe {
@@ -108,6 +113,33 @@ pub fn syscall_entry_info(pid: Pid) -> Result<SyscallInfo, Errno> {
   Ok(SyscallInfo { arch, number })
 }
 
+/// Get syscall result on ptrace syscall exit stop
+///
+/// # Precondition
+///
+/// The caller is the tracer thread and at the syscall exit stop.
+pub fn syscall_exit_result(pid: Pid) -> Result<isize, Errno> {
+  let mut info = MaybeUninit::<ptrace_syscall_info>::uninit();
+  let info = unsafe {
+    let ret = nix::libc::ptrace(
+      PTRACE_GET_SYSCALL_INFO,
+      pid.as_raw(),
+      size_of::<ptrace_syscall_info>(),
+      info.as_mut_ptr(),
+    );
+    if ret < 0 {
+      return Err(Errno::last());
+    } else {
+      info.assume_init()
+    }
+  };
+  if info.op == PTRACE_SYSCALL_INFO_EXIT {
+    Ok(unsafe { info.u.exit.sval } as isize)
+  } else {
+    Err(Errno::EINVAL)
+  }
+}
+
 pub fn ptrace_getregs(pid: Pid) -> Result<Regs, Errno> {
   // https://github.com/torvalds/linux/blob/v6.9/include/uapi/linux/elf.h#L378
   // libc crate doesn't provide this constant when using musl libc.