llmdbg/src/target.rs

//! Linux x86-64 ptrace wrapper for v0.1.
//!
//! Spawn-only, single-threaded debuggee. The `Target` owns the child PID
//! and tracks whether the child is currently stopped or has exited.

use nix::sys::ptrace;
use nix::sys::signal::Signal;
use nix::sys::uio::{process_vm_readv, RemoteIoVec};
use nix::sys::wait::{waitpid, WaitStatus};
use nix::unistd::{execvp, fork, ForkResult, Pid};
use std::ffi::CString;
use std::io::IoSliceMut;

use crate::breakpoint::BreakpointSet;
use crate::symbols::SymbolTable;

/// Why the target stopped after a step/continue, or that it has exited.
#[derive(Debug, Clone)]
pub enum StopReason {
    /// Hit a SIGTRAP from single-stepping.
    Step,
    /// Hit a software breakpoint. Carries the breakpoint id.
    Breakpoint(u32),
    /// Caught some other signal — we don't yet inject it back, just report.
    Signal(Signal),
    /// Process exited normally with this code.
    Exited(i32),
    /// Process was killed by a signal.
    Killed(Signal),
}

impl StopReason {
    /// Short tag for the REPL caller to print.
    pub fn tag(&self) -> String {
        match self {
            StopReason::Step => "step".into(),
            StopReason::Breakpoint(id) => format!("breakpoint:{}", id),
            StopReason::Signal(s) => format!("signal:{}", s),
            StopReason::Exited(c) => format!("exited:{}", c),
            StopReason::Killed(s) => format!("killed:{}", s),
        }
    }

    pub fn is_alive(&self) -> bool {
        matches!(self, StopReason::Step | StopReason::Breakpoint(_) | StopReason::Signal(_))
    }
}

#[derive(Debug)]
#[allow(dead_code)]
pub enum State {
    /// Stopped, ready to read regs/memory and step.
    Stopped(StopReason),
    /// Dead. Subsequent operations should fail cleanly.
    Dead(StopReason),
}

pub struct Target {
    pid: Pid,
    state: State,
    /// Snapshot of registers from the previous stop. None until at least one
    /// step/cont has happened after spawn.
    prev_regs: Option<libc::user_regs_struct>,
    /// ELF symbols of the main executable. None if parsing failed (e.g.
    /// stripped binary, unreadable /proc) — symbol lookups become no-ops.
    symbols: Option<SymbolTable>,
    /// Software (INT3) breakpoints.
    breakpoints: BreakpointSet,
}

impl Target {
    pub fn pid(&self) -> i32 {
        self.pid.as_raw()
    }

    #[allow(dead_code)]
    pub fn state(&self) -> &State {
        &self.state
    }

    /// True if the target is currently stopped and inspectable.
    pub fn is_stopped(&self) -> bool {
        matches!(self.state, State::Stopped(_))
    }

    /// fork() + PTRACE_TRACEME + execvp(). Returns once the child has been
    /// caught at its initial post-exec SIGTRAP.
    pub fn spawn(path: &str, args: &[String]) -> Result<Target, String> {
        let c_path = CString::new(path).map_err(|e| format!("bad path: {}", e))?;
        let mut c_args: Vec<CString> = Vec::with_capacity(args.len() + 1);
        c_args.push(c_path.clone());
        for a in args {
            c_args.push(CString::new(a.as_str()).map_err(|e| format!("bad arg: {}", e))?);
        }

        // SAFETY: fork() must only call async-signal-safe functions in the
        // child. We use ptrace::traceme and execvp, both fine.
        let fork_result = unsafe { fork() }.map_err(|e| format!("fork failed: {}", e))?;

        match fork_result {
            ForkResult::Child => {
                // In the child. Anything that goes wrong here we report via
                // exit code; the parent will see the abnormal exit on waitpid.
                if ptrace::traceme().is_err() {
                    unsafe { libc::_exit(127) };
                }
                // execvp replaces the process image; on success it does not return.
                let _ = execvp(&c_path, &c_args);
                unsafe { libc::_exit(127) };
            }
            ForkResult::Parent { child } => {
                // Wait for the initial post-exec SIGTRAP.
                let status = waitpid(child, None)
                    .map_err(|e| format!("waitpid (initial) failed: {}", e))?;
                match status {
                    WaitStatus::Stopped(_, Signal::SIGTRAP) => {
                        let symbols = SymbolTable::from_pid(child.as_raw()).ok();
                        Ok(Target {
                            pid: child,
                            state: State::Stopped(StopReason::Step),
                            prev_regs: None,
                            symbols,
                            breakpoints: BreakpointSet::new(),
                        })
                    }
                    WaitStatus::Exited(_, code) => {
                        Err(format!("child exited immediately with code {} (exec failed?)", code))
                    }
                    other => Err(format!("unexpected initial wait status: {:?}", other)),
                }
            }
        }
    }

    pub fn prev_regs(&self) -> Option<libc::user_regs_struct> {
        self.prev_regs
    }

