[PATCH bpf-next 1/4] bpf: Allow get_func_[arg|arg_cnt] helpers in raw tracepoint programs

[KaFai Wan]

Adding support to use get_func_[arg|arg_cnt] helpers in raw_tp/tp_btf
programs.

We can use get_func_[arg|ret|arg_cnt] helpers in fentry/fexit/fmod_ret
programs currently. If we try to use get_func_[arg|arg_cnt] helpers in
raw_tp/tp_btf programs, verifier will fail to load the program with:

; __u64 cnt = bpf_get_func_arg_cnt(ctx);
3: (85) call bpf_get_func_arg_cnt#185
unknown func bpf_get_func_arg_cnt#185

Adding get_func_[arg|arg_cnt] helpers in raw_tp_prog_func_proto and
tracing_prog_func_proto for raw tracepoint.

Adding 1 arg on ctx of raw tracepoint program and make it stores number of
arguments on ctx-8, so it's easy to verify argument index and find
argument's position.

Signed-off-by: KaFai Wan <mannkafai@xxxxxxxxx>
---
 kernel/trace/bpf_trace.c | 17 ++++++++++++++---
 net/bpf/test_run.c       | 13 +++++--------
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
index 52c432a44aeb..eb4c56013493 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -1892,6 +1892,10 @@ raw_tp_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 		return &bpf_get_stackid_proto_raw_tp;
 	case BPF_FUNC_get_stack:
 		return &bpf_get_stack_proto_raw_tp;
+	case BPF_FUNC_get_func_arg:
+		return &bpf_get_func_arg_proto;
+	case BPF_FUNC_get_func_arg_cnt:
+		return &bpf_get_func_arg_cnt_proto;
 	case BPF_FUNC_get_attach_cookie:
 		return &bpf_get_attach_cookie_proto_tracing;
 	default:
@@ -1950,10 +1954,16 @@ tracing_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 	case BPF_FUNC_d_path:
 		return &bpf_d_path_proto;
 	case BPF_FUNC_get_func_arg:
+		if (prog->type == BPF_PROG_TYPE_TRACING &&
+		    prog->expected_attach_type == BPF_TRACE_RAW_TP)
+			return &bpf_get_func_arg_proto;
 		return bpf_prog_has_trampoline(prog) ? &bpf_get_func_arg_proto : NULL;
 	case BPF_FUNC_get_func_ret:
 		return bpf_prog_has_trampoline(prog) ? &bpf_get_func_ret_proto : NULL;
 	case BPF_FUNC_get_func_arg_cnt:
+		if (prog->type == BPF_PROG_TYPE_TRACING &&
+		    prog->expected_attach_type == BPF_TRACE_RAW_TP)
+			return &bpf_get_func_arg_cnt_proto;
 		return bpf_prog_has_trampoline(prog) ? &bpf_get_func_arg_cnt_proto : NULL;
 	case BPF_FUNC_get_attach_cookie:
 		if (prog->type == BPF_PROG_TYPE_TRACING &&
@@ -2312,7 +2322,7 @@ void __bpf_trace_run(struct bpf_raw_tp_link *link, u64 *args)
 #define REPEAT(X, FN, DL, ...)		REPEAT_##X(FN, DL, __VA_ARGS__)
 
 #define SARG(X)		u64 arg##X
-#define COPY(X)		args[X] = arg##X
+#define COPY(X)		args[X + 1] = arg##X
 
 #define __DL_COM	(,)
 #define __DL_SEM	(;)
@@ -2323,9 +2333,10 @@ void __bpf_trace_run(struct bpf_raw_tp_link *link, u64 *args)
 	void bpf_trace_run##x(struct bpf_raw_tp_link *link,		\
 			      REPEAT(x, SARG, __DL_COM, __SEQ_0_11))	\
 	{								\
-		u64 args[x];						\
+		u64 args[x + 1];					\
+		args[0] = x;						\
 		REPEAT(x, COPY, __DL_SEM, __SEQ_0_11);			\
-		__bpf_trace_run(link, args);				\
+		__bpf_trace_run(link, args + 1);			\
 	}								\
 	EXPORT_SYMBOL_GPL(bpf_trace_run##x)
 BPF_TRACE_DEFN_x(1);
diff --git a/net/bpf/test_run.c b/net/bpf/test_run.c
index aaf13a7d58ed..8cb285187270 100644
--- a/net/bpf/test_run.c
+++ b/net/bpf/test_run.c
@@ -760,6 +760,7 @@ int bpf_prog_test_run_raw_tp(struct bpf_prog *prog,
 	void __user *ctx_in = u64_to_user_ptr(kattr->test.ctx_in);
 	__u32 ctx_size_in = kattr->test.ctx_size_in;
 	struct bpf_raw_tp_test_run_info info;
+	u64 args[MAX_BPF_FUNC_ARGS + 1] = {};
 	int cpu = kattr->test.cpu, err = 0;
 	int current_cpu;
 
@@ -776,14 +777,11 @@ int bpf_prog_test_run_raw_tp(struct bpf_prog *prog,
 	if ((kattr->test.flags & BPF_F_TEST_RUN_ON_CPU) == 0 && cpu != 0)
 		return -EINVAL;
 
-	if (ctx_size_in) {
-		info.ctx = memdup_user(ctx_in, ctx_size_in);
-		if (IS_ERR(info.ctx))
-			return PTR_ERR(info.ctx);
-	} else {
-		info.ctx = NULL;
-	}
+	if (ctx_size_in && copy_from_user(args + 1, ctx_in, ctx_size_in))
+		return -EFAULT;
 
+	args[0] = ctx_size_in / sizeof(u64);
+	info.ctx = args + 1;
 	info.prog = prog;
 
 	current_cpu = get_cpu();
@@ -807,7 +805,6 @@ int bpf_prog_test_run_raw_tp(struct bpf_prog *prog,
 	    copy_to_user(&uattr->test.retval, &info.retval, sizeof(u32)))
 		err = -EFAULT;
 
-	kfree(info.ctx);
 	return err;
 }
 
-- 
2.43.0