Stop tracking skb in the end of its lifetime

`kfree_skbmem` marks the end of an skb's lifetime, so we stop tracking it at that time when `--filter-track-skb` is enabled. Signed-off-by: Zhichuan Liang <[email protected]>
cilium · Aug 30, 2023 · cc9c63f · cc9c63f
1 parent 3f39b56
commit cc9c63f
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 0 deletions.
diff --git a/bpf/kprobe_pwru.c b/bpf/kprobe_pwru.c
@@ -321,4 +321,13 @@ PWRU_ADD_KPROBE(5)
 #undef PWRU_HAS_GET_FUNC_IP
 #undef PWRU_KPROBE_TYPE
 
+SEC("kprobe/skb_lifetime_termination")
+int kprobe_skb_lifetime_termination(struct pt_regs *ctx) {
+	u64 skb = (u64) PT_REGS_PARM1(ctx);
+	if (cfg->track_skb)
+		bpf_map_delete_elem(&skb_addresses, &skb);
+
+	return 0;
+}
+
 char __license[] SEC("license") = "Dual BSD/GPL";
diff --git a/internal/pwru/types.go b/internal/pwru/types.go
@@ -138,6 +138,7 @@ type KProbePrograms interface {
 	GetKprobeSkb3() *ebpf.Program
 	GetKprobeSkb4() *ebpf.Program
 	GetKprobeSkb5() *ebpf.Program
+	GetKprobeSkbLifetimeTermination() *ebpf.Program
 }
 
 type KProbeObjects interface {

diff --git a/main.go b/main.go
@@ -124,6 +124,9 @@ func main() {
 	}
 
 	for name, program := range bpfSpec.Programs {
+		if name == "kprobe_skb_lifetime_termination" {
+			continue
+		}
 		if err = libpcap.InjectFilter(program, flags.FilterPcap); err != nil {
 			log.Fatalf("Failed to inject filter ebpf for %s: %v", name, err)
 		}
@@ -153,6 +156,7 @@ func main() {
 	kprobe3 := objs.GetKprobeSkb3()
 	kprobe4 := objs.GetKprobeSkb4()
 	kprobe5 := objs.GetKprobeSkb5()
+	kprobeLifetimeTermination := objs.GetKprobeSkbLifetimeTermination()
 
 	events := objs.GetEvents()
 	printStackMap := objs.GetPrintStackMap()
@@ -187,6 +191,22 @@ func main() {
 	log.Printf("Attaching kprobes (via %s)...\n", msg)
 	ignored := 0
 	bar := pb.StartNew(len(funcs))
+
+	if flags.FilterTrackSkb {
+		kp, err := link.Kprobe("kfree_skbmem", kprobeLifetimeTermination, nil)
+		bar.Increment()
+		if err != nil {
+			if !errors.Is(err, os.ErrNotExist) {
+				log.Fatalf("Opening kprobe kfree_skbmem: %s\n", err)
+			} else {
+				ignored += 1
+				log.Printf("Warn: kfree_skbmem not found, pwru is likely to mismatch skb due to lack of skb lifetime management\n")
+			}
+		} else {
+			kprobes = append(kprobes, kp)
+		}
+	}
+
 	funcsByPos := pwru.GetFuncsByPos(funcs)
 	for pos, fns := range funcsByPos {
 		var fn *ebpf.Program