Add initial arm64 support

.cargo/config.toml

@@ -4,5 +4,8 @@ rustflags = ["-Cforce-frame-pointers=yes"]
 [target.x86_64-unknown-linux-gnu]
 runner = "sudo -E"
 
+[target.aarch64-unknown-linux-gnu]
+runner = "sudo -E"
+
 [alias]
 xtask = "run --package xtask --"

.github/workflows/build.yml

@@ -18,7 +18,7 @@
     steps:
     - uses: actions/checkout@main
     - uses: DeterminateSystems/nix-installer-action@main
     - uses: DeterminateSystems/magic-nix-cache-action@main
     - name: Set up nix dev env
       run: nix develop --command echo 0
     - name: Run `cargo check`
@@ -42,8 +42,10 @@
     steps:
       - uses: actions/checkout@main
       - uses: DeterminateSystems/nix-installer-action@main
       - uses: DeterminateSystems/magic-nix-cache-action@main
       - name: Set up nix dev env
         run: nix develop --command echo 0
       - name: Run `cargo check`
-        run: nix develop --ignore-environment --command cargo check
+        run: nix develop --ignore-environment --command cargo check
+      - name: Run `cargo test`
+        run: nix develop --command cargo test --workspace

src/bpf/profiler.bpf.c

@@ -209,6 +209,25 @@ static __always_inline void send_event(Event *event, struct bpf_perf_event_data
   bpf_map_update_elem(&rate_limits, event, &rate_limited, BPF_ANY);
 }
 
+#ifdef __TARGET_ARCH_x86
+static __always_inline u64 remove_pac(u64 addr) {
+  return addr;
+}
+#endif
+
+#ifdef __TARGET_ARCH_arm64
+// Arm64 supports pointer authentication, we need to remove the signatured during
+// unwinding.
+static __always_inline u64 remove_pac(u64 addr) {
+  // The signature is stored in the top 55 - virtual address size bits [0], which
+  // is typically 48 bytes, hence we need to clear the top 7 bits. Clearing 8 bits
+  // as they are all the non-addressable anyways.
+  // - [0]: https://docs.kernel.org/arch/arm64/pointer-authentication.html#basic-support
+  addr &= 0x0000FFFFFFFFFFFF;
+  return addr;
+}
+#endif
+
 // Kernel addresses have the top bits set.
 static __always_inline bool in_kernel(u64 ip) { return ip & (1UL << 63); }
 