    pub fn symbols(&self) -> Option<&SymbolTable> {
        self.symbols.as_ref()
    }

    pub fn breakpoints(&self) -> &BreakpointSet {
        &self.breakpoints
    }

    // ----- breakpoint ops -------------------------------------------------

    /// Install a software breakpoint at `addr`. Reads the original byte,
    /// writes 0xCC. Returns the breakpoint id.
    pub fn breakpoint_set(&mut self, addr: u64, name: Option<String>) -> Result<u32, String> {
        if !self.is_stopped() {
            return Err("target is not stopped".into());
        }
        let orig = read_byte(self.pid, addr)?;
        let id = self.breakpoints.add(addr, orig, name)?;
        if let Err(e) = write_byte(self.pid, addr, 0xCC) {
            // Roll back: the bp record was added but writing failed.
            let _ = self.breakpoints.remove(id);
            return Err(format!("install bp at 0x{:x}: {}", addr, e));
        }
        Ok(id)
    }

    /// Remove a breakpoint by id, restoring the original byte if it was
    /// enabled. Returns Ok(()) on success, Err if the id doesn't exist.
    pub fn breakpoint_remove(&mut self, id: u32) -> Result<(), String> {
        if !self.is_stopped() {
            return Err("target is not stopped".into());
        }
        let bp = self.breakpoints.remove(id)
            .ok_or_else(|| format!("no breakpoint with id {}", id))?;
        if bp.enabled {
            write_byte(self.pid, bp.addr, bp.orig_byte)
                .map_err(|e| format!("remove bp 0x{:x}: {}", bp.addr, e))?;
        }
        Ok(())
    }

    /// Look up a bp id by name; convenience for callers.
    pub fn breakpoint_id_by_name(&self, name: &str) -> Option<u32> {
        self.breakpoints.id_by_name(name)
    }

    // ----- step / cont ----------------------------------------------------

    /// PTRACE_SINGLESTEP + waitpid. Updates internal state. If the current PC
    /// has an enabled breakpoint, we temporarily disarm it, step over the
    /// original instruction, then re-arm.
    pub fn step(&mut self) -> Result<StopReason, String> {
        if !self.is_stopped() {
            return Err("target is not stopped".into());
        }
        self.prev_regs = ptrace::getregs(self.pid).ok();

        let pc = self.prev_regs.map(|r| r.rip).unwrap_or(0);
        let bp_at_pc = self.breakpoints.find_by_addr(pc).map(|b| (b.id, b.orig_byte, b.enabled));

        if let Some((_id, orig, enabled)) = bp_at_pc {
            if enabled {
                // Disarm, step, re-arm.
                write_byte(self.pid, pc, orig)
                    .map_err(|e| format!("disarm bp at 0x{:x}: {}", pc, e))?;
                ptrace::step(self.pid, None)
                    .map_err(|e| format!("ptrace step: {}", e))?;
                let reason = self.wait_for_stop_raw()?;
                // Re-arm even if the step exited or signalled (best-effort).
                if reason.is_alive() {
                    let _ = write_byte(self.pid, pc, 0xCC);
                }
                self.classify_stop(reason);
                let r = match &self.state {
                    State::Stopped(r) => r.clone(),
                    State::Dead(r) => r.clone(),
                };
                return Ok(r);
            }
        }

        ptrace::step(self.pid, None).map_err(|e| format!("ptrace step: {}", e))?;
        let reason = self.wait_for_stop_raw()?;
        self.classify_stop(reason);
        match &self.state {
            State::Stopped(r) => Ok(r.clone()),
            State::Dead(r) => Ok(r.clone()),
        }
    }

    /// PTRACE_CONT + waitpid. If the current PC is a breakpoint, we have to
    /// step over it first (same disarm/step/re-arm dance), then continue.
    pub fn cont(&mut self) -> Result<StopReason, String> {
        if !self.is_stopped() {
            return Err("target is not stopped".into());
        }
        self.prev_regs = ptrace::getregs(self.pid).ok();

        // Step over current PC if it has an active breakpoint.
        let pc = self.prev_regs.map(|r| r.rip).unwrap_or(0);
        if let Some(bp) = self.breakpoints.find_by_addr(pc).cloned() {
            if bp.enabled {
                write_byte(self.pid, pc, bp.orig_byte)
                    .map_err(|e| format!("disarm bp at 0x{:x}: {}", pc, e))?;
                ptrace::step(self.pid, None)
                    .map_err(|e| format!("ptrace step (over-bp): {}", e))?;
                let inner = self.wait_for_stop_raw()?;
                // If the step itself exited/signalled, classify and return.
                if !inner.is_alive() {
                    self.classify_stop(inner);
                    return match &self.state {
                        State::Stopped(r) | State::Dead(r) => Ok(r.clone()),
                    };
                }
                // Re-arm and fall through to PTRACE_CONT.
                let _ = write_byte(self.pid, pc, 0xCC);
            }
        }

        ptrace::cont(self.pid, None).map_err(|e| format!("ptrace cont: {}", e))?;
        let reason = self.wait_for_stop_raw()?;
        // If we hit a breakpoint, rewind RIP to bp_addr (kernel left it at
        // bp_addr+1) so the user sees pc == bp_addr, which is what they expect.
        self.handle_possible_breakpoint(reason)?;
        match &self.state {
            State::Stopped(r) | State::Dead(r) => Ok(r.clone()),
        }
    }

