Support tracking skb clones

With --filter-track-skb option, we're able to track sbk by its address.

Then, as for skb's clones/copies, we're able to track them too. We track
the skb clones/copies by kprobing `__skb_clone()` and `skb_copy_header()`,
whose callers include `skb_clone()` and `skb_copy()`.

Signed-off-by: Leon Hwang <[email protected]>

Author: Asphaltt

bpf/kprobe_pwru.c

@@ -344,7 +344,28 @@ int kprobe_skb_lifetime_termination(struct pt_regs *ctx) {
 
 	bpf_map_delete_elem(&skb_addresses, &skb);
 
-	return 0;
+	return BPF_OK;
+}
+
+static __always_inline int
+track_skb_clone(struct pt_regs *ctx) {
+	u64 new = (u64) PT_REGS_PARM1(ctx);
+	u64 old = (u64) PT_REGS_PARM2(ctx);
+
+	if (bpf_map_lookup_elem(&skb_addresses, &old))
+		bpf_map_update_elem(&skb_addresses, &new, &TRUE, BPF_ANY);
+
+	return BPF_OK;
+}
+
+SEC("kprobe/__skb_clone")
+int kprobe___skb_clone(struct pt_regs *ctx) {
+	return track_skb_clone(ctx);
+}
+
+SEC("kprobe/skb_copy_header")
+int kprobe_skb_copy_header(struct pt_regs *ctx) {
+	return track_skb_clone(ctx);
 }
 
 SEC("fentry/tc")

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 == "kprobe___skb_clone" ||
+			name == "kprobe_skb_copy_header" {
 			continue
 		}
 		if err = libpcap.InjectFilters(program, flags.FilterPcap); err != nil {
@@ -152,8 +155,12 @@ 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")
+		delete(bpfSpec.Programs, "kprobe___skb_clone")
+		delete(bpfSpec.Programs, "kprobe_skb_copy_header")
 	}
 
 	coll, err := ebpf.NewCollectionWithOptions(bpfSpec, opts)
@@ -216,17 +223,28 @@ func main() {
 	bar := pb.StartNew(len(funcs))
 
 	if flags.FilterTrackSkb {
-		kp, err := link.Kprobe("kfree_skbmem", coll.Programs["kprobe_skb_lifetime_termination"], nil)
-		bar.Increment()
-		if err != nil {
-			if !errors.Is(err, os.ErrNotExist) {
-				log.Fatalf("Opening kprobe kfree_skbmem: %s\n", err)
+		progs := []struct {
+			hook string
+			prog *ebpf.Program
+			hint string
+		}{
+			{"kfree_skbmem", coll.Programs["kprobe_skb_lifetime_termination"], "skb lifetime termination"},
+			{"__skb_clone", coll.Programs["kprobe___skb_clone"], "skb clone tracking"},
+			{"skb_copy_header", coll.Programs["kprobe_skb_copy_header"], "skb copy tracking"},
+		}
+		for _, prog := range progs {
+			kp, err := link.Kprobe(prog.hook, prog.prog, nil)
+			bar.Increment()
+			if err != nil {
+				if !errors.Is(err, os.ErrNotExist) {
+					log.Fatalf("Opening kprobe %s: %s\n", prog.hook, err)
+				} else {
+					ignored += 1
+					log.Printf("Warn: %s not found, pwru is likely to mismatch skb due to lack of %s\n", prog.hook, prog.hint)
+				}
 			} else {
-				ignored += 1
-				log.Printf("Warn: kfree_skbmem not found, pwru is likely to mismatch skb due to lack of skb lifetime management\n")
+				kprobes = append(kprobes, kp)
 			}
-		} else {
-			kprobes = append(kprobes, kp)
 		}
 	}