Add ARM64 Support

Ported from odigos-io/opentelemetry-go-instrumentation#53

Co-authored-by: Eden Federman <eden@keyval.dev>

.github/workflows/release.yaml

@@ -42,4 +42,4 @@ jobs:
           push: true
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
-          platforms: linux/amd64
+          platforms: linux/amd64,linux/arm64

Dockerfile

@@ -1,10 +1,11 @@
 FROM fedora:37 as builder
+ARG TARGETARCH
 RUN dnf install clang llvm make libbpf-devel -y
-RUN curl -LO https://go.dev/dl/go1.20.linux-amd64.tar.gz && tar -C /usr/local -xzf go*.linux-amd64.tar.gz
+RUN curl -LO https://go.dev/dl/go1.20.linux-$TARGETARCH.tar.gz && tar -C /usr/local -xzf go*.linux-$TARGETARCH.tar.gz
 ENV PATH="/usr/local/go/bin:${PATH}"
 WORKDIR /app
 COPY . .
-RUN make build
+RUN TARGET=$TARGETARCH make build
 
 FROM registry.fedoraproject.org/fedora-minimal:37
 COPY --from=builder /app/otel-go-instrumentation /

Makefile

@@ -4,7 +4,7 @@ REPODIR := $(shell dirname $(realpath $(firstword $(MAKEFILE_LIST))))
 
 # Build the list of include directories to compile the bpf program
 BPF_INCLUDE += -I${REPODIR}/include/libbpf
-BPF_INCLUDE+= -I${REPODIR}/include
+BPF_INCLUDE += -I${REPODIR}/include
 
 # Tools
 TOOLS_MOD_DIR := ./internal/tools
@@ -30,7 +30,7 @@ generate:
 
 .PHONY: build
 build: generate
-	GOOS=linux GOARCH=amd64 go build -o otel-go-instrumentation cli/main.go
+	GOOS=linux go build -o otel-go-instrumentation cli/main.go
 
 .PHONY: docker-build
 docker-build:

include/alloc.h

