Support tracking skb clones

bpf/kprobe_pwru.c

@@ -341,10 +341,28 @@ PWRU_ADD_KPROBE(5)
 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;
+	bpf_map_delete_elem(&skb_addresses, &skb);
+
+	return BPF_OK;
+}
+
+static __always_inline int
+track_skb_clone(u64 old, u64 new) {
+	if (bpf_map_lookup_elem(&skb_addresses, &old))
+		bpf_map_update_elem(&skb_addresses, &new, &TRUE, BPF_ANY);
+
+	return BPF_OK;
+}
+
+SEC("fexit/skb_clone")
+int BPF_PROG(fexit_skb_clone, u64 old, gfp_t mask, u64 new) {
+	return track_skb_clone(old, new);
+}
+
+SEC("fexit/skb_copy")
+int BPF_PROG(fexit_skb_copy, u64 old, gfp_t mask, u64 new) {
+	return track_skb_clone(old, new);
 }
 
 SEC("fentry/tc")

internal/pwru/utils.go

@@ -155,3 +155,32 @@ func HaveBPFLinkKprobeMulti() bool {
 
 	return true
 }
+
+// Very hacky way to check whether tracing link is supported.
+func HaveBPFLinkTracing() bool {
+	prog, err := ebpf.NewProgram(&ebpf.ProgramSpec{
+		Name: "fexit_skb_clone",
+		Type: ebpf.Tracing,
+		Instructions: asm.Instructions{
+			asm.Mov.Imm(asm.R0, 0),
+			asm.Return(),
+		},
+		AttachType: ebpf.AttachTraceFExit,
+		AttachTo:   "skb_clone",
+		License:    "MIT",
+	})
+	if err != nil {
+		return false
+	}
+	defer prog.Close()
+
+	link, err := link.AttachTracing(link.TracingOptions{
+		Program: prog,
+	})
+	if err != nil {
+		return false
+	}
+	defer link.Close()
+
+	return true
+}

main.go

@@ -121,7 +121,10 @@ func main() {
 	}
 
 	for name, program := range bpfSpec.Programs {
-		if name == "kprobe_skb_lifetime_termination" {
+		// Skip the skb-tracking ones that should not inject pcap-filter.
+		if name == "kprobe_skb_lifetime_termination" ||
+			name == "fexit_skb_clone" ||
+			name == "fexit_skb_copy" {
 			continue
 		}
 		if err = libpcap.InjectFilters(program, flags.FilterPcap); err != nil {
@@ -152,6 +155,18 @@ func main() {
 	// deleted from the spec.
 	delete(bpfSpec.Programs, "fentry_tc")
 
+	// If not tracking skb, deleting the skb-tracking programs to reduce loading
+	// time.
+	if !flags.FilterTrackSkb {
+		delete(bpfSpec.Programs, "kprobe_skb_lifetime_termination")
+	}
+
+	haveFexit := pwru.HaveBPFLinkTracing()
+	if !flags.FilterTrackSkb || !haveFexit {
+		delete(bpfSpec.Programs, "fexit_skb_clone")
+		delete(bpfSpec.Programs, "fexit_skb_copy")
+	}
+
 	coll, err := ebpf.NewCollectionWithOptions(bpfSpec, opts)
 	if err != nil {
 		var (
@@ -171,7 +186,6 @@ func main() {
 	kprobe3 := coll.Programs["kprobe_skb_3"]
 	kprobe4 := coll.Programs["kprobe_skb_4"]
 	kprobe5 := coll.Programs["kprobe_skb_5"]
-	kprobeLifetimeTermination := coll.Programs["kprobe_skb_lifetime_termination"]
 
 	events := coll.Maps["events"]
 	printStackMap := coll.Maps["print_stack_map"]
@@ -213,7 +227,7 @@ func main() {
 	bar := pb.StartNew(len(funcs))
 
 	if flags.FilterTrackSkb {
-		kp, err := link.Kprobe("kfree_skbmem", kprobeLifetimeTermination, nil)
+		kp, err := link.Kprobe("kfree_skbmem", coll.Programs["kprobe_skb_lifetime_termination"], nil)
 		bar.Increment()
 		if err != nil {
 			if !errors.Is(err, os.ErrNotExist) {
@@ -225,6 +239,28 @@ func main() {
 		} else {
 			kprobes = append(kprobes, kp)
 		}
+
+		if haveFexit {
+			progs := []*ebpf.Program{
+				coll.Programs["fexit_skb_clone"],
+				coll.Programs["fexit_skb_copy"],
+			}
+			for _, prog := range progs {
+				fexit, err := link.AttachTracing(link.TracingOptions{
+					Program: prog,
+				})
+				bar.Increment()
+				if err != nil {
+					if !errors.Is(err, os.ErrNotExist) {
+						log.Fatalf("Opening tracing(%s): %s\n", prog, err)
+					} else {
+						ignored += 1
+					}
+				} else {
+					kprobes = append(kprobes, fexit)
+				}
+			}
+		}
 	}
 
 	funcsByPos := pwru.GetFuncsByPos(funcs)