bpf: Remove HAS_KPROBE_MULTI macro

Simplify bpf code. And use Go to determine loading kprobe or kprobe.multi bpf prog. Signed-off-by: Leon Hwang <hffilwlqm@gmail.com>
cilium · Nov 10, 2024 · 2571b11 · 2571b11
1 parent 3a05678
commit 2571b11
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 27 deletions.
diff --git a/bpf/kprobe_pwru.c b/bpf/kprobe_pwru.c
@@ -490,7 +490,7 @@ 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, 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) {
 	struct event_t event = {};
 
 	if (!handle_everything(skb, ctx, &event, _stackid, true))
@@ -508,19 +508,17 @@ kprobe_skb(struct sk_buff *skb, struct pt_regs *ctx, bool has_get_func_ip, u64 *
 	return BPF_OK;
 }
 
-#ifdef HAS_KPROBE_MULTI
-#define PWRU_KPROBE_TYPE "kprobe.multi"
-#define PWRU_HAS_GET_FUNC_IP true
-#else
-#define PWRU_KPROBE_TYPE "kprobe"
-#define PWRU_HAS_GET_FUNC_IP false
-#endif /* HAS_KPROBE_MULTI */
-
-#define PWRU_ADD_KPROBE(X)                                                     \
-  SEC(PWRU_KPROBE_TYPE "/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, PWRU_HAS_GET_FUNC_IP, 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);				   				\
+  }																			\
+																			\
+  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);				   				\
   }
 
 PWRU_ADD_KPROBE(1)
@@ -529,21 +527,19 @@ PWRU_ADD_KPROBE(3)
 PWRU_ADD_KPROBE(4)
 PWRU_ADD_KPROBE(5)
 
+#undef PWRU_ADD_KPROBE
+
 SEC("kprobe/skb_by_stackid")
 int kprobe_skb_by_stackid(struct pt_regs *ctx) {
 	u64 stackid = get_stackid(ctx);
 
 	struct sk_buff **skb = bpf_map_lookup_elem(&stackid_skb, &stackid);
 	if (skb && *skb)
-		return kprobe_skb(*skb, ctx, PWRU_HAS_GET_FUNC_IP, &stackid);
+		return kprobe_skb(*skb, ctx, false, &stackid);
 
 	return BPF_OK;
 }
 
-#undef PWRU_KPROBE
-#undef PWRU_HAS_GET_FUNC_IP
-#undef PWRU_KPROBE_TYPE
-
 SEC("kprobe/skb_lifetime_termination")
 int kprobe_skb_lifetime_termination(struct pt_regs *ctx) {
 	struct sk_buff *skb = (typeof(skb)) PT_REGS_PARM1(ctx);

diff --git a/build.go b/build.go
@@ -3,6 +3,5 @@
 
 //go:generate sh -c "echo Generating for $TARGET_GOARCH"
 //go:generate go run github.com/cilium/ebpf/cmd/bpf2go -target $TARGET_GOARCH -cc clang -no-strip KProbePWRU ./bpf/kprobe_pwru.c -- -I./bpf/headers -Wno-address-of-packed-member
-//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -target $TARGET_GOARCH -cc clang -no-strip KProbeMultiPWRU ./bpf/kprobe_pwru.c -- -DHAS_KPROBE_MULTI -I./bpf/headers -Wno-address-of-packed-member
 
 package main
diff --git a/main.go b/main.go
@@ -115,17 +115,21 @@ func main() {
 	opts.Programs.LogLevel = ebpf.LogLevelInstruction
 	opts.Programs.LogSize = ebpf.DefaultVerifierLogSize * 100
 
-	var bpfSpec *ebpf.CollectionSpec
-	switch {
-	case (flags.OutputSkb || flags.OutputShinfo) && useKprobeMulti:
-		bpfSpec, err = LoadKProbeMultiPWRU()
-	default:
-		bpfSpec, err = LoadKProbePWRU()
-	}
+	bpfSpec, err := LoadKProbePWRU()
 	if err != nil {
 		log.Fatalf("Failed to load bpf spec: %v", err)
 	}
 
+	if useKprobeMulti {
+		for i := 1; i <= 5; i++ {
+			delete(bpfSpec.Programs, fmt.Sprintf("kprobe_skb_%d", i))
+		}
+	} else {
+		for i := 1; i <= 5; i++ {
+			delete(bpfSpec.Programs, fmt.Sprintf("kprobe_multi_skb_%d", i))
+		}
+	}
+
 	for name, program := range bpfSpec.Programs {
 		// Skip the skb-tracking ones that should not inject pcap-filter.
 		switch name {