From 1fee05821004f437094cbefe04b44b65eb578229 Mon Sep 17 00:00:00 2001
From: gray <greyschwinger@gmail.com>
Date: Fri, 3 Jan 2025 17:22:45 +0800
Subject: [PATCH] Events from kprobe.multi have eventTypeKprobeMulti set in
 .type

Even with --backend kprobe-multi, pwru still uses kprobe for
--filter-track-bpf-helpers. This patch makes pwru capable of
distinguishing event types in order to adjust addressres for symbol
resolution.

Fixes: https://github.com/cilium/pwru/issues/462

Signed-off-by: gray <gray.liang@isovalent.com>
---
 bpf/kprobe_pwru.c       | 38 +++++++++++++++++++++-----------------
 internal/pwru/output.go |  9 +++++----
 2 files changed, 26 insertions(+), 21 deletions(-)

diff --git a/bpf/kprobe_pwru.c b/bpf/kprobe_pwru.c
index 42fca157..d3d84bc1 100644
--- a/bpf/kprobe_pwru.c
+++ b/bpf/kprobe_pwru.c
@@ -68,9 +68,10 @@ struct tuple {
 } __attribute__((packed));
 
 enum event_type {
-	EVENT_TYPE_KPROBE = 0,
-	EVENT_TYPE_TC     = 1,
-	EVENT_TYPE_XDP    = 2,
+	EVENT_TYPE_KPROBE	= 0,
+	EVENT_TYPE_KPROBE_MULTI	= 1,
+	EVENT_TYPE_TC		= 2,
+	EVENT_TYPE_XDP		= 3,
 };
 
 struct event_t {
@@ -515,7 +516,8 @@ handle_everything(struct sk_buff *skb, void *ctx, struct event_t *event, u64 *_s
 }
 
 static __always_inline int
-kprobe_skb(struct sk_buff *skb, struct pt_regs *ctx, const bool has_get_func_ip, u64 *_stackid) {
+kprobe_skb(struct sk_buff *skb, struct pt_regs *ctx, const bool has_get_func_ip,
+	   u64 *_stackid, const bool kprobe_multi) {
 	struct event_t event = {};
 
 	if (!handle_everything(skb, ctx, &event, _stackid, true))
@@ -523,28 +525,30 @@ kprobe_skb(struct sk_buff *skb, struct pt_regs *ctx, const bool has_get_func_ip,
 
 	event.skb_addr = (u64) skb;
 	event.addr = has_get_func_ip ? bpf_get_func_ip(ctx) : PT_REGS_IP(ctx);
+	event.type = kprobe_multi ? EVENT_TYPE_KPROBE_MULTI: EVENT_TYPE_KPROBE;
 	event.param_second = PT_REGS_PARM2(ctx);
 	event.param_third = PT_REGS_PARM3(ctx);
 	if (CFG.output_caller)
 		bpf_probe_read_kernel(&event.caller_addr, sizeof(event.caller_addr), (void *)PT_REGS_SP(ctx));
 
+
 	bpf_map_push_elem(&events, &event, BPF_EXIST);
 
 	return BPF_OK;
 }
 
-#define PWRU_ADD_KPROBE(X)                                                  \
-  SEC("kprobe/skb-" #X)                                                     \
-  int kprobe_skb_##X(struct pt_regs *ctx) {                                 \
-    struct sk_buff *skb = (struct sk_buff *) PT_REGS_PARM##X(ctx);          \
-    return kprobe_skb(skb, ctx, false, NULL);                               \
-  }                                                                         \
-                                                                            \
-  SEC("kprobe.multi/skb-" #X)                                               \
-  int kprobe_multi_skb_##X(struct pt_regs *ctx) {                           \
-    struct sk_buff *skb = (struct sk_buff *) PT_REGS_PARM##X(ctx);          \
-    return kprobe_skb(skb, ctx, true, NULL);                                \
-  }
+#define PWRU_ADD_KPROBE(X)							\
+SEC("kprobe/skb-" #X)								\
+	int kprobe_skb_##X(struct pt_regs *ctx) {				\
+		struct sk_buff *skb = (struct sk_buff *) PT_REGS_PARM##X(ctx);	\
+		return kprobe_skb(skb, ctx, false, NULL, false);		\
+	}									\
+										\
+	SEC("kprobe.multi/skb-" #X)						\
+	int kprobe_multi_skb_##X(struct pt_regs *ctx) {				\
+		struct sk_buff *skb = (struct sk_buff *) PT_REGS_PARM##X(ctx);	\
+		return kprobe_skb(skb, ctx, true, NULL, true);			\
+	}
 
 PWRU_ADD_KPROBE(1)
 PWRU_ADD_KPROBE(2)
@@ -560,7 +564,7 @@ int kprobe_skb_by_stackid(struct pt_regs *ctx) {
 
 	struct sk_buff **skb = bpf_map_lookup_elem(&stackid_skb, &stackid);
 	if (skb && *skb)
-		return kprobe_skb(*skb, ctx, false, &stackid);
+		return kprobe_skb(*skb, ctx, false, &stackid, false);
 
 	return BPF_OK;
 }
diff --git a/internal/pwru/output.go b/internal/pwru/output.go
index 877053f3..595c1bbc 100644
--- a/internal/pwru/output.go
+++ b/internal/pwru/output.go
@@ -31,9 +31,10 @@ import (
 const absoluteTS string = "15:04:05.000"
 
 const (
-	eventTypeKprobe     = 0
-	eventTypeTracingTc  = 1
-	eventTypeTracingXdp = 2
+	eventTypeKprobe = iota
+	eventTypeKprobeMulti
+	eventTypeTracingTc
+	eventTypeTracingXdp
 )
 
 type output struct {
@@ -261,7 +262,7 @@ func getAddrByArch(event *Event, o *output) (addr uint64) {
 	switch runtime.GOARCH {
 	case "amd64":
 		addr = event.Addr
-		if !o.kprobeMulti && event.Type == eventTypeKprobe {
+		if event.Type == eventTypeKprobe {
 			addr -= 1
 		}
 	case "arm64":