stracers_core/

tracer.rs

use std::collections::{HashMap, HashSet};
use std::ffi::CString;
use std::io;

use nix::sys::ptrace;
use nix::sys::wait::WaitStatus;
use nix::unistd::{ForkResult, Pid, execvp, fork};

use crate::arch;
use crate::event::SyscallEvent;
use crate::platform;

/// Options for the tracer.
pub struct TracerOptions {
    /// Follow child processes created by fork/vfork/clone.
    pub follow_forks: bool,
}

impl Default for TracerOptions {
    fn default() -> Self {
        Self {
            follow_forks: false,
        }
    }
}

pub struct Tracer {
    /// The initial tracee PID (either spawned or attached).
    root_pid: Pid,
    /// All currently-traced PIDs.
    traced: HashSet<i32>,
    /// Whether each PID is currently inside a syscall (entry seen, exit pending).
    in_syscall: HashMap<i32, bool>,
    /// Syscall entry events waiting for the corresponding exit.
    pending_entry: HashMap<i32, SyscallEvent>,
    /// Whether we attached to an existing process (vs spawned).
    attached: bool,
    options: TracerOptions,
}

impl Tracer {
    /// Fork a child process, set up ptrace, and exec the given command.
    pub fn spawn(command: &[String], options: TracerOptions) -> io::Result<Self> {
        if command.is_empty() {
            return Err(io::Error::new(io::ErrorKind::InvalidInput, "empty command"));
        }

        let c_args: Vec<CString> = command
            .iter()
            .map(|s| CString::new(s.as_str()).unwrap())
            .collect();

        match unsafe { fork() }.map_err(|e| io::Error::from_raw_os_error(e as i32))? {
            ForkResult::Child => {
                platform::traceme()?;
                execvp(&c_args[0], &c_args)
                    .map_err(|e| io::Error::from_raw_os_error(e as i32))?;
                unreachable!()
            }
            ForkResult::Parent { child } => {
                // Wait for the initial SIGTRAP from the child's exec.
                platform::wait(child)?;
                platform::set_options(child, Self::ptrace_options(&options))?;
                platform::syscall_continue(child, None)?;

                let mut traced = HashSet::new();
                traced.insert(child.as_raw());

                Ok(Tracer {
                    root_pid: child,
                    traced,
                    in_syscall: HashMap::new(),
                    pending_entry: HashMap::new(),
                    attached: false,
                    options,
                })
            }
        }
    }

    /// Attach to an already-running process.
    pub fn attach(pid: i32, options: TracerOptions) -> io::Result<Self> {
        let target = Pid::from_raw(pid);
        platform::attach(target)?;
        // Wait for the SIGSTOP delivered by PTRACE_ATTACH.
        platform::wait(target)?;
        platform::set_options(target, Self::ptrace_options(&options))?;
        platform::syscall_continue(target, None)?;

        let mut traced = HashSet::new();
        traced.insert(pid);

        Ok(Tracer {
            root_pid: target,
            traced,
            in_syscall: HashMap::new(),
            pending_entry: HashMap::new(),
            attached: true,
            options,
        })
    }

    fn ptrace_options(opts: &TracerOptions) -> ptrace::Options {
        let mut flags = ptrace::Options::PTRACE_O_TRACESYSGOOD
            | ptrace::Options::PTRACE_O_TRACEEXEC;
        if opts.follow_forks {
            flags |= ptrace::Options::PTRACE_O_TRACEFORK
                | ptrace::Options::PTRACE_O_TRACEVFORK
                | ptrace::Options::PTRACE_O_TRACECLONE;
        }
        flags
    }

