Add initial arm64 support

This commit adds support for the arm64 (aarch64) architecture. This is
achived by adapting the compact unwind information generation to deal
with the DWARF registers that this architecture maps to and by modifying
the unwinder.

These changes support PAC (Pointer Authentication Code), see test plan
for more details.

While working on this feature I realised that the vDSO included in
Ubuntu 24.10 (6.11.0-12-generic) does not contain unwind information. In
the future we should synthetise it which should be easy enough as it
seems to be compiled with frame pointers.

Test Plan
=========

Ran tests locally on the arm64 machine mentioned above, and the
PAC-enabled binary (see tests/testprogs/). The tests will also be run
for this arch.

```
2025-01-22T12:59:32.726079Z  INFO lightswitch::profiler: unwinder stats: unwinder_stats_t { total: 44, success_dwarf: 44, error_truncated: 0, error_unsupported_expression: 0, error_unsupported_frame_pointer_action: 0, error_unsupported_cfa_register: 0, error_previous_rsp_zero: 0, error_previous_rip_zero: 0, error_previous_rbp_zero: 0, error_should_never_happen: 0, error_mapping_not_found: 0, error_mapping_does_not_contain_pc: 0, error_chunk_not_found: 0, error_binary_search_exhausted_iterations: 0, error_sending_new_process_event: 0, error_cfa_offset_did_not_fit: 0, error_rbp_offset_did_not_fit: 0, bp_non_zero_for_bottom_frame: 0, vdso_encountered: 0, jit_encountered: 0 }
```

.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

@@ -46,4 +46,6 @@ jobs:
       - 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;