aarch64 register mappings

Author: nia-e

src/discrete_alloc.rs

@@ -42,7 +42,7 @@ impl MachineAlloc {
         let mut alloc = ALLOCATOR.lock().unwrap();
         alloc.enabled = true;
         // This needs to specifically be the system pagesize!
-        alloc.page_size = unsafe { 
+        alloc.page_size = unsafe {
             let ret = libc::sysconf(libc::_SC_PAGE_SIZE);
             if ret > 0 {
                 ret as _

src/shims/trace.rs

@@ -4,24 +4,30 @@ use nix::unistd;
 
 use crate::discrete_alloc;
 
+#[cfg(target_pointer_width = "64")]
+const BITS: u32 = 64;
+#[cfg(target_pointer_width = "32")]
+const BITS: u32 = 32;
+
+#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+const BREAKPT_INSTR: i64 = 0xCC;
+#[cfg(target_arch = "aarch64")]
+const BREAKPT_INSTR: i64 = 0xD420;
+
 // We do NOT ever want to block the child from accessing this!!
 static SUPERVISOR: std::sync::Mutex<Option<Supervisor>> = std::sync::Mutex::new(None);
 static mut PAGE_ADDR: *mut libc::c_void = std::ptr::null_mut();
 static mut PAGE_SIZE: usize = 4096;
 static mut CLICK_HERE_4_FREE_STACK: [u8; 1024] = [0; 1024];
 
-#[cfg(target_pointer_width = "64")]
-const BITS: u32 = 64;
-#[cfg(target_pointer_width = "32")]
-const BITS: u32 = 32;
-
 trait ArchIndependentRegs {
-    fn ax(&self) -> usize;
-    fn di(&self) -> usize;
-    fn si(&self) -> usize;
+    // see https://man7.org/linux/man-pages/man2/syscall.2.html
+    fn retval(&self) -> usize;
+    fn arg1(&self) -> usize;
+    fn arg2(&self) -> usize;
+    fn syscall_nr(&self) -> usize;
     fn ip(&self) -> usize;
     fn sp(&self) -> usize;
-    fn orig_ax(&self) -> usize;
     fn set_ip(&mut self, ip: usize);
     fn set_sp(&mut self, sp: usize);
 }
@@ -30,12 +36,12 @@ trait ArchIndependentRegs {
 #[allow(clippy::cast_possible_truncation)]
 #[rustfmt::skip]
 impl ArchIndependentRegs for libc::user_regs_struct {
-    fn ax(&self) -> usize { self.rax as _ }
-    fn di(&self) -> usize { self.rdi as _ }
-    fn si(&self) -> usize { self.rsi as _ }
+    fn retval(&self) -> usize { self.rax as _ }
+    fn arg1(&self) -> usize { self.rdi as _ }
+    fn arg2(&self) -> usize { self.rsi as _ }
+    fn syscall_nr(&self) -> usize { self.orig_rax as _ }
     fn ip(&self) -> usize { self.rip as _ }
     fn sp(&self) -> usize { self.rsp as _ }
-    fn orig_ax(&self) -> usize { self.orig_rax as _ }
     fn set_ip(&mut self, ip: usize) { self.rip = ip as _ }
     fn set_sp(&mut self, sp: usize) { self.rsp = sp as _ }
 }
@@ -44,16 +50,30 @@ impl ArchIndependentRegs for libc::user_regs_struct {
 #[allow(clippy::cast_possible_truncation)]
 #[rustfmt::skip]
 impl ArchIndependentRegs for libc::user_regs_struct {
-    fn ax(&self) -> usize { self.eax as _ }
-    fn di(&self) -> usize { self.edi as _ }
-    fn si(&self) -> usize { self.esi as _ }
+    fn retval(&self) -> usize { self.eax as _ }
+    fn arg1(&self) -> usize { self.edi as _ }
+    fn arg2(&self) -> usize { self.esi as _ }
+    fn syscall_nr(&self) -> usize { self.orig_eax as _ }
     fn ip(&self) -> usize { self.eip as _ }
     fn sp(&self) -> usize { self.esp as _ }
-    fn orig_ax(&self) -> usize { self.orig_eax as _ }
     fn set_ip(&mut self, ip: usize) { self.eip = ip as _ }
     fn set_sp(&mut self, sp: usize) { self.esp = sp as _ }
 }
 
+#[cfg(target_arch = "aarch64")]
+#[allow(clippy::cast_possible_truncation)]
+#[rustfmt::skip]
+impl ArchIndependentRegs for libc::user_regs_struct {
+    fn retval(&self) -> usize { self.regs[0] as _ }
+    fn arg1(&self) -> usize { self.regs[0] as _ }
+    fn arg2(&self) -> usize { self.regs[1] as _ }
+    fn syscall_nr(&self) -> usize { self.regs[8] as _ }
+    fn ip(&self) -> usize { self.pc as _ }
+    fn sp(&self) -> usize { self.sp as _ }
+    fn set_ip(&mut self, ip: usize) { self.pc = ip as _ }
+    fn set_sp(&mut self, sp: usize) { self.sp = sp as _ }
+}
+
 pub struct Supervisor {
     t_message: ipc::IpcSender<TraceRequest>,
     r_event: ipc::IpcReceiver<MemEvents>,
@@ -295,16 +315,16 @@ fn sv_loop(listener: ChildListener, t_event: ipc::IpcSender<MemEvents>) -> ! {
                     }
                     wait::WaitStatus::PtraceSyscall(pid) => {
                         let regs = ptrace::getregs(pid).unwrap();
-                        let syscall_nr = regs.orig_ax();
+                        let syscall_nr = regs.syscall_nr();
                         if enter {
                             if syscall_nr as i64 == libc::SYS_mmap {
-                                let len = regs.si();
+                                let len = regs.arg2();
                                 flush_mapped = Some(len);
                             } else if syscall_nr as i64 == libc::SYS_munmap {
                                 // The unmap call might hit multiple mappings we've saved,
                                 // or overlap with them partially (or both)
-                                let um_start = regs.di();
-                                let um_len = regs.si();
+                                let um_start = regs.arg1();
+                                let um_len = regs.arg2();
                                 let um_end = um_start + um_len;
                                 let mut idxes = vec![];
                                 for (idx, &(mp_start, len)) in mappings.iter().enumerate() {
@@ -326,9 +346,9 @@ fn sv_loop(listener: ChildListener, t_event: ipc::IpcSender<MemEvents>) -> ! {
                             } // TODO: handle brk/sbrk
                         } else {
                             if syscall_nr as i64 == libc::SYS_mmap {
-                                if regs.ax() as isize > 0 {
+                                if regs.retval() as isize > 0 {
                                     if let Some(len) = flush_mapped.take() {
-                                        let addr = regs.ax();
+                                        let addr = regs.retval();
                                         mappings.push((addr, len as _));
                                     } else {
                                         eprintln!(
@@ -338,7 +358,7 @@ fn sv_loop(listener: ChildListener, t_event: ipc::IpcSender<MemEvents>) -> ! {
                                 }
                             } else if syscall_nr as i64 == libc::SYS_munmap {
                                 if let Some((idxes, um_start, um_len)) = munmap_args.take() {
-                                    if regs.ax() as isize > 0 {
+                                    if regs.retval() as isize > 0 {
                                         // Unmap succeeded, so take out the mapping(s) from our list
                                         // but it may be only partial so we may readd some sections
                                         for idx in idxes {
@@ -373,9 +393,9 @@ fn sv_loop(listener: ChildListener, t_event: ipc::IpcSender<MemEvents>) -> ! {
                         ptrace::seize(main_pid, ptrace::Options::PTRACE_O_TRACESYSGOOD).unwrap();
                         let _ = wait::waitpid(main_pid, None).unwrap();
 
-                        ptrace::write(main_pid, malloc_addr as _, 0xcc).unwrap();
-                        ptrace::write(main_pid, realloc_addr as _, 0xcc).unwrap();
-                        ptrace::write(main_pid, free_addr as _, 0xcc).unwrap();
+                        ptrace::write(main_pid, malloc_addr as _, BREAKPT_INSTR).unwrap();
+                        ptrace::write(main_pid, realloc_addr as _, BREAKPT_INSTR).unwrap();
+                        ptrace::write(main_pid, free_addr as _, BREAKPT_INSTR).unwrap();
 
                         ptrace::syscall(main_pid, None).unwrap();
                     }
@@ -515,13 +535,13 @@ fn handle_sigtrap(
     let regs = ptrace::getregs(pid).unwrap();
     match regs.ip() - 1 {
         a if a == malloc_addr => {
-            let size = regs.di(); // !
+            let size = regs.arg1(); // !
             let ptr = intercept_retptr(pid, regs, malloc_addr, malloc_bytes);
             mappings.push((ptr as _, size));
         }
         a if a == realloc_addr => {
-            let old_ptr = regs.di();
-            let size = regs.si();
+            let old_ptr = regs.arg1();
+            let size = regs.arg2();
             let pos = mappings
                 .iter()
                 .position(|&(ptr, size)| ptr <= old_ptr as _ && (old_ptr as usize) < ptr + size);
@@ -532,7 +552,7 @@ fn handle_sigtrap(
             mappings.push((ptr as _, size as _));
         }
         a if a == free_addr => {
-            let old_ptr = regs.di();
+            let old_ptr = regs.arg1();
             //let size = regs.rdi;
             let pos =
                 mappings.iter().position(|&(ptr, size)| ptr <= old_ptr && old_ptr < ptr + size);
@@ -565,7 +585,7 @@ fn intercept_retptr(
     regs.set_ip(regs.ip() - 1);
     let ret_addr = ptrace::read(pid, regs.sp() as _).unwrap();
     let ret_bytes = ptrace::read(pid, ret_addr as _).unwrap();
-    ptrace::write(pid, ret_addr as _, 0xcc).unwrap();
+    ptrace::write(pid, ret_addr as _, BREAKPT_INSTR).unwrap();
     ptrace::write(pid, fn_addr as _, fn_bytes).unwrap();
     ptrace::setregs(pid, regs).unwrap();
 
@@ -574,9 +594,9 @@ fn intercept_retptr(
 
     // now we're getting the return hopefully
     let mut regs = ptrace::getregs(pid).unwrap();
-    let ptr = regs.ax(); // !
+    let ptr = regs.retval(); // !
     regs.set_ip(regs.ip() - 1);
-    ptrace::write(pid, fn_addr as _, 0xcc).unwrap();
+    ptrace::write(pid, fn_addr as _, BREAKPT_INSTR).unwrap();
     ptrace::write(pid, ret_addr as _, ret_bytes).unwrap();
     ptrace::setregs(pid, regs).unwrap();