From 7a336831849c3ed38c7143921417c8f459e6ef6b Mon Sep 17 00:00:00 2001
From: gray <greyschwinger@gmail.com>
Date: Sat, 29 Jun 2024 19:55:55 +0800
Subject: [PATCH] --filter-track-skb can track skbs re-built from
 veth_convert_skb_to_xdp_buff

When XDP is attached to a veth, skbs will be consumed and re-created on
that veth. This is done in the function veth_convert_skb_to_xdp_buff():

```
// drivers/net/veth.c
static int veth_convert_skb_to_xdp_buff(struct veth_rq *rq,
					struct xdp_buff *xdp,
					struct sk_buff **pskb)
{
	struct sk_buff *skb = *pskb;
[...]
		nskb = build_skb(page_address(page), PAGE_SIZE);
[...]
		skb_copy_header(nskb, skb);
[...]
		consume_skb(skb);
		skb = nskb;
[...]
}
```

This causes problems for pwru --filter-track-skb because of the new skb
addresses. I ran into a lot of situations where I lost track of NAT-ed
traffic at veth when cilium kind cluster is created by "kind.sh --xdp".

This patch allows pwru to keep track of the new skbs at XDP-attached
veth devices.

Signed-off-by: gray <gray.liang@isovalent.com>
---
 bpf/kprobe_pwru.c            | 32 ++++++++++++++++++++++++++++++++
 internal/pwru/skb_tracker.go | 18 ++++++++++++++++++
 main.go                      |  9 ++++++---
 3 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/bpf/kprobe_pwru.c b/bpf/kprobe_pwru.c
index 8d85416c..99fc10c9 100644
--- a/bpf/kprobe_pwru.c
+++ b/bpf/kprobe_pwru.c
@@ -105,6 +105,13 @@ struct {
 	__uint(max_entries, MAX_TRACK_SIZE);
 } skb_addresses SEC(".maps");
 
+struct {
+	__uint(type, BPF_MAP_TYPE_HASH);
+	__type(key, __u64); // pid_tgid
+	__type(value, __u64); // struct sk_buff **
+	__uint(max_entries, MAX_TRACK_SIZE);
+} veth_skbs SEC(".maps");
+
 struct {
 	__uint(type, BPF_MAP_TYPE_HASH);
 	__type(key, struct sk_buff *);
@@ -686,4 +693,29 @@ int BPF_PROG(fentry_xdp, struct xdp_buff *xdp) {
 	return BPF_OK;
 }
 
+SEC("kprobe/veth_convert_skb_to_xdp_buff")
+int kprobe_veth_convert_skb_to_xdp_buff(struct pt_regs *ctx) {
+	struct sk_buff **skb = (struct sk_buff **)PT_REGS_PARM3(ctx);
+	u64 skb_addr;
+	bpf_probe_read_kernel(&skb_addr, sizeof(skb_addr), (void *)skb);
+	if (bpf_map_lookup_elem(&skb_addresses, &skb_addr)) {
+		u64 pid_tgid = bpf_get_current_pid_tgid();
+		bpf_map_update_elem(&veth_skbs, &pid_tgid, &skb, BPF_ANY);
+	}
+	return BPF_OK;
+}
+
+SEC("kretprobe/veth_convert_skb_to_xdp_buff")
+int kretprobe_veth_convert_skb_to_xdp_buff(struct pt_regs *ctx) {
+	u64 pid_tgid = bpf_get_current_pid_tgid();
+	struct sk_buff ***skb = (struct sk_buff ***)bpf_map_lookup_elem(&veth_skbs, &pid_tgid);
+	if (skb && *skb) {
+		u64 skb_addr;
+		bpf_probe_read_kernel(&skb_addr, sizeof(skb_addr), (void *)*skb);
+		bpf_map_update_elem(&skb_addresses, &skb_addr, &TRUE, BPF_ANY);
+		bpf_map_delete_elem(&veth_skbs, &pid_tgid);
+	}
+	return BPF_OK;
+}
+
 char __license[] SEC("license") = "Dual BSD/GPL";
diff --git a/internal/pwru/skb_tracker.go b/internal/pwru/skb_tracker.go
index 4dfe8450..2f9995ad 100644
--- a/internal/pwru/skb_tracker.go
+++ b/internal/pwru/skb_tracker.go
@@ -38,6 +38,24 @@ func TrackSkb(coll *ebpf.Collection, haveFexit, trackSkbClone bool) *skbTracker
 		t.links = append(t.links, kp)
 	}
 
+	kp, err = link.Kretprobe("veth_convert_skb_to_xdp_buff", coll.Programs["kretprobe_veth_convert_skb_to_xdp_buff"], nil)
+	if err != nil {
+		if !errors.Is(err, os.ErrNotExist) {
+			log.Fatalf("Opening kretprobe veth_convert_skb_to_xdp_buff: %s\n", err)
+		}
+	} else {
+		t.links = append(t.links, kp)
+	}
+
+	kp, err = link.Kprobe("veth_convert_skb_to_xdp_buff", coll.Programs["kprobe_veth_convert_skb_to_xdp_buff"], nil)
+	if err != nil {
+		if !errors.Is(err, os.ErrNotExist) {
+			log.Fatalf("Opening kprobe veth_convert_skb_to_xdp_buff: %s\n", err)
+		}
+	} else {
+		t.links = append(t.links, kp)
+	}
+
 	if haveFexit && trackSkbClone {
 		progs := []*ebpf.Program{
 			coll.Programs["fexit_skb_clone"],
diff --git a/main.go b/main.go
index a7cc5b7c..06341821 100644
--- a/main.go
+++ b/main.go
@@ -126,9 +126,12 @@ func main() {
 
 	for name, program := range bpfSpec.Programs {
 		// 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" {
+		switch name {
+		case "kprobe_skb_lifetime_termination",
+			"fexit_skb_clone",
+			"fexit_skb_copy",
+			"kprobe_veth_convert_skb_to_xdp_buff",
+			"kretprobe_veth_convert_skb_to_xdp_buff":
 			continue
 		}
 		if name == "fentry_xdp" {