    /// Run the ptrace loop until all tracees exit.
    /// Calls `on_event` for each completed syscall (after both entry and exit).
    /// Returns the root tracee's exit code.
    pub fn run<F: FnMut(&SyscallEvent)>(&mut self, mut on_event: F) -> io::Result<i32> {
        let mut root_exit_code: Option<i32> = None;

        loop {
            if self.traced.is_empty() {
                return Ok(root_exit_code.unwrap_or(0));
            }

            // Wait for any traced child when following forks, otherwise wait for root.
            let status = if self.options.follow_forks {
                platform::wait_any()
            } else {
                platform::wait(self.root_pid)
            };

            let status = match status {
                Ok(s) => s,
                Err(e) if e.raw_os_error() == Some(10) /* ECHILD */ => {
                    return Ok(root_exit_code.unwrap_or(0));
                }
                Err(e) => return Err(e),
            };

            match status {
                WaitStatus::PtraceSyscall(pid) => {
                    let regs = platform::get_registers(pid)?;
                    let raw_pid = pid.as_raw();

                    if !self.in_syscall.get(&raw_pid).copied().unwrap_or(false) {
                        // Syscall entry
                        self.in_syscall.insert(raw_pid, true);
                        let (number, args) = arch::extract_syscall_entry(&regs);
                        let name = arch::syscall_name(number);
                        let decoded_args = arch::decode_entry_args(pid, number, &args);
                        let event = SyscallEvent {
                            pid: raw_pid,
                            number,
                            name,
                            args,
                            ret: None,
                            decoded_args,
                        };
                        self.pending_entry.insert(raw_pid, event);
                    } else {
                        // Syscall exit
                        self.in_syscall.insert(raw_pid, false);
                        if let Some(mut event) = self.pending_entry.remove(&raw_pid) {
                            let ret = arch::extract_return_value(&regs);
                            event.ret = Some(ret);
                            arch::decode_exit_args(
                                pid,
                                event.number,
                                &event.args,
                                ret,
                                &mut event.decoded_args,
                            );
                            on_event(&event);
                        }
                    }
                    platform::syscall_continue(pid, None)?;
                }
                WaitStatus::PtraceEvent(pid, _sig, event) => {
                    if event == nix::libc::PTRACE_EVENT_FORK as i32
                        || event == nix::libc::PTRACE_EVENT_VFORK as i32
                        || event == nix::libc::PTRACE_EVENT_CLONE as i32
                    {
                        // A new child was created.
                        if let Ok(new_pid) = platform::get_event(pid) {
                            let new_pid_raw = new_pid as i32;
                            self.traced.insert(new_pid_raw);
                            let new_pid_nix = Pid::from_raw(new_pid_raw);
                            // Wait for the new child to stop, then configure and resume.
                            let _ = platform::wait(new_pid_nix);
                            let _ = platform::set_options(
                                new_pid_nix,
                                Self::ptrace_options(&self.options),
                            );
                            let _ = platform::syscall_continue(new_pid_nix, None);
                        }
                    } else if event == nix::libc::PTRACE_EVENT_EXEC as i32 {
                        // After PTRACE_EVENT_EXEC, the kernel still delivers a
                        // syscall-exit-stop for the execve. Keep in_syscall=true
                        // so that exit is correctly paired with the pending entry.
                    }
                    platform::syscall_continue(pid, None)?;
                }
                WaitStatus::Exited(pid, code) => {
                    let raw_pid = pid.as_raw();
                    self.traced.remove(&raw_pid);
                    self.in_syscall.remove(&raw_pid);
                    self.pending_entry.remove(&raw_pid);
                    if pid == self.root_pid {
                        root_exit_code = Some(code);
                    }
                    if !self.options.follow_forks || self.traced.is_empty() {
                        return Ok(root_exit_code.unwrap_or(code));
                    }
                }
                WaitStatus::Signaled(pid, sig, _core_dumped) => {
                    let raw_pid = pid.as_raw();
                    self.traced.remove(&raw_pid);
                    self.in_syscall.remove(&raw_pid);
                    self.pending_entry.remove(&raw_pid);
                    if pid == self.root_pid {
                        root_exit_code = Some(128 + sig as i32);
                    }
                    if !self.options.follow_forks || self.traced.is_empty() {
                        return Ok(root_exit_code.unwrap_or(128 + sig as i32));
                    }
                }
                WaitStatus::Stopped(pid, sig) => {
                    // Real signal — forward it to the tracee.
                    platform::syscall_continue(pid, Some(sig))?;
                }
                _ => {
                    // StillAlive or other — shouldn't happen with blocking wait.
                }
            }
        }
    }
}

impl Drop for Tracer {
    fn drop(&mut self) {
        // If we attached, try to detach cleanly from any remaining tracees.
        if self.attached {
            for &pid in &self.traced {
                let _ = platform::detach(Pid::from_raw(pid), None);
            }
        }
    }
}