    /// Raw waitpid: just decodes the status into a tentative StopReason::Step
    /// or Signal/Exited/Killed. Doesn't yet check for breakpoint matches.
    fn wait_for_stop_raw(&self) -> Result<StopReason, String> {
        let status = waitpid(self.pid, None).map_err(|e| format!("waitpid: {}", e))?;
        let reason = match status {
            WaitStatus::Stopped(_, Signal::SIGTRAP) => StopReason::Step,
            WaitStatus::Stopped(_, sig) => StopReason::Signal(sig),
            WaitStatus::Exited(_, code) => StopReason::Exited(code),
            WaitStatus::Signaled(_, sig, _) => StopReason::Killed(sig),
            other => return Err(format!("unexpected wait status: {:?}", other)),
        };
        Ok(reason)
    }

    /// After a wait, if the reason is a SIGTRAP, check whether PC-1 matches a
    /// known breakpoint; if so, rewind PC and mark the stop as
    /// Breakpoint(id). Update self.state.
    fn handle_possible_breakpoint(&mut self, mut reason: StopReason) -> Result<(), String> {
        if let StopReason::Step = reason {
            // After INT3 the kernel leaves RIP at addr+1.
            if let Ok(r) = ptrace::getregs(self.pid) {
                let candidate = r.rip.wrapping_sub(1);
                if let Some(bp) = self.breakpoints.find_by_addr(candidate).cloned() {
                    // Rewind RIP to the breakpoint address.
                    let mut new_regs = r;
                    new_regs.rip = candidate;
                    ptrace::setregs(self.pid, new_regs)
                        .map_err(|e| format!("setregs (rewind bp): {}", e))?;
                    if let Some(bpm) = self.breakpoints.get_mut(bp.id) {
                        bpm.hits += 1;
                    }
                    reason = StopReason::Breakpoint(bp.id);
                }
            }
        }
        self.classify_stop(reason);
        Ok(())
    }

    fn classify_stop(&mut self, reason: StopReason) {
        self.state = if reason.is_alive() {
            State::Stopped(reason)
        } else {
            State::Dead(reason)
        };
    }

    pub fn regs(&self) -> Result<libc::user_regs_struct, String> {
        if !self.is_stopped() {
            return Err("target is not stopped".into());
        }
        ptrace::getregs(self.pid).map_err(|e| format!("getregs: {}", e))
    }

    /// Read `len` bytes from the target's memory at `addr` via process_vm_readv.
    /// Much faster than peek loops for >8 bytes, and works for any size.
    pub fn read_mem(&self, addr: u64, len: usize) -> Result<Vec<u8>, String> {
        if !self.is_stopped() {
            return Err("target is not stopped".into());
        }
        let mut buf = vec![0u8; len];
        let mut local = [IoSliceMut::new(&mut buf)];
        let remote = [RemoteIoVec { base: addr as usize, len }];
        let n = process_vm_readv(self.pid, &mut local, &remote)
            .map_err(|e| format!("read_mem(0x{:x}, {}): {}", addr, len, e))?;
        if n != len {
            return Err(format!(
                "short read at 0x{:x}: got {} of {}",
                addr, n, len
            ));
        }
        Ok(buf)
    }
}

impl Drop for Target {
    fn drop(&mut self) {
        // Best effort: if the target is still alive, detach so we don't leave
        // it stopped. If it's already dead, nothing to do.
        if self.is_stopped() {
            let _ = ptrace::detach(self.pid, None);
        }
    }
}

// --- ptrace byte-level peek/poke for breakpoint patching ------------------

/// Read a single byte at `addr` via PTRACE_PEEKDATA.
fn read_byte(pid: Pid, addr: u64) -> Result<u8, String> {
    let word = ptrace::read(pid, addr as ptrace::AddressType)
        .map_err(|e| format!("peek 0x{:x}: {}", addr, e))?;
    let bytes = (word as u64).to_le_bytes();
    Ok(bytes[0])
}

/// Write a single byte at `addr` via PEEK+POKE (read-modify-write of one word).
/// Writes the 8-byte word starting at addr; the upper 7 bytes are unchanged
/// because we read them first.
fn write_byte(pid: Pid, addr: u64, value: u8) -> Result<(), String> {
    let word = ptrace::read(pid, addr as ptrace::AddressType)
        .map_err(|e| format!("peek 0x{:x} (for write): {}", addr, e))?;
    let mut bytes = (word as u64).to_le_bytes();
    bytes[0] = value;
    let new_word = u64::from_le_bytes(bytes) as i64;
    ptrace::write(pid, addr as ptrace::AddressType, new_word as libc::c_long)
        .map_err(|e| format!("poke 0x{:x}: {}", addr, e))?;
    Ok(())
}