Fix tracing tc-bpf with --filter-track-skb-by-stackid

For tracing tc-bpf (same for XDP), it is able to get FP from r10 to
track skb by stack ID.

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

bpf/kprobe_pwru.c

@@ -395,9 +395,25 @@ set_shinfo_btf(struct sk_buff *skb, u64 *event_id) {
 }
 
 static __always_inline u64
-get_stackid(struct pt_regs *ctx) {
+get_tracing_fp(void)
+{
+	u64 fp;
+
+	/* get frame pointer */
+	asm volatile ("%[fp] = r10" : [fp] "+r"(fp) :);
+	return fp;
+}
+
+static __always_inline u64
+get_kprobe_fp(struct pt_regs *ctx)
+{
+	return PT_REGS_FP(ctx);
+}
+
+static __always_inline u64
+get_stackid(void *ctx, const bool is_kprobe) {
 	u64 caller_fp;
-	u64 fp = PT_REGS_FP(ctx);
+	u64 fp = is_kprobe ? get_kprobe_fp(ctx) : get_tracing_fp();
 	for (int depth = 0; depth < MAX_STACK_DEPTH; depth++) {
 		if (bpf_probe_read_kernel(&caller_fp, sizeof(caller_fp), (void *)fp) < 0)
 			break;
@@ -440,8 +456,8 @@ handle_everything(struct sk_buff *skb, void *ctx, struct event_t *event, u64 *_s
 	u64 skb_head = (u64) BPF_CORE_READ(skb, head);
 	u64 stackid;
 
-	if (cfg->track_skb_by_stackid && is_kprobe)
-		stackid = _stackid ? *_stackid : get_stackid(ctx);
+	if (cfg->track_skb_by_stackid)
+		stackid = _stackid ? *_stackid : get_stackid(ctx, is_kprobe);
 
 	if (cfg->is_set) {
 		if (cfg->track_xdp && cfg->track_skb) {
@@ -457,7 +473,7 @@ handle_everything(struct sk_buff *skb, void *ctx, struct event_t *event, u64 *_s
 			goto cont;
 		}
 
-		if (cfg->track_skb_by_stackid && is_kprobe && bpf_map_lookup_elem(&stackid_skb, &stackid)) {
+		if (cfg->track_skb_by_stackid && bpf_map_lookup_elem(&stackid_skb, &stackid)) {
 			tracked_by = TRACKED_BY_STACKID;
 			goto cont;
 		}
@@ -479,7 +495,7 @@ handle_everything(struct sk_buff *skb, void *ctx, struct event_t *event, u64 *_s
 			bpf_map_update_elem(&xdp_dhs_skb_heads, &skb_head, &skb_addr, BPF_ANY);
 	}
 
-	if (cfg->track_skb_by_stackid && is_kprobe && tracked_by != TRACKED_BY_STACKID) {
+	if (cfg->track_skb_by_stackid && tracked_by != TRACKED_BY_STACKID) {
 		u64 *old_stackid = bpf_map_lookup_elem(&skb_stackid, &skb);
 		if (old_stackid && *old_stackid != stackid) {
 			bpf_map_delete_elem(&stackid_skb, old_stackid);
@@ -537,7 +553,7 @@ PWRU_ADD_KPROBE(5)
 
 SEC("kprobe/skb_by_stackid")
 int kprobe_skb_by_stackid(struct pt_regs *ctx) {
-	u64 stackid = get_stackid(ctx);
+	u64 stackid = get_stackid(ctx, true);
 
 	struct sk_buff **skb = bpf_map_lookup_elem(&stackid_skb, &stackid);
 	if (skb && *skb)
@@ -558,7 +574,7 @@ int kprobe_skb_lifetime_termination(struct pt_regs *ctx) {
 	bpf_map_delete_elem(&skb_addresses, &skb_addr);
 
 	if (cfg->track_skb_by_stackid) {
-		u64 stackid = get_stackid(ctx);
+		u64 stackid = get_stackid(ctx, true);
 		bpf_map_delete_elem(&stackid_skb, &stackid);
 		bpf_map_delete_elem(&skb_stackid, &skb_addr);
 	}