[RFC PATCH bpf-next 2/3] bpf: Implement bpf_get_kern_btf_id() kfunc

[Yonghong Song <yhs@fb.com>]

The signature of bpf_get_kern_btf_id() function looks like
  void *bpf_get_kern_btf_id(obj, expected_btf_id)
The obj has a pointer type. The expected_btf_id is 0 or
a btf id to be returned by the kfunc. The function
currently supports two kinds of obj:
  - obj: ptr_to_ctx, expected_btf_id: 0
    return the expected kernel ctx btf id
  - obj: ptr to char/unsigned char, expected_btf_id: a struct btf id
    return expected_btf_id
The second case looks like a type casting, e.g., in kernel we have
  #define skb_shinfo(SKB) ((struct skb_shared_info *)(skb_end_pointer(SKB)))
bpf program can get a skb_shared_info btf id ptr with bpf_get_kern_btf_id()
kfunc.

The current prototype is incomplete in the following areas:
  - ptr to char/unsigned char is not supported. In btf_struct_walk(),
    for a pointer not pointing to a struct, a WALK_SCALAR is returned.
    this needs to be improved to recognize walking ptr to non-struct.
  - the current implementation didn't validate non-zero expected_btf_id
    parameter has to be a struct. This can be added easily later.
  - The forced type casting case may not be reliable, so the resulted
    ptr and later walked ptr cannot be passed to helper/kfunc's.
    We need to resolve this some how. David Vernet has some work ([1])
    in this area and it is not finalized yet.

  [1] https://lore.kernel.org/bpf/20221020222416.3415511-1-void@manifault.com/

Signed-off-by: Yonghong Song <yhs@fb.com>
---
 include/linux/bpf.h   |  2 ++
 kernel/bpf/btf.c      | 67 ++++++++++++++++++++++++++++++++++++++++++-
 kernel/bpf/helpers.c  |  5 ++++
 kernel/bpf/verifier.c |  8 +++++-
 4 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 798aec816970..d6a1c463b87e 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -2087,6 +2087,7 @@ struct bpf_kfunc_arg_meta {
 	bool r0_rdonly;
 	int ref_obj_id;
 	u32 flags;
+	u32 expected_ret_btf_id;
 };
 
 struct bpf_reg_state;
@@ -2113,6 +2114,7 @@ bool bpf_prog_has_kfunc_call(const struct bpf_prog *prog);
 const struct btf_func_model *
 bpf_jit_find_kfunc_model(const struct bpf_prog *prog,
 			 const struct bpf_insn *insn);
