Correct func IP in bpf

It's better to correct func IP in bpf then in user space, because in bpf
it is able to distinguish every case, include endbr case.

As a result, it gets `addr` by this `get_addr()` function for **kprobe**,
**kprobe-multi**, **fentry** and **fexit**.

Signed-off-by: Leon Hwang <hffilwlqm@gmail.com>

bpf/kprobe_pwru.c

@@ -511,6 +511,62 @@ handle_everything(struct sk_buff *skb, void *ctx, struct event_t *event, u64 *_s
 	return true;
 }
 
+#ifdef bpf_target_x86
+/* The following gen_endbr_poison() and is_endbr() functions are taken from the
+ * kernel's arch/x86/include/asm/ibt.h file.
+ */
+
+static __always_inline u32
+gen_endbr_poison(void)
+{
+	/*
+	 * 4 byte NOP that isn't NOP4 (in fact it is OSP NOP3), such that it
+	 * will be unique to (former) ENDBR sites.
+	 */
+	return 0x001f0f66; /* osp nopl (%rax) */
+}
+
+static __always_inline bool
+is_endbr(u32 val)
+{
+	static const u32 endbr64 = 0xfa1e0ff3;
+	static const u32 endbr32 = 0xfb1e0ff3;
+
+	if (val == gen_endbr_poison())
+		return true;
+
+	val &= ~0x01000000U; /* ENDBR32 -> ENDBR64 */
+	return val == endbr64;
+}
+
+static __always_inline u64
+correct_func_ip(u64 ip) {
+	u32 endbr;
+
+	bpf_probe_read_kernel(&endbr, sizeof(endbr), (void *) (ip - 4));
+	return is_endbr(endbr) ? (ip - 4) : ip;
+}
+#endif
+
+static __always_inline u64 get_func_ip(void);
+
+static __always_inline u64
+get_addr(void *ctx, const bool is_kprobe, const bool has_get_func_ip) {
+	u64 ip;
+
+	if (has_get_func_ip) {
+		ip = bpf_get_func_ip(ctx); /* endbr has been handled in helper. */
+	} else {
+		ip = is_kprobe ? PT_REGS_IP((struct pt_regs *) ctx) : get_func_ip();
+#ifdef bpf_target_x86
+		ip = correct_func_ip(ip);
+		ip -= is_kprobe; /* -1 always on x86 if kprobe. */
+#endif
+	}
+
+	return ip;
+}
+
 static __always_inline int
 kprobe_skb(struct sk_buff *skb, struct pt_regs *ctx, const bool has_get_func_ip, u64 *_stackid) {
 	struct event_t event = {};
@@ -519,7 +575,7 @@ kprobe_skb(struct sk_buff *skb, struct pt_regs *ctx, const bool has_get_func_ip,
 		return BPF_OK;
 
 	event.skb_addr = (u64) skb;
-	event.addr = has_get_func_ip ? bpf_get_func_ip(ctx) : PT_REGS_IP(ctx);
+	event.addr = get_addr(ctx, true, has_get_func_ip);
 	event.param_second = PT_REGS_PARM2(ctx);
 	event.param_third = PT_REGS_PARM3(ctx);
 	if (CFG.output_caller)
@@ -691,7 +747,7 @@ int BPF_PROG(fentry_tc, struct sk_buff *skb) {
 		return BPF_OK;
 
 	event.skb_addr = (u64) skb;
-	event.addr = get_func_ip();
+	event.addr = get_addr(ctx, false, false);
 	event.type = EVENT_TYPE_TC;
 	bpf_map_push_elem(&events, &event, BPF_EXIST);
 
@@ -799,7 +855,7 @@ int BPF_PROG(fentry_xdp, struct xdp_buff *xdp) {
 	event.ts = bpf_ktime_get_ns();
 	event.cpu_id = bpf_get_smp_processor_id();
 	event.skb_addr = (u64) &xdp;
-	event.addr = get_func_ip();
+	event.addr = get_addr(ctx, false, false);
 	event.type = EVENT_TYPE_XDP;
 	bpf_map_push_elem(&events, &event, BPF_EXIST);
 

internal/pwru/output.go

@@ -257,19 +257,6 @@ func getExecName(pid int) string {
 	return execName
 }
 
-func getAddrByArch(event *Event, o *output) (addr uint64) {
-	switch runtime.GOARCH {
-	case "amd64":
-		addr = event.Addr
-		if !o.kprobeMulti && event.Type == eventTypeKprobe {
-			addr -= 1
-		}
-	case "arm64":
-		addr = event.Addr
-	}
-	return addr
-}
-
 func getTupleData(event *Event) (tupleData string) {
 	tupleData = fmt.Sprintf("%s:%d->%s:%d(%s)",
 		addrToStr(event.Tuple.L3Proto, event.Tuple.Saddr), byteorder.NetworkToHost16(event.Tuple.Sport),
@@ -350,11 +337,6 @@ func getOutFuncName(o *output, event *Event, addr uint64) string {
 
 	if ksym, ok := o.addr2name.Addr2NameMap[addr]; ok {
 		funcName = ksym.name
-	} else if ksym, ok := o.addr2name.Addr2NameMap[addr-4]; runtime.GOARCH == "amd64" && ok {
-		// Assume that function has ENDBR in its prelude (enabled by CONFIG_X86_KERNEL_IBT).
-		// See https://lore.kernel.org/bpf/20220811091526.172610-5-jolsa@kernel.org/
-		// for more ctx.
-		funcName = ksym.name
 	} else {
 		funcName = fmt.Sprintf("0x%x", addr)
 	}
@@ -411,10 +393,7 @@ func (o *output) Print(event *Event) {
 		ts = getRelativeTs(event, o)
 	}
 
-	// XXX: not sure why the -1 offset is needed on x86 but not on arm64
-	addr := getAddrByArch(event, o)
-
-	outFuncName := getOutFuncName(o, event, addr)
+	outFuncName := getOutFuncName(o, event, event.Addr)
 
 	fmt.Fprintf(o.writer, "%-18s %-3s %-16s", fmt.Sprintf("%#x", event.SkbAddr),
 		fmt.Sprintf("%d", event.CPU), fmt.Sprintf("%s", execName))