From 0dc1f2cd8abcc5fe3d9a41a87fde3d8257bd700a Mon Sep 17 00:00:00 2001
From: Leon Hwang <hffilwlqm@gmail.com>
Date: Tue, 24 Oct 2023 20:32:06 +0800
Subject: [PATCH] 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, by the
way `fexit` on `skb_clone()` and `skb_copy()`.

Signed-off-by: Leon Hwang <hffilwlqm@gmail.com>
---
 bpf/kprobe_pwru.c      | 18 ++++++++++++++++++
 internal/pwru/utils.go | 29 +++++++++++++++++++++++++++++
 main.go                | 35 ++++++++++++++++++++++++++++++++++-
 3 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/bpf/kprobe_pwru.c b/bpf/kprobe_pwru.c
index 5d6039ab..b6f6d205 100644
--- a/bpf/kprobe_pwru.c
+++ b/bpf/kprobe_pwru.c
@@ -347,6 +347,24 @@ int kprobe_skb_lifetime_termination(struct pt_regs *ctx) {
 	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")
 int BPF_PROG(fentry_tc, struct sk_buff *skb) {
 	struct event_t event = {};
diff --git a/internal/pwru/utils.go b/internal/pwru/utils.go
index 7036f52a..ad88a07f 100644
--- a/internal/pwru/utils.go
+++ b/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
+}
diff --git a/main.go b/main.go
index b0dc0e91..1c01d254 100644
--- a/main.go
+++ b/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,10 +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 (
@@ -228,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)