From 2e392573234d0468a7268b13ad2792fc204050bc Mon Sep 17 00:00:00 2001
From: R <irgstg@gmail.com>
Date: Sun, 10 Apr 2022 12:53:55 +0100
Subject: [PATCH] PTRACE_GETREGSET and PTRACE_SETREGSET basic implementation on
 aarch64

---
 src/sys/ptrace/linux.rs | 89 ++++++++++++++++++++++++++++++++++++++++-
 test/sys/test_ptrace.rs | 41 +++++++++++++++++++
 2 files changed, 128 insertions(+), 2 deletions(-)

diff --git a/src/sys/ptrace/linux.rs b/src/sys/ptrace/linux.rs
index 24152d7d5a..be3b475645 100644
--- a/src/sys/ptrace/linux.rs
+++ b/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! {
@@ -481,3 +482,87 @@ 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)
+    }
+}
diff --git a/test/sys/test_ptrace.rs b/test/sys/test_ptrace.rs
index 89c4e2ddad..28852fa13e 100644
--- a/test/sys/test_ptrace.rs
+++ b/test/sys/test_ptrace.rs
@@ -217,3 +217,44 @@ fn test_ptrace_syscall() {
         },
     }
 }
+
+#[cfg(any(all(target_os = "linux",
+        target_arch = "aarch64", target_env = "gnu")))]
+#[test]
+fn test_ptrace_regsets() {
+    use nix::sys::signal::*;
+    use nix::sys::wait::{waitpid, WaitStatus};
+    use nix::unistd::fork;
+    use nix::sys::ptrace::{self, getregset, setregset};
+    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"),
+            }
+        },
+    }
+}