+void *bpf_get_kern_btf_id(void *, u32 expected_btf_id);
 struct bpf_core_ctx {
 	struct bpf_verifier_log *log;
 	const struct btf *btf;
diff --git a/kernel/bpf/btf.c b/kernel/bpf/btf.c
index 84f09235857c..b0a2555c8ca3 100644
--- a/kernel/bpf/btf.c
+++ b/kernel/bpf/btf.c
@@ -6309,6 +6309,9 @@ static bool btf_is_kfunc_arg_mem_size(const struct btf *btf,
 	return true;
 }
 
+BTF_ID_LIST_SINGLE(bpf_get_kern_btf_id_id, func, bpf_get_kern_btf_id)
+
+
 static int btf_check_func_arg_match(struct bpf_verifier_env *env,
 				    const struct btf *btf, u32 func_id,
 				    struct bpf_reg_state *regs,
@@ -6318,7 +6321,7 @@ static int btf_check_func_arg_match(struct bpf_verifier_env *env,
 {
 	enum bpf_prog_type prog_type = resolve_prog_type(env->prog);
 	bool rel = false, kptr_get = false, trusted_args = false;
-	bool sleepable = false;
+	bool sleepable = false, get_btf_id_kfunc = false;
 	struct bpf_verifier_log *log = &env->log;
 	u32 i, nargs, ref_id, ref_obj_id = 0;
 	bool is_kfunc = btf_is_kernel(btf);
@@ -6357,6 +6360,67 @@ static int btf_check_func_arg_match(struct bpf_verifier_env *env,
 		kptr_get = kfunc_meta->flags & KF_KPTR_GET;
 		trusted_args = kfunc_meta->flags & KF_TRUSTED_ARGS;
 		sleepable = kfunc_meta->flags & KF_SLEEPABLE;
+		get_btf_id_kfunc = func_id == *bpf_get_kern_btf_id_id;
+	}
+
+	/* special processing for bpf_get_btf_id kfunc.
+	 * arg1:
+	 *   must be a ptr_to_ctx or ptr_to_u8/s8.
+	 * arg2:
+	 *   must be a constant, if non-zero representing an user specified expected
+	 *   ret_btf_id.
+	 * If ptr_to_ctx, arg2 must be 0 or a value equals to corresponding kctx btf_id
+	 * and the ret ptr can be passed to a helper/kfunc. Otherwise, arg2 must be a
+	 * valid struct type btf_id, and the ret ptr cannot be passed to a helper/kfunc.
+	 */
+	if (get_btf_id_kfunc) {
+		struct bpf_reg_state *reg = &regs[1];
+		int kctx_btf_id = 0;
+		s64 val;
+
+		if (nargs != 2) {
+			bpf_log(log, "Incorrect number of arguments %s, actual %d expect 2\n",
+				func_name, nargs);
+			return -EINVAL;
+		}
+
+		/* arg1 */
+		if (reg->type == PTR_TO_CTX) {
+			enum bpf_prog_type prog_type = resolve_prog_type(env->prog);
+			const struct btf_type *conv_struct;
+			const struct btf_member *ctx_type;
+
+			conv_struct = bpf_ctx_convert.t;
+			ctx_type = btf_type_member(conv_struct) + bpf_ctx_convert_map[prog_type] * 2;
+			ctx_type++;
+
+			/* find the kctx type */
+			kctx_btf_id = ctx_type->type;
+		} else if (reg->type != SCALAR_VALUE) {
+			/* FIXME: we actually expects a pointer to char/unsigned_char */
+			bpf_log(log, "Incorrect type %x\n", reg->type);
+			return -EINVAL;
+		}
+
+		reg = &regs[2];
+		if (!tnum_is_const(reg->var_off)) {
+			bpf_log(log, "arg 2 is not constant\n");
+			return -EINVAL;
+		}
+
+		val = reg->var_off.value;
+		if (kctx_btf_id == 0) {
+			/* FIXME: ensure val is a btf_id pointing to a struct */
+			kctx_btf_id = val;
+		} else {
+			if (val != 0 && val != kctx_btf_id) {
+				bpf_log(log, "Incorrect expected_btf_id %lld, expect 0 or %d\n", val, kctx_btf_id);
+				return -EINVAL;
+			}
+		}
+
+		kfunc_meta->expected_ret_btf_id = kctx_btf_id;
+		goto check;
 	}
 
 	/* check that BTF function arguments match actual types that the
@@ -6639,6 +6703,7 @@ static int btf_check_func_arg_match(struct bpf_verifier_env *env,
 		return -EINVAL;
 	}
 
+check:
 	if (sleepable && !env->prog->aux->sleepable) {
 		bpf_log(log, "kernel function %s is sleepable but the program is not\n",
 			func_name);
diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index 2f11e22eefba..f8995947a790 100644
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -1717,7 +1717,12 @@ static const struct btf_kfunc_id_set tracing_kfunc_set = {
 	.set   = &tracing_btf_ids,
 };
 
+void *bpf_get_kern_btf_id(void *obj, u32 expected_btf_id) {
+	return obj;
+}
+
 BTF_SET8_START(generic_btf_ids)
+BTF_ID_FLAGS(func, bpf_get_kern_btf_id)
 BTF_SET8_END(generic_btf_ids)
 
 static const struct btf_kfunc_id_set generic_kfunc_set = {
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 07c0259dfc1a..cf284af382a3 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -7868,7 +7868,13 @@ static int check_kfunc_call(struct bpf_verifier_env *env, struct bpf_insn *insn,
 	} else if (btf_type_is_ptr(t)) {
 		ptr_type = btf_type_skip_modifiers(desc_btf, t->type,
 						   &ptr_type_id);
-		if (!btf_type_is_struct(ptr_type)) {
+		if (meta.expected_ret_btf_id) {
+			mark_reg_known_zero(env, regs, BPF_REG_0);
+			regs[BPF_REG_0].btf = desc_btf;
+			regs[BPF_REG_0].type = PTR_TO_BTF_ID;
+			regs[BPF_REG_0].btf_id = meta.expected_ret_btf_id;
+
+		} else if (!btf_type_is_struct(ptr_type)) {
 			if (!meta.r0_size) {
 				ptr_type_name = btf_name_by_offset(desc_btf,
 								   ptr_type->name_off);
-- 
2.30.2