Track and output hard start of packet data

As we know, the hard start of packet data is same between `struct
sk_buff` and `struct xdp_buff`. So, we can track skb by it, and output
it as skb address.

Signed-off-by: Leon Hwang <hffilwlqm@gmail.com>

bpf/kprobe_pwru.c

@@ -289,27 +289,24 @@ set_output(void *ctx, struct sk_buff *skb, struct event_t *event) {
 
 static __noinline bool
 handle_everything(struct sk_buff *skb, void *ctx, struct event_t *event) {
-	bool tracked = false;
-	u64 skb_addr = (u64) skb;
+	u64 skb_addr = (u64) BPF_CORE_READ(skb, head);
 
 	if (cfg->is_set) {
-		if (cfg->track_skb && bpf_map_lookup_elem(&skb_addresses, &skb_addr)) {
-			tracked = true;
-			goto cont;
-		}
+		if (cfg->track_skb) {
+			if (!bpf_map_lookup_elem(&skb_addresses, &skb_addr)) {
+				if (!filter(skb))
+					return false;
+
+				bpf_map_update_elem(&skb_addresses, &skb_addr, &TRUE, BPF_ANY);
+			}
 
-		if (!filter(skb)) {
+		} else if (!filter(skb)) {
 			return false;
 		}
 
-cont:
 		set_output(ctx, skb, event);
 	}
 
-	if (cfg->track_skb && !tracked) {
-		bpf_map_update_elem(&skb_addresses, &skb_addr, &TRUE, BPF_ANY);
-	}
-
 	event->pid = bpf_get_current_pid_tgid() >> 32;
 	event->ts = bpf_ktime_get_ns();
 	event->cpu_id = bpf_get_smp_processor_id();
@@ -324,7 +321,7 @@ kprobe_skb(struct sk_buff *skb, struct pt_regs *ctx, bool has_get_func_ip) {
 	if (!handle_everything(skb, ctx, &event))
 		return BPF_OK;
 
-	event.skb_addr = (u64) skb;
+	event.skb_addr = (u64) BPF_CORE_READ(skb, head);
 	event.addr = has_get_func_ip ? bpf_get_func_ip(ctx) : PT_REGS_IP(ctx);
 	event.param_second = PT_REGS_PARM2(ctx);
 	bpf_map_push_elem(&events, &event, BPF_EXIST);
@@ -359,29 +356,38 @@ 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);
+	struct sk_buff *skb = (void *) PT_REGS_PARM1(ctx);
+	u64 skb_addr = (u64) BPF_CORE_READ(skb, head);
 
-	bpf_map_delete_elem(&skb_addresses, &skb);
+	bpf_map_delete_elem(&skb_addresses, &skb_addr);
 
 	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);
+track_skb_clone(struct sk_buff *old, struct sk_buff *new) {
+	u64 skb_addr_old = (u64) BPF_CORE_READ(old, head);
+	u64 skb_addr_new = (u64) BPF_CORE_READ(new, head);
+	if (bpf_map_lookup_elem(&skb_addresses, &skb_addr_old))
+		bpf_map_update_elem(&skb_addresses, &skb_addr_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);
+int BPF_PROG(fexit_skb_clone, struct sk_buff *old, gfp_t mask, struct sk_buff *new) {
+	if (new)
+		return track_skb_clone(old, new);
+
+	return BPF_OK;
 }
 
 SEC("fexit/skb_copy")
-int BPF_PROG(fexit_skb_copy, u64 old, gfp_t mask, u64 new) {
-	return track_skb_clone(old, new);
+int BPF_PROG(fexit_skb_copy, struct sk_buff *old, gfp_t mask, struct sk_buff *new) {
+	if (new)
+		return track_skb_clone(old, new);
+
+	return BPF_OK;
 }
 
 SEC("fentry/tc")
@@ -391,7 +397,7 @@ int BPF_PROG(fentry_tc, struct sk_buff *skb) {
 	if (!handle_everything(skb, ctx, &event))
 		return BPF_OK;
 
-	event.skb_addr = (u64) skb;
+	event.skb_addr = (u64) BPF_CORE_READ(skb, head);
 	event.addr = BPF_PROG_ADDR;
 	event.type = EVENT_TYPE_TC;
 	bpf_map_push_elem(&events, &event, BPF_EXIST);
@@ -476,19 +482,29 @@ set_xdp_output(void *ctx, struct xdp_buff *xdp, struct event_t *event) {
 
 SEC("fentry/xdp")
 int BPF_PROG(fentry_xdp, struct xdp_buff *xdp) {
+	u64 skb_addr = (u64) BPF_CORE_READ(xdp, data_hard_start);
 	struct event_t event = {};
 
 	if (cfg->is_set) {
-		if (!filter_xdp(xdp))
+		if (cfg->track_skb) {
+			if (!bpf_map_lookup_elem(&skb_addresses, &skb_addr)) {
+				if (!filter_xdp(xdp))
+					return BPF_OK;
+
+				bpf_map_update_elem(&skb_addresses, &skb_addr, &TRUE, BPF_ANY);
+			}
+
+		} else if (!filter_xdp(xdp)) {
 			return BPF_OK;
+		}
 
 		set_xdp_output(ctx, xdp, &event);
 	}
 
 	event.pid = bpf_get_current_pid_tgid() >> 32;
 	event.ts = bpf_ktime_get_ns();
 	event.cpu_id = bpf_get_smp_processor_id();
-	event.skb_addr = (u64) xdp;
+	event.skb_addr = (u64) skb_addr;
 	event.addr = BPF_PROG_ADDR;
 	event.type = EVENT_TYPE_XDP;
 	bpf_map_push_elem(&events, &event, BPF_EXIST);

main.go

@@ -97,7 +97,8 @@ func main() {
 	if len(funcs) == 0 && !flags.FilterTraceTc && !flags.FilterTraceXdp {
 		log.Fatalf("Cannot find a matching kernel function")
 	}
-	// If --filter-trace-tc, it's to retrieve and print bpf prog's name.
+	// If --filter-trace-tc/--filter-trace-xdp, it's to retrieve and print bpf
+	// prog's name.
 	addr2name, name2addr, err := pwru.ParseKallsyms(funcs, flags.OutputStack ||
 		len(flags.KMods) != 0 || flags.FilterTraceTc || flags.FilterTraceXdp)
 	if err != nil {