wip arm64 support

.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+    "files.associations": {
+        "vmlinux_arm64.h": "c"
+    }
+}

src/bpf/profiler.bpf.c

@@ -170,6 +170,15 @@ static __always_inline void event_new_process(struct bpf_perf_event_data *ctx, i
   bpf_map_update_elem(&rate_limits, &event, &rate_limited, BPF_ANY);
 }
 
+#ifdef __TARGET_ARCH_arm64
+// Arm64 supports pointer authentication, we need to remove the signatured during
+// unwinding.
+static __always_inline u64 remove_pac(u64 addr) {
+  //addr &= 0x0000FFFFFFFFFFFF; // do not harcode
+  return addr;
+}
+#endif
+
 // Kernel addresses have the top bits set.
 static __always_inline bool in_kernel(u64 ip) { return ip & (1UL << 63); }
 
@@ -195,8 +204,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;
   }
 
@@ -208,6 +217,8 @@ static __always_inline bool retrieve_task_registers(u64 *ip, u64 *sp, u64 *bp) {
     return false;
   }
 
+
+ /// is this needed?
   if (is_kthread()) {
     return false;
   }
@@ -219,11 +230,12 @@ 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);
+  *lr = remove_pac(PT_REGS_RET_CORE(regs));
 
   return true;
 }
@@ -412,6 +424,11 @@ int dwarf_unwind(struct bpf_perf_event_data *ctx) {
       break;
     }
 
+    if (found_rbp_type == RBP_TYPE_OFFSET_DID_NOT_FIT) {
+      // bump counter
+      return 1;
+    }
+
     // Add address to stack.
     u64 len = unwind_state->stack.len;
     // Appease the verifier.
@@ -470,15 +487,53 @@ int dwarf_unwind(struct bpf_perf_event_data *ctx) {
       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 {
+      LOG("\t(bp_offset: %d, bp value stored at %llx)", found_rbp_offset,
+          previous_rbp_addr);
+      int ret =
+          bpf_probe_read_user(&previous_rbp, 8, (void *)(previous_rbp_addr));
+      if (ret != 0) {
+        LOG("[error] previous_rbp should not be zero. This can mean "
+            "that the read has failed %d.",
+            ret);
+        bump_unwind_error_catchall();
+        return 1;
+      }
+    }
+
+
+    // Recover the return address.
+    u64 previous_rip = 0;
+    u64 previous_rip_addr = 0;
+    int err = 0;
+
+#ifdef __TARGET_ARCH_arm64
+    // Special handling for leaf frame.
+    if (unwind_state->stack.len == 0) {
+      previous_rip = unwind_state->lr;
+    } else {
+      previous_rip_addr = previous_rbp_addr + 8;  // the saved return address is 8 bytes ahead of the
+                                                      // previous frame pointer
+      err = bpf_probe_read_user(&previous_rip, 8, (void *)(previous_rip_addr));
+    }
+#endif
+#ifdef __TARGET_ARCH_x86
     // 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_rip_addr =
         previous_rsp - 8; // the saved return address is 8 bytes ahead of the
                           // previous stack pointer
     u64 previous_rip = 0;
-    int err =
+    err =
         bpf_probe_read_user(&previous_rip, 8, (void *)(previous_rip_addr));
+#endif
 
     if (previous_rip == 0) {
       if (err == 0) {
@@ -492,29 +547,10 @@ int dwarf_unwind(struct bpf_perf_event_data *ctx) {
       return 1;
     }
 
-    // Set rbp register.
-    u64 previous_rbp = 0;
-    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 =
-          bpf_probe_read_user(&previous_rbp, 8, (void *)(previous_rbp_addr));
-      if (ret != 0) {
-        LOG("[error] previous_rbp should not be zero. This can mean "
-            "that the read has failed %d.",
-            ret);
-        bump_unwind_error_catchall();
-        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);
@@ -570,16 +606,17 @@ 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;
     }
   } else {
     // Currently executing userspace code.
-    unwind_state->ip = PT_REGS_IP(regs);
+    unwind_state->ip = remove_pac(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

@@ -53,6 +53,7 @@ _Static_assert(1 << MAX_BINARY_SEARCH_DEPTH >= MAX_UNWIND_TABLE_SIZE,
 #define RBP_TYPE_EXPRESSION 3
 // Special values.
 #define RBP_TYPE_UNDEFINED_RETURN_ADDRESS 4
+#define RBP_TYPE_OFFSET_DID_NOT_FIT 5
 
 // Binary search error codes.
 #define BINARY_SEARCH_DEFAULT 0xFABADAFABADAULL
@@ -178,6 +179,8 @@ 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/unwind_info.rs

@@ -6,10 +6,12 @@ use object::{Object, ObjectSection, Section};
 use std::fs::File;
 use std::path::PathBuf;
 use thiserror::Error;
+use tracing::span;
 
 use crate::bpf::profiler_bindings::stack_unwind_row_t;
 use anyhow::anyhow;
-use tracing::{error, span, Level};
+use tracing::error;
+use tracing::Level;
 
 #[repr(u8)]
 pub enum CfaType {
@@ -29,6 +31,7 @@ enum RbpType {
     Register = 2,
     Expression = 3,
     UndefinedReturnAddress = 4,
+    OffsetDidNotFit = 5,
 }
 
 #[repr(u16)]
@@ -88,8 +91,12 @@ pub enum Error {
     ErrorNoTextSection,
 }
 
-const RBP_X86: gimli::Register = gimli::Register(6);
-const RSP_X86: gimli::Register = gimli::Register(7);
+/* const RBP_X86: gimli::Register = gimli::Register(6);
+const RSP_X86: gimli::Register = gimli::Register(7); */
+
+const ARM64_FP: gimli::Register = gimli::Register(29);
+const ARM64_SP: gimli::Register = gimli::Register(31);
+
 
 pub fn end_of_function_marker(last_addr: u64) -> CompactUnwindRow {
     CompactUnwindRow {
@@ -309,9 +316,9 @@ impl<'a> UnwindInfoBuilder<'a> {
                         compact_row.pc = row.start_address();
                         match row.cfa() {
                             CfaRule::RegisterAndOffset { register, offset } => {
-                                if register == &RBP_X86 {
+                                if register == &ARM64_FP {
                                     compact_row.cfa_type = CfaType::FramePointerOffset as u8;
-                                } else if register == &RSP_X86 {
+                                } else if register == &ARM64_SP {
                                     compact_row.cfa_type = CfaType::StackPointerOffset as u8;
                                 } else {
                                     compact_row.cfa_type = CfaType::UnsupportedRegisterOffset as u8;
@@ -339,12 +346,19 @@ impl<'a> UnwindInfoBuilder<'a> {
                             }
                         };
 
-                        match row.register(RBP_X86) {
+                        match row.register(ARM64_FP) {
                             gimli::RegisterRule::Undefined => {}
                             gimli::RegisterRule::Offset(offset) => {
                                 compact_row.rbp_type = RbpType::CfaOffset as u8;
-                                compact_row.rbp_offset =
-                                    i16::try_from(offset).expect("convert rbp offset");
+
+                                match i16::try_from(offset) {
+                                    Ok(off) => {
+                                        compact_row.rbp_offset = off;
+                                    }
+                                    Err(_) => {
+                                        compact_row.rbp_type = RbpType::OffsetDidNotFit as u8;
+                                    }
+                                }
                             }
                             gimli::RegisterRule::Register(_reg) => {
                                 compact_row.rbp_type = RbpType::Register as u8;