ptrace_do/
lib.rs

1mod arch;
2use arch::UserRegs;
3
4use libc::{pid_t, ptrace, PTRACE_ATTACH, PTRACE_CONT, PTRACE_DETACH};
5use std::{mem, process::Child};
6use thiserror::Error;
7
8/// Utility function that converts the usize inputs individually with the appropriate endianness into a Vec<u8>
9fn usize_arr_to_u8(data: &[usize]) -> Vec<u8> {
10    let mut arr: Vec<u8> = Vec::new();
11    for p in data {
12        if cfg!(target_endian = "big") {
13            arr.extend_from_slice(&p.to_be_bytes());
14        } else {
15            arr.extend_from_slice(&p.to_le_bytes());
16        }
17    }
18    return arr;
19}
20
21/// Enum containing all errors tracing can witness
22#[derive(Debug, Error)]
23pub enum TraceError {
24    /// Error spawning from syscall interactions
25    #[error("Ptrace error: `{0}`")]
26    Ptrace(std::io::Error),
27
28    /// Error during read and writing of process's memory
29    #[error("IO Error: `{0}`")]
30    Io(#[from] std::io::Error),
31
32    /// Contexted error
33    #[error("General `{0}`")]
34    General(&'static str),
35}
36
37/// Internal Result type
38pub type TraceResult<T> = Result<T, TraceError>;
39
40/// Trait representing the type represents a process and has a unique process identifier, pid.
41pub trait ProcessIdentifier {
42    /// Acces the pid
43    fn pid(&self) -> pid_t;
44}
45
46/// A raw process, initialized by an explicit pid. Unsafe and prone to permission errors beware,
47/// know your environment and security restrictions.
48pub struct RawProcess {
49    pid: pid_t,
50}
51
52impl ProcessIdentifier for RawProcess {
53    fn pid(&self) -> pid_t {
54        self.pid
55    }
56}
57
58impl RawProcess {
59    /// Initialize raw process with the explicit pid
60    pub fn new(pid: pid_t) -> Self {
61        Self { pid }
62    }
63}
64
65/// An owned process.
66pub struct OwnedProcess {
67    child: Child,
68}
69
70/// An owned process can be initialized from a Child os process
71impl From<Child> for OwnedProcess {
72    fn from(value: Child) -> Self {
73        Self { child: value }
74    }
75}
76
77/// Will attempt to kill the child on drop.
78/// Never panics only logs an error
79impl Drop for OwnedProcess {
80    fn drop(&mut self) {
81        if let Err(e) = self.child.kill() {
82            tracing::error!("Unable to kill owned process's child {e:?}");
83        } else {
84            tracing::info!("Owned process has been killed.");
85        }
86    }
87}
88
89impl ProcessIdentifier for OwnedProcess {
90    fn pid(&self) -> pid_t {
91        self.child.id() as pid_t
92    }
93}
94
95/// A process actively being traced.
96/// Simply a wrapper around a type with an identifiable process identifier.
97pub struct TracedProcess<T>
98where
99    T: ProcessIdentifier,
100{
101    process: T,
102}
103
104impl<T> ProcessIdentifier for TracedProcess<T>
105where
106    T: ProcessIdentifier,
107{
108    fn pid(&self) -> pid_t {
109        self.process.pid()
110    }
111}
112
113/// The available wait options
114#[allow(unused)]
115enum WaitOptions {
116    None,
117    NoHang,
118    Untraced,
119    Continued,
120}
121
122/// A wait option can be extracted from an i32
123impl From<WaitOptions> for i32 {
124    fn from(val: WaitOptions) -> Self {
125        match val {
126            WaitOptions::None => 0,
127            WaitOptions::NoHang => libc::WNOHANG,
128            WaitOptions::Untraced => libc::WUNTRACED,
129            WaitOptions::Continued => libc::WCONTINUED,
130        }
131    }
132}
133
134/// Wait status result of a ptrace step
135struct WaitStatus(i32);
136
137impl WaitStatus {
138    /// is stopped status
139    fn is_stop(&self) -> bool {
140        libc::WIFSTOPPED(self.0)
141    }
142
143    /// is signaled status
144    fn is_signaled(&self) -> bool {
145        libc::WIFSIGNALED(self.0)
146    }
147
148    /// is continued status
149    fn is_continued(&self) -> bool {
150        libc::WIFCONTINUED(self.0)
151    }
152
153    /// is exited status
154    fn is_exited(&self) -> bool {
155        libc::WIFEXITED(self.0)
156    }
157
158    /// is stopcode status
159    fn stop_code(&self) -> i32 {
160        libc::WSTOPSIG(self.0)
161    }
162}
163
164impl std::fmt::Debug for WaitStatus {
165    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
166        f.debug_struct("WaitStatus")
167            .field("is_stopped", &self.is_stop())
168            .field("is_signaled", &self.is_signaled())
169            .field("is_continued", &self.is_continued())
170            .field("is_exited", &self.is_exited())
171            .field("stop_code", &self.stop_code())
172            .finish()
173    }
174}
175
176/// Represents a traced process which is in the process of building a `frame`.
177/// `frame` can be thought of a mutable view into a stopped process's execution.
178/// Given you own a process frame, it is an appropriate time to edit registers, change instructions, and edit memory.
179pub struct ProcessFrame<T>
180where
181    T: ProcessIdentifier,
182{
183    process: TracedProcess<T>,
184}
185
186impl<T> ProcessFrame<T>
187where
188    T: ProcessIdentifier,
189{
190    /// aarch64 specific get registers functionality.
191    /// uses iovec's and GETREGSET
192    #[cfg(target_arch = "aarch64")]
193    pub fn query_registers(&mut self) -> TraceResult<UserRegs> {
194        let mut registers: UserRegs = unsafe { mem::zeroed() };
195        let mut iovec = libc::iovec {
196            iov_base: &mut registers as *mut _ as *mut core::ffi::c_void,
197            iov_len: std::mem::size_of::<UserRegs>() as libc::size_t,
198        };
199
200        let result = unsafe { ptrace(libc::PTRACE_GETREGSET, self.process.pid(), 1, &mut iovec) };
201
202        match result == -1 {
203            true => Err(TraceError::Ptrace(std::io::Error::last_os_error())),
204            false => Ok(registers),
205        }
206    }
207
208    /// aarch64 specific set registers functionality.
209    /// uses iovec's and SETREGSET
210    #[cfg(target_arch = "aarch64")]
211    pub fn set_registers(&mut self, mut registers: UserRegs) -> TraceResult<()> {
212        let mut iovec = libc::iovec {
213            iov_base: &mut registers as *mut _ as *mut core::ffi::c_void,
214            iov_len: std::mem::size_of::<UserRegs>() as libc::size_t,
215        };
216
217        let result = unsafe { ptrace(libc::PTRACE_SETREGSET, self.process.pid(), 1, &mut iovec) };
218
219        match result == -1 {
220            true => Err(TraceError::Ptrace(std::io::Error::last_os_error())),
221            false => Ok(()),
222        }
223    }
224
225    /// Attempts to invoke a remote function with the provided parameters.
226    /// Internally this is an os cfg controlled function that write the inputted parameters according to the architectures expectations.
227    /// Additionally it prepares the provided return address
228    #[cfg(target_arch = "aarch64")]
229    pub fn invoke_remote(
230        mut self,
231        func_address: usize,
232        return_address: usize,
233        parameters: &[usize],
234    ) -> TraceResult<(UserRegs, ProcessFrame<T>)> {
235        use std::mem::size_of;
236
237        const REGISTER_ARGUMENTS: usize = 8;
238        let mut current_registers = self.query_registers()?;
239        tracing::trace!(
240            "Initial registers acquired Current PC: {:X?}",
241            current_registers.program_counter()
242        );
243
244        let cached_registers = current_registers.clone();
245        current_registers.set_stack_pointer(current_registers.stack_pointer() & !0xfusize);
246        for (i, param) in parameters[..std::cmp::min(parameters.len(), REGISTER_ARGUMENTS)]
247            .iter()
248            .enumerate()
249        {
250            let reg: usize = *param;
251            current_registers.regs[i] = reg as u64;
252            tracing::trace!("Applying register {i} with param {}", reg);
253        }
254
255        if parameters.len() > REGISTER_ARGUMENTS {
256            let stack_arguments = &parameters[REGISTER_ARGUMENTS..];
257            tracing::trace!("Remaining stack arguments: {:?}", stack_arguments);
258
259            // adjust stack pointer
260            current_registers.set_stack_pointer(
261                current_registers.stack_pointer()
262                    - (((stack_arguments.len() + 1) & !1usize) * size_of::<usize>()),
263            );
264
265            self.write_memory(
266                current_registers.stack_pointer(),
267                usize_arr_to_u8(stack_arguments).as_slice(),
268            )?;
269        };
270
271        // set registers cached_registers
272        current_registers.set_program_counter(func_address);
273        if (current_registers.program_counter() & 1) != 0 {
274            current_registers.set_program_counter(current_registers.program_counter() & !1);
275            current_registers
276                .set_cpsr((current_registers.cpsr() as u32 | (arch::CPSR_T_MASK)) as usize);
277        } else {
278            current_registers
279                .set_cpsr((current_registers.cpsr() as u32 & !(arch::CPSR_T_MASK)) as usize);
280        }
281        current_registers.set_lr(return_address);
282        tracing::trace!(
283            "Executing with PC: {:X?}, and arguments {parameters:?}",
284            func_address
285        );
286
287        self.set_registers(current_registers)?;
288        tracing::trace!("Registers successfully injected.");
289
290        let mut frame = self.step_cont()?;
291        let result_regs = frame.query_registers()?;
292        tracing::trace!("Result {result_regs:#?}");
293
294        frame.set_registers(cached_registers)?;
295        Ok((result_regs, frame))
296    }
297
298    /// gets process frame registers.
299    /// internally uses GETREGS
300    #[cfg(any(target_arch = "x86", target_arch = "x86_64", target_arch = "arm"))]
301    pub fn query_registers(&mut self) -> TraceResult<UserRegs> {
302        let mut registers: UserRegs = unsafe { mem::zeroed() };
303        let result = unsafe { ptrace(libc::PTRACE_GETREGS, self.process.pid(), 0, &mut registers) };
304
305        match result == -1 {
306            true => Err(TraceError::Ptrace(std::io::Error::last_os_error())),
307            false => Ok(registers),
308        }
309    }
310
311    /// sets a process frame registers.
312    /// internally uses SETREGS
313    #[cfg(any(target_arch = "x86", target_arch = "x86_64", target_arch = "arm"))]
314    pub fn set_registers(&mut self, registers: UserRegs) -> TraceResult<()> {
315        let result = unsafe { ptrace(libc::PTRACE_SETREGS, self.process.pid(), 0, &registers) };
316
317        match result == -1 {
318            true => Err(TraceError::Ptrace(std::io::Error::last_os_error())),
319            false => Ok(()),
320        }
321    }
322
323    /// Attempts to invoke a remote function with the provided parameters.
324    /// Internally this is an os cfg controlled function that write the inputted parameters according to the architectures expectations.
325    /// Additionally it prepares the provided return address
326    #[cfg(target_arch = "arm")]
327    pub fn invoke_remote(
328        mut self,
329        func_address: usize,
330        return_address: usize,
331        parameters: &[usize],
332    ) -> TraceResult<(UserRegs, ProcessFrame<T>)> {
333        use std::mem::size_of;
334
335        const REGISTER_ARGUMENTS: usize = 4;
336        let mut current_registers = self.query_registers()?;
337        tracing::info!(
338            "Initial registers acquired Current PC: {:X?}",
339            current_registers.program_counter()
340        );
341
342        let cached_registers = current_registers.clone();
343        current_registers.set_stack_pointer(current_registers.stack_pointer() & !0xfusize);
344        for (i, param) in parameters[..std::cmp::min(parameters.len(), REGISTER_ARGUMENTS)]
345            .iter()
346            .enumerate()
347        {
348            let reg: u32 = (*param).try_into().unwrap();
349            current_registers.regs[i] = reg;
350            tracing::trace!("Applying register {i} with param {}", reg);
351        }
352
353        if parameters.len() > REGISTER_ARGUMENTS {
354            let stack_arguments = &parameters[REGISTER_ARGUMENTS..];
355            tracing::trace!("Remaining stack arguments: {:?}", stack_arguments);
356
357            // adjust stack pointer
358            current_registers.set_stack_pointer(
359                current_registers.stack_pointer()
360                    - (((stack_arguments.len() + 3) & !3usize) * size_of::<usize>()),
361            );
362
363            self.write_memory(
364                current_registers.stack_pointer(),
365                usize_arr_to_u8(stack_arguments).as_slice(),
366            )?;
367        };
368
369        // set registers cached_registers
370        current_registers.set_program_counter(func_address);
371        if (current_registers.program_counter() & 1) != 0 {
372            current_registers.set_program_counter(current_registers.program_counter() & !1);
373            current_registers
374                .set_cpsr((current_registers.cpsr() as u32 | (arch::CPSR_T_MASK)) as usize);
375        } else {
376            current_registers
377                .set_cpsr((current_registers.cpsr() as u32 & !(arch::CPSR_T_MASK)) as usize);
378        }
379        current_registers.set_lr(return_address);
380        tracing::trace!(
381            "Executing with PC: {:X?}, and arguments {parameters:?}",
382            func_address
383        );
384
385        self.set_registers(current_registers)?;
386        tracing::trace!("Registers successfully injected.");
387
388        let mut frame = self.step_cont()?;
389        let result_regs = frame.query_registers()?;
390        tracing::trace!("Result {result_regs:#?}");
391
392        frame.set_registers(cached_registers)?;
393        Ok((result_regs, frame))
394    }
395
396    /// Attempts to invoke a remote function with the provided parameters.
397    /// Internally this is an os cfg controlled function that write the inputted parameters according to the architectures expectations.
398    /// Additionally it prepares the provided return address
399    #[cfg(target_arch = "x86")]
400    pub fn invoke_remote(
401        mut self,
402        func_address: usize,
403        return_address: usize,
404        parameters: &[usize],
405    ) -> TraceResult<(UserRegs, ProcessFrame<T>)> {
406        use std::mem::size_of;
407
408        let mut current_registers = self.query_registers()?;
409        tracing::trace!(
410            "Initial registers acquired Current PC: {:X?}",
411            current_registers.program_counter()
412        );
413
414        let cached_registers = current_registers.clone();
415        let param_count = parameters.len();
416        current_registers.set_stack_pointer(current_registers.stack_pointer() & !0xfusize);
417
418        tracing::trace!("Function parameters: {:?}", parameters);
419
420        if param_count > 0 {
421            // adjust stack pointer
422            current_registers.set_stack_pointer(
423                current_registers.stack_pointer()
424                    - (((param_count + 3) & !3usize) * size_of::<usize>()),
425            );
426            self.write_memory(
427                current_registers.stack_pointer(),
428                usize_arr_to_u8(parameters).as_slice(),
429            )?;
430        }
431
432        // return address is bottom of stack!
433        current_registers.set_stack_pointer(current_registers.stack_pointer() - size_of::<usize>());
434        self.write_memory(
435            current_registers.stack_pointer(),
436            &return_address.to_le_bytes(),
437        )?;
438
439        current_registers.eax = 0;
440        current_registers.orig_eax = 0;
441
442        // set registers cached_registers
443        current_registers.set_program_counter(func_address);
444        tracing::trace!(
445            "Executing with PC: {:X?}, and arguments {parameters:?}",
446            func_address
447        );
448
449        self.set_registers(current_registers)?;
450        tracing::trace!("Registers successfully injected.");
451
452        let mut frame = self.step_cont()?;
453        let result_regs = frame.query_registers()?;
454        tracing::trace!("Result {result_regs:#?}");
455
456        frame.set_registers(cached_registers)?;
457        Ok((result_regs, frame))
458    }
459
460    /// Attempts to invoke a remote function with the provided parameters.
461    /// Internally this is an os cfg controlled function that write the inputted parameters according to the architectures expectations.
462    /// Additionally it prepares the provided return address
463    #[cfg(target_arch = "x86_64")]
464    pub fn invoke_remote(
465        mut self,
466        func_address: usize,
467        return_address: usize,
468        parameters: &[usize],
469    ) -> TraceResult<(UserRegs, ProcessFrame<T>)> {
470        use std::mem::size_of;
471
472        const REGISTER_ARGUMENTS: usize = 6;
473        let mut current_registers = self.query_registers()?;
474        tracing::info!(
475            "Initial registers acquired Current PC: {:X?}",
476            current_registers.program_counter()
477        );
478
479        let cached_registers = current_registers.clone();
480        let param_count = parameters.len();
481        current_registers.set_stack_pointer(current_registers.stack_pointer() & !0xfusize);
482
483        // You gotta a better idea???????
484        if param_count > 0 {
485            current_registers.rdi = parameters[0] as u64;
486        }
487        if param_count > 1 {
488            current_registers.rsi = parameters[1] as u64;
489        }
490        if param_count > 2 {
491            current_registers.rdx = parameters[2] as u64;
492        }
493        if param_count > 3 {
494            current_registers.rcx = parameters[3] as u64;
495        }
496        if param_count > 4 {
497            current_registers.r8 = parameters[4] as u64;
498        }
499        if param_count > 5 {
500            current_registers.r9 = parameters[5] as u64;
501        }
502
503        if parameters.len() > REGISTER_ARGUMENTS {
504            let stack_arguments = &parameters[REGISTER_ARGUMENTS..];
505            tracing::trace!("Remaining stack arguments: {:?}", stack_arguments);
506
507            // adjust stack pointer
508            current_registers.set_stack_pointer(
509                current_registers.stack_pointer()
510                    - (((stack_arguments.len() + 1) & !1usize) * size_of::<usize>()),
511            );
512            self.write_memory(
513                current_registers.stack_pointer(),
514                usize_arr_to_u8(stack_arguments).as_slice(),
515            )?;
516        };
517
518        // return address is bottom of stack!
519        current_registers.set_stack_pointer(current_registers.stack_pointer() - size_of::<usize>());
520        self.write_memory(
521            current_registers.stack_pointer(),
522            &return_address.to_le_bytes(),
523        )?;
524
525        current_registers.rax = 0;
526        current_registers.orig_rax = 0;
527
528        // set registers cached_registers
529        current_registers.set_program_counter(func_address);
530        tracing::trace!(
531            "Executing with PC: {:X?}, and arguments {parameters:?}",
532            func_address
533        );
534
535        self.set_registers(current_registers)?;
536        tracing::trace!("Registers successfully injected.");
537
538        let mut frame = self.step_cont()?;
539        let result_regs = frame.query_registers()?;
540        tracing::info!("Result {result_regs:#?}");
541
542        frame.set_registers(cached_registers)?;
543        Ok((result_regs, frame))
544    }
545
546    /// Attempts to read a process's memory from fs
547    pub fn read_memory(&mut self, addr: usize, len: usize) -> TraceResult<Vec<u8>> {
548        let mut data = vec![0; len];
549        let len_read = self.read_memory_mut(addr, &mut data)?;
550        data.truncate(len_read);
551        Ok(data)
552    }
553
554    /// Attempts to read a mutable section of a process's memory from fs
555    pub fn read_memory_mut(&self, addr: usize, data: &mut [u8]) -> TraceResult<usize> {
556        use std::os::unix::fs::FileExt;
557        let mem = std::fs::File::open(self.process.proc_mem_path())?;
558        let len = mem.read_at(data, addr.try_into().unwrap())?;
559        Ok(len)
560    }
561
562    /// Attempts to write to a section of the process's memory
563    pub fn write_memory(&mut self, addr: usize, data: &[u8]) -> TraceResult<usize> {
564        use std::os::unix::fs::FileExt;
565        let mem = std::fs::OpenOptions::new()
566            .read(true)
567            .write(true)
568            .open(self.process.proc_mem_path())?;
569        let len = mem.write_at(data, addr.try_into().unwrap())?;
570        Ok(len)
571    }
572
573    /// Continue the process frame, consuming self.
574    fn step_cont(mut self) -> TraceResult<ProcessFrame<T>> {
575        self.process.cont()?;
576        self.process.next_frame()
577    }
578}
579
580/// Drop implementation that attempts to ptrace detach from the process.
581/// On failure there is a warning, but this does not ever panic.
582impl<T> Drop for TracedProcess<T>
583where
584    T: ProcessIdentifier,
585{
586    fn drop(&mut self) {
587        let pid = self.pid();
588        match self.detach() {
589            Ok(()) => tracing::info!("Successfully detached from Pid: {pid}"),
590            Err(e) => tracing::error!("Failed to detach from Pid: {pid}, {e:#?}"),
591        }
592    }
593}
594
595impl<T> TracedProcess<T>
596where
597    T: ProcessIdentifier,
598{
599    /// Attempt to detach from the traced process
600    fn detach(&mut self) -> TraceResult<()> {
601        let result = unsafe { ptrace(PTRACE_DETACH, self.pid(), 0, 0) };
602        match result == -1 {
603            true => Err(TraceError::Ptrace(std::io::Error::last_os_error())),
604            false => Ok(()),
605        }
606    }
607
608    /// Attempt to attach from the traced process
609    pub fn attach(process: T) -> TraceResult<Self> {
610        let result = unsafe { ptrace(PTRACE_ATTACH, process.pid(), 0, 0) };
611        match result == -1 {
612            true => Err(TraceError::Ptrace(std::io::Error::last_os_error())),
613            false => Ok(Self { process }),
614        }
615    }
616
617    /// pid of the actively traced process
618    pub fn pid(&self) -> pid_t {
619        self.process.pid()
620    }
621
622    /// continue execution
623    fn cont(&mut self) -> TraceResult<()> {
624        let result = unsafe { ptrace(PTRACE_CONT, self.pid(), 0, 0) };
625        match result == -1 {
626            true => Err(TraceError::Ptrace(std::io::Error::last_os_error())),
627            false => Ok(()),
628        }
629    }
630
631    /// wait for a status
632    fn wait(&mut self, options: WaitOptions) -> TraceResult<WaitStatus> {
633        let mut raw_status = 0;
634        let result = unsafe { libc::waitpid(self.process.pid(), &mut raw_status, options.into()) };
635        if result == -1 {
636            return Err(TraceError::Ptrace(std::io::Error::last_os_error()));
637        }
638
639        let status = WaitStatus(raw_status);
640        tracing::info!("{status:?}");
641        Ok(status)
642    }
643
644    /// wait for an untraced status consuming self and opening a process frame
645    pub fn next_frame(mut self) -> TraceResult<ProcessFrame<T>> {
646        let wait_status = self.wait(WaitOptions::Untraced)?;
647        if wait_status.is_exited() {
648            return Err(TraceError::General("Waiting stop received an exit signal"));
649        }
650
651        Ok(ProcessFrame { process: self })
652    }
653
654    /// path to the process's memory
655    fn proc_mem_path(&self) -> String {
656        format!("/proc/{}/mem", self.process.pid())
657    }
658}
ptrace_do/lib.rs

ptrace_do/
lib.rs
