PTRACE_GETREGSET and PTRACE_SETREGSET basic implementation on aarch64

src/sys/ptrace/linux.rs

@@ -14,8 +14,9 @@ pub type AddressType = *mut ::libc::c_void;
     target_os = "linux",
     any(all(target_arch = "x86_64",
             any(target_env = "gnu", target_env = "musl")),
-        all(target_arch = "x86", target_env = "gnu"))
-))]
+        all(target_arch = "x86", target_env = "gnu"),
+        all(target_arch = "aarch64", target_os = "linux"),
+)))]
 use libc::user_regs_struct;
 
 cfg_if! {
@@ -482,6 +483,90 @@ pub unsafe fn write(
     ptrace_other(Request::PTRACE_POKEDATA, pid, addr, data).map(drop)
 }
 
+/// Read the tracee's registers.
+///
+/// as with `ptrace(PTRACE_GETREGSET, ...)`
+///
+/// # Arguments
+///
+/// * `pid` - tracee's `nix::unistd::Pid`
+#[cfg(any(all(target_os = "linux", target_env = "gnu", target_arch = "aarch64")))]
+pub fn getregset(pid: Pid) -> Result<user_regs_struct> {
+    ptrace_get_iovec_data::<user_regs_struct>(
+        Request::PTRACE_GETREGSET,
+        pid,
+        libc::NT_PRSTATUS,
+    )
+}
+
+/// Modify the tracee's registers.
+///
+/// as with `ptrace(PTRACE_SETREGSET, ...)`
+///
+/// # Arguments
+///
+/// * `pid` - tracee's `nix::unistd::Pid`
+///
+/// * `regs` - `libc::user_regs_struct` to set
+#[cfg(any(all(target_os = "linux", target_env = "gnu", target_arch = "aarch64")))]
+pub fn setregset(pid: Pid, regs: user_regs_struct) -> Result<()> {
+    ptrace_set_iovec_data(
+        Request::PTRACE_SETREGSET,
+        pid,
+        libc::NT_PRSTATUS,
+        regs,
+    )
+}
+
+/// As with `ptrace_get_data` but with an `iovec`
+#[cfg(any(all(target_os = "linux", target_env = "gnu", target_arch = "aarch64")))]
+fn ptrace_get_iovec_data<T>(
+    request: Request,
+    pid: Pid,
+    nt_req: libc::c_int,
+) -> Result<T> {
+    let mut data = mem::MaybeUninit::<T>::uninit();
+    let mut iov = libc::iovec {
+        iov_base: data.as_mut_ptr() as *mut c_void,
+        iov_len: mem::size_of::<T>(),
+    };
+
+    let res = unsafe {
+        libc::ptrace(
+            request as RequestType,
+            pid,
+            nt_req as AddressType,
+            &mut iov as *mut _ as *mut c_void,
+        )
+    };
+
+    Errno::result(res)?;
+    Ok(unsafe { data.assume_init() })
+}
+
+#[cfg(any(all(target_os = "linux", target_env = "gnu", target_arch = "aarch64")))]
+fn ptrace_set_iovec_data<T>(
+    request: Request,
+    pid: Pid,
+    nt_req: libc::c_int,
+    data: T,
+) -> Result<()> {
+    let iov = libc::iovec {
+        iov_base: &data as *const _ as *mut c_void,
+        iov_len: mem::size_of::<T>(),
+    };
+
+    unsafe {
+        ptrace_other(
+            request,
+            pid,
+            nt_req as AddressType,
+            &iov as *const _ as *mut c_void,
+        )
+        .map(drop)
+    }
+}
+
 /// Reads a word from a user area at `offset`.
 /// The user struct definition can be found in `/usr/include/sys/user.h`.
 pub fn read_user(pid: Pid, offset: AddressType) -> Result<c_long> {

test/sys/test_ptrace.rs

@@ -273,3 +273,49 @@ fn test_ptrace_syscall() {
         }
     }
 }
+
+#[cfg(any(all(
+    target_os = "linux",
+    target_arch = "aarch64",
+    target_env = "gnu"
+)))]
+#[test]
+fn test_ptrace_regsets() {
+    use nix::sys::ptrace::{self, getregset, setregset};
+    use nix::sys::signal::*;
+    use nix::sys::wait::{waitpid, WaitStatus};
+    use nix::unistd::fork;
+    use nix::unistd::ForkResult::*;
+
+    require_capability!("test_ptrace_regsets", CAP_SYS_PTRACE);
+
+    let _m = crate::FORK_MTX.lock();
+
+    match unsafe { fork() }.expect("Error: Fork Failed") {
+        Child => {
+            ptrace::traceme().unwrap();
+            // As recommended by ptrace(2), raise SIGTRAP to pause the child
+            // until the parent is ready to continue
+            loop {
+                raise(Signal::SIGTRAP).unwrap();
+            }
+        }
+
+        Parent { child } => {
+            assert_eq!(
+                waitpid(child, None),
+                Ok(WaitStatus::Stopped(child, Signal::SIGTRAP))
+            );
+            let mut regstruct = getregset(child).unwrap();
+            regstruct.regs[16] = 0xdeadbeef;
+            let _ = setregset(child, regstruct);
+            assert_eq!(0xdeadbeef, getregset(child).unwrap().regs[16]);
+            ptrace::cont(child, Some(Signal::SIGKILL)).unwrap();
+            match waitpid(child, None) {
+                Ok(WaitStatus::Signaled(pid, Signal::SIGKILL, _))
+                    if pid == child => {}
+                _ => panic!("The process should have been killed"),
+            }
+        }
+    }
+}