@@ -52,7 +52,7 @@ static __always_inline u64 get_area_end(u64 start)
 {
     s64 partition_size = (end_addr - start_addr) / total_cpus;
     s32 end_index = 1;
-    u64 *end = (u64 *)bpf_map_lookup_elem(&alloc_map, &end_index);
+    u64* end = (u64*)bpf_map_lookup_elem(&alloc_map, &end_index);
     if (end == NULL || *end == 0)
     {
         u64 current_end_addr = start + partition_size;
@@ -88,6 +88,12 @@ static __always_inline void *write_target_data(void *data, s32 size)
     {
         s32 start_index = 0;
         u64 updated_start = start + size;
+
+        // align updated_start to 8 bytes
+        if (updated_start % 8 != 0) {
+            updated_start += 8 - (updated_start % 8);
+        }
+
         bpf_map_update_elem(&alloc_map, &start_index, &updated_start, BPF_ANY);
         return target;
     }

include/arguments.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "common.h"
+#include "bpf_tracing.h"
 #include "bpf_helpers.h"
 #include <stdbool.h>
 
@@ -24,32 +25,32 @@ void *get_argument_by_reg(struct pt_regs *ctx, int index)
     switch (index)
     {
     case 1:
-        return (void *)(ctx->rax);
+        return (void *)GO_PARAM1(ctx);
     case 2:
-        return (void *)(ctx->rbx);
+        return (void *)GO_PARAM2(ctx);
     case 3:
-        return (void *)(ctx->rcx);
+        return (void *)GO_PARAM3(ctx);
     case 4:
-        return (void *)(ctx->rdi);
+        return (void *)GO_PARAM4(ctx);
     case 5:
-        return (void *)(ctx->rsi);
+        return (void *)GO_PARAM5(ctx);
     case 6:
-        return (void *)(ctx->r8);
+        return (void *)GO_PARAM6(ctx);
     case 7:
-        return (void *)(ctx->r9);
+        return (void *)GO_PARAM7(ctx);
     case 8:
-        return (void *)(ctx->r10);
+        return (void *)GO_PARAM8(ctx);
     case 9:
-        return (void *)(ctx->r11);
+        return (void *)GO_PARAM9(ctx);
     default:
         return NULL;
     }
 }
 
 void *get_argument_by_stack(struct pt_regs *ctx, int index)
 {
-    void *ptr = 0;
-    bpf_probe_read(&ptr, sizeof(ptr), (void *)(ctx->rsp + (index * 8)));
+    void* ptr = 0;
+    bpf_probe_read(&ptr, sizeof(ptr), (void *)(PT_REGS_SP(ctx)+(index*8)));
     return ptr;
 }
 

include/common.h

@@ -84,41 +84,57 @@ enum
 #define BPF_F_INDEX_MASK 0xffffffffULL
 #define BPF_F_CURRENT_CPU BPF_F_INDEX_MASK
 
-#define PT_REGS_RC(x) ((x)->rax)
-struct pt_regs
-{
-	/*
-	 * C ABI says these regs are callee-preserved. They aren't saved on kernel entry
-	 * unless syscall needs a complete, fully filled "struct pt_regs".
-	 */
-	unsigned long r15;
-	unsigned long r14;
-	unsigned long r13;
-	unsigned long r12;
-	unsigned long rbp;
-	unsigned long rbx;
-	/* These regs are callee-clobbered. Always saved on kernel entry. */
-	unsigned long r11;
-	unsigned long r10;
-	unsigned long r9;
-	unsigned long r8;
-	unsigned long rax;
-	unsigned long rcx;
-	unsigned long rdx;
-	unsigned long rsi;
-	unsigned long rdi;
-	/*
-	 * On syscall entry, this is syscall#. On CPU exception, this is error code.
-	 * On hw interrupt, it's IRQ number:
-	 */
-	unsigned long orig_rax;
-	/* Return frame for iretq */
-	unsigned long rip;
-	unsigned long cs;
-	unsigned long eflags;
-	unsigned long rsp;
-	unsigned long ss;
-	/* top of stack page */
+#if defined(__TARGET_ARCH_x86)
+struct pt_regs {
+	long unsigned int r15;
+	long unsigned int r14;
+	long unsigned int r13;
+	long unsigned int r12;
+	long unsigned int bp;
+	long unsigned int bx;
+	long unsigned int r11;
+	long unsigned int r10;
+	long unsigned int r9;
+	long unsigned int r8;
+	long unsigned int ax;
+	long unsigned int cx;
+	long unsigned int dx;
+	long unsigned int si;
+	long unsigned int di;
+	long unsigned int orig_ax;
+	long unsigned int ip;
+	long unsigned int cs;
+	long unsigned int flags;
+	long unsigned int sp;
+	long unsigned int ss;
+};
+#elif defined(__TARGET_ARCH_arm64)
+struct user_pt_regs {
+	__u64 regs[31];
+	__u64 sp;
+	__u64 pc;
+	__u64 pstate;
+};
+
+struct pt_regs {
+	union {
+		struct user_pt_regs user_regs;
+		struct {
+			u64 regs[31];
+			u64 sp;
+			u64 pc;
+			u64 pstate;
+		};
+	};
+	u64 orig_x0;
+	s32 syscallno;
+	u32 unused2;
+	u64 orig_addr_limit;
+	u64 pmr_save;
+	u64 stackframe[2];
+	u64 lockdep_hardirqs;
+	u64 exit_rcu;
 };
+#endif
 
 #endif /* __VMLINUX_H__ */

include/go_types.h

@@ -54,7 +54,11 @@ static __always_inline struct go_string write_user_go_string(char *str, u32 len)
     new_string.len = len;
 
     // Copy new string struct to userspace
-    write_target_data((void *)&new_string, sizeof(new_string));
+    void* res = write_target_data((void*)&new_string, sizeof(new_string));
+    if (res == NULL) {
+        new_string.len = 0;
+    }
+
     return new_string;
 }
 

include/libbpf/bpf_tracing.h

@@ -24,6 +24,9 @@
 #elif defined(__TARGET_ARCH_sparc)
 	#define bpf_target_sparc
 	#define bpf_target_defined
+#elif defined(__TARGET_ARCH_riscv)
+	#define bpf_target_riscv
+	#define bpf_target_defined
 #else
 
 /* Fall back to what the compiler says */
@@ -48,6 +51,9 @@
 #elif defined(__sparc__)
 	#define bpf_target_sparc
 	#define bpf_target_defined
+#elif defined(__riscv) && __riscv_xlen == 64
+	#define bpf_target_riscv
+	#define bpf_target_defined
 #endif /* no compiler target */
 
 #endif
@@ -60,6 +66,17 @@
 
 #if defined(__KERNEL__) || defined(__VMLINUX_H__)
 
+#define GO_PARAM1(x) ((x)->ax)
+#define GO_PARAM2(x) ((x)->bx)
+#define GO_PARAM3(x) ((x)->cx)
+#define GO_PARAM4(x) ((x)->di)
+#define GO_PARAM5(x) ((x)->si)
+#define GO_PARAM6(x) ((x)->r8)
+#define GO_PARAM7(x) ((x)->r9)
+#define GO_PARAM8(x) ((x)->r10)
+#define GO_PARAM9(x) ((x)->r11)
+#define GOROUTINE(x) ((x)->r14)
+
 #define PT_REGS_PARM1(x) ((x)->di)
 #define PT_REGS_PARM2(x) ((x)->si)
 #define PT_REGS_PARM3(x) ((x)->dx)
@@ -192,6 +209,18 @@ struct pt_regs;
 /* arm64 provides struct user_pt_regs instead of struct pt_regs to userspace */
 struct pt_regs;
 #define PT_REGS_ARM64 const volatile struct user_pt_regs
+
+#define GO_PARAM1(x) (((PT_REGS_ARM64 *)(x))->regs[0])
+#define GO_PARAM2(x) (((PT_REGS_ARM64 *)(x))->regs[1])
+#define GO_PARAM3(x) (((PT_REGS_ARM64 *)(x))->regs[2])
+#define GO_PARAM4(x) (((PT_REGS_ARM64 *)(x))->regs[3])
+#define GO_PARAM5(x) (((PT_REGS_ARM64 *)(x))->regs[4])
+#define GO_PARAM6(x) (((PT_REGS_ARM64 *)(x))->regs[5])
+#define GO_PARAM7(x) (((PT_REGS_ARM64 *)(x))->regs[6])
+#define GO_PARAM8(x) (((PT_REGS_ARM64 *)(x))->regs[7])
+#define GO_PARAM9(x) (((PT_REGS_ARM64 *)(x))->regs[8])
+#define GOROUTINE(x) (((PT_REGS_ARM64 *)(x))->regs[28])
+
 #define PT_REGS_PARM1(x) (((PT_REGS_ARM64 *)(x))->regs[0])
 #define PT_REGS_PARM2(x) (((PT_REGS_ARM64 *)(x))->regs[1])
 #define PT_REGS_PARM3(x) (((PT_REGS_ARM64 *)(x))->regs[2])
@@ -288,6 +317,32 @@ struct pt_regs;
 #define PT_REGS_IP_CORE(x) BPF_CORE_READ((x), pc)
 #endif
 
+#elif defined(bpf_target_riscv)
+
+struct pt_regs;
+#define PT_REGS_RV const volatile struct user_regs_struct
+#define PT_REGS_PARM1(x) (((PT_REGS_RV *)(x))->a0)
+#define PT_REGS_PARM2(x) (((PT_REGS_RV *)(x))->a1)
+#define PT_REGS_PARM3(x) (((PT_REGS_RV *)(x))->a2)
+#define PT_REGS_PARM4(x) (((PT_REGS_RV *)(x))->a3)
+#define PT_REGS_PARM5(x) (((PT_REGS_RV *)(x))->a4)
+#define PT_REGS_RET(x) (((PT_REGS_RV *)(x))->ra)
+#define PT_REGS_FP(x) (((PT_REGS_RV *)(x))->s5)
+#define PT_REGS_RC(x) (((PT_REGS_RV *)(x))->a5)
+#define PT_REGS_SP(x) (((PT_REGS_RV *)(x))->sp)
+#define PT_REGS_IP(x) (((PT_REGS_RV *)(x))->epc)
+
+#define PT_REGS_PARM1_CORE(x) BPF_CORE_READ((PT_REGS_RV *)(x), a0)
+#define PT_REGS_PARM2_CORE(x) BPF_CORE_READ((PT_REGS_RV *)(x), a1)
+#define PT_REGS_PARM3_CORE(x) BPF_CORE_READ((PT_REGS_RV *)(x), a2)
+#define PT_REGS_PARM4_CORE(x) BPF_CORE_READ((PT_REGS_RV *)(x), a3)
+#define PT_REGS_PARM5_CORE(x) BPF_CORE_READ((PT_REGS_RV *)(x), a4)
+#define PT_REGS_RET_CORE(x) BPF_CORE_READ((PT_REGS_RV *)(x), ra)
+#define PT_REGS_FP_CORE(x) BPF_CORE_READ((PT_REGS_RV *)(x), fp)
+#define PT_REGS_RC_CORE(x) BPF_CORE_READ((PT_REGS_RV *)(x), a5)
+#define PT_REGS_SP_CORE(x) BPF_CORE_READ((PT_REGS_RV *)(x), sp)
+#define PT_REGS_IP_CORE(x) BPF_CORE_READ((PT_REGS_RV *)(x), epc)
+
 #endif
 
 #if defined(bpf_target_powerpc)

pkg/instrumentors/bpf/github.com/gorilla/mux/probe.go

@@ -35,7 +35,7 @@ import (
 	"golang.org/x/sys/unix"
 )
 
-//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -target bpfel -cc clang -cflags $CFLAGS bpf ./bpf/probe.bpf.c
+//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -target $TARGET -cc clang -cflags $CFLAGS bpf ./bpf/probe.bpf.c
 
 type HttpEvent struct {
 	StartTime   uint64

pkg/instrumentors/bpf/google/golang/org/grpc/bpf/probe.bpf.c

@@ -132,31 +132,35 @@ int uprobe_Http2Client_CreateHeaderFields(struct pt_regs *ctx)
     struct go_slice_user_ptr slice_user_ptr = {};
     if (is_registers_abi)
     {
-        slice.array = (void *)ctx->rax;
-        slice.len = (s32)ctx->rbx;
-        slice.cap = (s32)ctx->rcx;
-        slice_user_ptr.array = &ctx->rax;
-        slice_user_ptr.len = &ctx->rbx;
-        slice_user_ptr.cap = &ctx->rcx;
+        slice.array = (void*) GO_PARAM1(ctx);
+        slice.len = (s32) GO_PARAM2(ctx);
+        slice.cap = (s32) GO_PARAM3(ctx);
+        slice_user_ptr.array = (void*) &GO_PARAM1(ctx);
+        slice_user_ptr.len = (void*) &GO_PARAM2(ctx);
+        slice_user_ptr.cap = (void*) &GO_PARAM3(ctx);
     }
     else
     {
         u64 slice_pointer_pos = 5;
         s32 slice_len_pos = 6;
         s32 slice_cap_pos = 7;
         slice.array = get_argument(ctx, slice_pointer_pos);
-        slice.len = (s32)get_argument(ctx, slice_len_pos);
-        slice.cap = (s32)get_argument(ctx, slice_cap_pos);
-        slice_user_ptr.array = (void *)ctx->rsp + (slice_pointer_pos * 8);
-        slice_user_ptr.len = (void *)ctx->rsp + (slice_len_pos * 8);
-        slice_user_ptr.cap = (void *)ctx->rsp + (slice_cap_pos * 8);
+        slice.len = (long)get_argument(ctx, slice_len_pos);
+        slice.cap = (long)get_argument(ctx, slice_cap_pos);
+        slice_user_ptr.array = (void *)(PT_REGS_SP(ctx)+(slice_pointer_pos*8));
+        slice_user_ptr.len = (void *)(PT_REGS_SP(ctx)+(slice_len_pos*8));
+        slice_user_ptr.cap = (void *)(PT_REGS_SP(ctx)+(slice_cap_pos*8));
     }
     char key[11] = "traceparent";
     struct go_string key_str = write_user_go_string(key, sizeof(key));
+    if (key_str.len == 0) {
+        bpf_printk("write failed, aborting ebpf probe");
+        return 0;
+    }
 
     // Get grpc request struct
     void *context_ptr = 0;
-    bpf_probe_read(&context_ptr, sizeof(context_ptr), (void *)(ctx->rsp + (context_pointer_pos * 8)));
+    bpf_probe_read(&context_ptr, sizeof(context_ptr), (void *)(PT_REGS_SP(ctx)+(context_pointer_pos*8)));
     void *parent_ctx = find_context_in_map(context_ptr, &context_to_grpc_events);
     void *grpcReq_ptr = bpf_map_lookup_elem(&context_to_grpc_events, &parent_ctx);
     struct grpc_request_t grpcReq = {};

pkg/instrumentors/bpf/google/golang/org/grpc/probe.go

@@ -36,7 +36,7 @@ import (
 	"golang.org/x/sys/unix"
 )
 
-//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -target bpfel -cc clang -cflags $CFLAGS bpf ./bpf/probe.bpf.c
+//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -target $TARGET -cc clang -cflags $CFLAGS bpf ./bpf/probe.bpf.c
 
 type GrpcEvent struct {
 	StartTime         uint64