@@ -234,8 +253,8 @@ static __always_inline bool is_kthread() {
 
 // avoid R0 invalid mem access 'scalar'
 // Port of `task_pt_regs` in BPF.
-static __always_inline bool retrieve_task_registers(u64 *ip, u64 *sp, u64 *bp) {
-  if (ip == NULL || sp == NULL || bp == NULL) {
+static __always_inline bool retrieve_task_registers(u64 *ip, u64 *sp, u64 *bp, u64 *lr) {
+  if (ip == NULL || sp == NULL || bp == NULL || lr == NULL) {
     return false;
   }
 
@@ -258,12 +277,14 @@ static __always_inline bool retrieve_task_registers(u64 *ip, u64 *sp, u64 *bp) {
   }
 
   void *ptr = stack + THREAD_SIZE - TOP_OF_KERNEL_STACK_PADDING;
-  bpf_user_pt_regs_t *regs = ((bpf_user_pt_regs_t *)ptr) - 1;
+  struct pt_regs *regs = ((struct pt_regs *)ptr) - 1;
 
   *ip = PT_REGS_IP_CORE(regs);
   *sp = PT_REGS_SP_CORE(regs);
   *bp = PT_REGS_FP_CORE(regs);
-
+#ifdef __TARGET_ARCH_arm64
+  *lr = PT_REGS_RET_CORE(regs);
+#endif
   return true;
 }
 
@@ -525,34 +546,13 @@ int dwarf_unwind(struct bpf_perf_event_data *ctx) {
       return 1;
     }
 
-    // HACK(javierhonduco): This is an architectural shortcut we can take. As we
-    // only support x86_64 at the minute, we can assume that the return address
-    // is *always* 8 bytes ahead of the previous stack pointer.
-    u64 previous_rip_addr =
-        previous_rsp - 8; // the saved return address is 8 bytes ahead of the
-                          // previous stack pointer
-    u64 previous_rip = 0;
-    int err =
-        bpf_probe_read_user(&previous_rip, 8, (void *)(previous_rip_addr));
-
-    if (previous_rip == 0) {
-      if (err == 0) {
-        LOG("[warn] previous_rip=0, maybe this is a JIT segment?");
-      } else {
-        LOG("[error] previous_rip should not be zero. This can mean that the "
-            "read failed, ret=%d while reading @ %llx.",
-            err, previous_rip_addr);
-        bump_unwind_error_previous_rip_zero();
-      }
-      return 1;
-    }
-
     // Set rbp register.
     u64 previous_rbp = 0;
+    u64 previous_rbp_addr = previous_rsp + found_rbp_offset;
+
     if (found_rbp_type == RBP_TYPE_UNCHANGED) {
       previous_rbp = unwind_state->bp;
     } else {
-      u64 previous_rbp_addr = previous_rsp + found_rbp_offset;
       LOG("\t(bp_offset: %d, bp value stored at %llx)", found_rbp_offset,
           previous_rbp_addr);
       int ret =
@@ -566,10 +566,45 @@ int dwarf_unwind(struct bpf_perf_event_data *ctx) {
       }
     }
 
+    u64 previous_rip = 0;
+    u64 previous_rip_addr = 0;
+
+#ifdef __TARGET_ARCH_x86
+    // The return address is guaranteed to be 8 bytes ahead of
+    // the previous stack pointer in x86_64.
+    previous_rip_addr = previous_rsp - 8;
+#endif
+
+#ifdef __TARGET_ARCH_arm64
+    // Special handling for leaf frame.
+    if (unwind_state->stack.len == 0) {
+      previous_rip = unwind_state->lr;
+    } else {
+      // This is guaranteed by the Aarch64 ABI.
+      previous_rip_addr = previous_rbp_addr + 8;
+    }
+#endif
+
+    int err =
+        bpf_probe_read_user(&previous_rip, 8, (void *)(previous_rip_addr));
+
+    if (previous_rip == 0) {
+      if (err == 0) {
+        LOG("[warn] previous_rip=0, maybe this is a JIT segment?");
+      } else {
+        LOG("[error] previous_rip should not be zero. This can mean that the "
+            "read failed, ret=%d while reading @ %llx.",
+            err, previous_rip_addr);
+        bump_unwind_error_previous_rip_zero();
+      }
+      return 1;
+    }
+
+
     LOG("\tprevious ip: %llx (@ %llx)", previous_rip, previous_rip_addr);
     LOG("\tprevious sp: %llx", previous_rsp);
     // Set rsp and rip registers
-    unwind_state->ip = previous_rip;
+    unwind_state->ip = remove_pac(previous_rip);
     unwind_state->sp = previous_rsp;
     // Set rbp
     LOG("\tprevious bp: %llx", previous_rbp);
@@ -627,7 +662,7 @@ static __always_inline bool set_initial_state(unwind_state_t *unwind_state, bpf_
   unwind_state->stack_key.kernel_stack_id = 0;
 
   if (in_kernel(PT_REGS_IP(regs))) {
-    if (!retrieve_task_registers(&unwind_state->ip, &unwind_state->sp, &unwind_state->bp)) {
+    if (!retrieve_task_registers(&unwind_state->ip, &unwind_state->sp, &unwind_state->bp, &unwind_state->lr)) {
       // in kernelspace, but failed, probs a kworker
       // todo: bump counter
       return false;
@@ -637,6 +672,7 @@ static __always_inline bool set_initial_state(unwind_state_t *unwind_state, bpf_
     unwind_state->ip = PT_REGS_IP(regs);
     unwind_state->sp = PT_REGS_SP(regs);
     unwind_state->bp = PT_REGS_FP(regs);
+    unwind_state->lr = remove_pac(PT_REGS_RET(regs));
   }
 
   return true;

src/bpf/profiler.h

@@ -177,6 +177,7 @@ typedef struct {
   unsigned long long ip;
   unsigned long long sp;
   unsigned long long bp;
+  unsigned long long lr;
   int tail_calls;
 
   stack_count_key_t stack_key;

src/process.rs

@@ -117,6 +117,7 @@ pub struct ObjectFileInfo {
     pub is_dyn: bool,
     pub references: i64,
     pub native_unwind_info_size: Option<u64>,
+    pub is_vdso: bool,
 }
 
 impl Clone for ObjectFileInfo {
@@ -128,6 +129,7 @@ impl Clone for ObjectFileInfo {
             is_dyn: self.is_dyn,
             references: self.references,
             native_unwind_info_size: self.native_unwind_info_size,
+            is_vdso: self.is_vdso,
         }
     }
 }
@@ -202,6 +204,7 @@ mod tests {
             is_dyn: false,
             references: 1,
             native_unwind_info_size: None,
+            is_vdso: false,
         };
 
         remove_file(file_path).unwrap();
@@ -221,6 +224,7 @@ mod tests {
             is_dyn: false,
             references: 0,
             native_unwind_info_size: None,
+            is_vdso: false,
         };
 
         let mapping = ExecutableMapping {

src/profiler.rs

@@ -45,7 +45,8 @@ use crate::profile::*;
 use crate::unwind_info::log_unwind_info_sections;
 use crate::unwind_info::manager::UnwindInfoManager;
 use crate::unwind_info::types::CompactUnwindRow;
-use crate::util::{get_online_cpus, summarize_address_range};
+use crate::util::Architecture;
+use crate::util::{architecture, get_online_cpus, summarize_address_range};
 use lightswitch_object::{ExecutableId, ObjectFile};
 
 pub enum TracerEvent {
@@ -1275,8 +1276,14 @@ impl Profiler {
         let executable_info = object_files.get(&executable_id).unwrap();
         let executable_path_open = executable_info.open_file_path();
         let executable_path = executable_info.path.to_string_lossy().to_string();
+        let needs_synthesis = executable_info.is_vdso && architecture() == Architecture::Arm64;
         std::mem::drop(object_files);
 
+        if needs_synthesis {
+            debug!("arm64 vDSO don't typically contain unwind information and synthesising it is not implemented yet");
+            return;
+        }
+
         let span = span!(
             Level::DEBUG,
             "calling in_memory_unwind_info",
@@ -1652,6 +1659,7 @@ impl Profiler {
                                     is_dyn: object_file.is_dynamic(),
                                     references: 1,
                                     native_unwind_info_size: None,
+                                    is_vdso: false,
                                 });
                             }
                             Err(e) => {
@@ -1715,6 +1723,7 @@ impl Profiler {
                                 is_dyn: object_file.is_dynamic(),
                                 references: 1,
                                 native_unwind_info_size: None,
+                                is_vdso: true,
                             },
                         );
                         mappings.push(ExecutableMapping {

src/unwind_info/convert.rs

@@ -3,6 +3,7 @@ use std::fs::File;
 use anyhow::Result;
 use gimli::{CfaRule, CieOrFde, EhFrame, UnwindContext, UnwindSection};
 use memmap2::Mmap;
+use object::Architecture;
 use object::{Object, ObjectSection};
 use thiserror::Error;
 use tracing::{debug, error, span, Level};
@@ -76,12 +77,26 @@ impl<'a> CompactUnwindInfoBuilder<'a> {
 
         let eh_frame_data = &eh_frame_section.uncompressed_data()?;
 
-        let eh_frame = EhFrame::new(eh_frame_data, endian);
+        let mut eh_frame = EhFrame::new(eh_frame_data, endian);
+        if object_file.architecture() == Architecture::Aarch64 {
+            eh_frame.set_vendor(gimli::Vendor::AArch64);
+        }
         let mut entries_iter = eh_frame.entries(&bases);
 
         let mut cur_cie = None;
         let mut pc_and_fde_offset = Vec::new();
 
+        let frame_pointer = if object_file.architecture() == Architecture::Aarch64 {
+            ARM64_FP
+        } else {
+            X86_FP
+        };
+        let stack_pointer = if object_file.architecture() == Architecture::Aarch64 {
+            ARM64_SP
+        } else {
+            X86_SP
+        };
+
         while let Ok(Some(entry)) = entries_iter.next() {
             match entry {
                 CieOrFde::Cie(cie) => {
@@ -137,9 +152,9 @@ impl<'a> CompactUnwindInfoBuilder<'a> {
                         compact_row.pc = row.start_address();
                         match row.cfa() {
                             CfaRule::RegisterAndOffset { register, offset } => {
-                                if register == &RBP_X86 {
+                                if register == &frame_pointer {
                                     compact_row.cfa_type = CfaType::FramePointerOffset;
-                                } else if register == &RSP_X86 {
+                                } else if register == &stack_pointer {
                                     compact_row.cfa_type = CfaType::StackPointerOffset;
                                 } else {
                                     compact_row.cfa_type = CfaType::UnsupportedRegisterOffset;
@@ -171,7 +186,7 @@ impl<'a> CompactUnwindInfoBuilder<'a> {
                             }
                         };
 
-                        match row.register(RBP_X86) {
+                        match row.register(frame_pointer) {
                             gimli::RegisterRule::Undefined => {}
                             gimli::RegisterRule::Offset(offset) => {
                                 compact_row.rbp_type = RbpType::CfaOffset;

src/unwind_info/types.rs

@@ -85,5 +85,11 @@ lazy_static! {
     ].map(|a| a.0);
 }
 
-pub const RBP_X86: gimli::Register = gimli::Register(6);
-pub const RSP_X86: gimli::Register = gimli::Register(7);
+// Source: https://gitlab.com/x86-psABIs/x86-64-ABI/-/jobs/artifacts/d725a372/raw/x86-64-ABI/abi.pdf?job=build
+// > Figure 3.36: DWARF Register Number Mapping
+pub const X86_FP: gimli::Register = gimli::Register(6); // Frame Pointer ($rbp)
+pub const X86_SP: gimli::Register = gimli::Register(7); // Stack Pointer ($rsp)
+
+// Source: https://github.com/ARM-software/abi-aa/blob/05abf4f7/aadwarf64/aadwarf64.rst#41dwarf-register-names
+pub const ARM64_FP: gimli::Register = gimli::Register(29); // Frame Pointer (x29)
+pub const ARM64_SP: gimli::Register = gimli::Register(31); // Stack Pointer (sp)

src/util/arch.rs

@@ -0,0 +1,15 @@
+#[derive(PartialEq)]
+pub enum Architecture {
+    Arm64,
+    X86,
+}
+
+#[cfg(target_arch = "aarch64")]
+pub fn architecture() -> Architecture {
+    Architecture::Arm64
+}
+
+#[cfg(target_arch = "x86_64")]
+pub fn architecture() -> Architecture {
+    Architecture::X86
+}

src/util/mod.rs

@@ -1,6 +1,9 @@
+mod arch;
 mod cpu;
 mod lpm;
 
+pub use arch::architecture;
+pub use arch::Architecture;
 pub use cpu::get_online_cpus;
 pub use lpm::summarize_address_range;
 pub use lpm::AddressBlockRange;