[RFC PATCH bpf-next v1 11/16] libbpf: Support the typed_arena_cast instruction

[Kumar Kartikeya Dwivedi <memxor@gmail.com>]

The typed_arena_cast instruction turns any 64-bit value into a pointer to an
object of a program-BTF struct in the struct's typed arena, the kernel-only
region of an arena map. It is encoded as a 64-bit register move with an off
of BPF_TYPED_ARENA_CAST and the struct's local BTF type ID in imm, and the
kernel takes the ID from there, so a compiler that emits the instruction
records a BPF_CORE_TYPE_ID_LOCAL relocation against it, the way it does for
the ld_imm64 behind __builtin_btf_type_id().

Let the relocation patcher write a local type ID into that instruction. It
is the one ALU form with a register source that carries a relocation, and it
takes no other relocation kind; nothing else about the patching of ALU
immediates changes.

Add bpf_typed_arena_cast() to bpf_helpers.h, wrapping the compiler builtin
under its feature macro, so that programs cast through one name whatever the
toolchain.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 tools/lib/bpf/bpf_helpers.h | 13 +++++++++++++
 tools/lib/bpf/relo_core.c   | 19 +++++++++++++++++--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/tools/lib/bpf/bpf_helpers.h b/tools/lib/bpf/bpf_helpers.h
index 9d160b5b9c0e..1a440e32ebeb 100644
--- a/tools/lib/bpf/bpf_helpers.h
+++ b/tools/lib/bpf/bpf_helpers.h
@@ -340,6 +340,19 @@ enum libbpf_tristate {
 /* Helper macro to print out debug messages */
 #define bpf_printk(fmt, args...) ___bpf_pick_printk(args)(fmt, ##args)
 
+/*
+ * Turn any 64-bit value into a pointer to an object of the given struct type
+ * in its typed arena, with the typed_arena_cast instruction: the kernel masks
+ * the value into the struct's slice of the arena map's typed region, so the
+ * result names an object of that type whatever the value was, and the
+ * verifier trusts it. The struct's local BTF type ID is relocated into the
+ * instruction. Requires compiler support.
+ */
+#ifdef __BPF_FEATURE_TYPED_ARENA_CAST
+#define bpf_typed_arena_cast(v, type)					\
+	((typeof(type) *)__builtin_bpf_typed_arena_cast((v), *(typeof(type) *)0))
+#endif
+
 struct bpf_iter_num;
 
 extern int bpf_iter_num_new(struct bpf_iter_num *it, int start, int end) __weak __ksym;
diff --git a/tools/lib/bpf/relo_core.c b/tools/lib/bpf/relo_core.c
index 2672623a4198..e872e432729b 100644
--- a/tools/lib/bpf/relo_core.c
+++ b/tools/lib/bpf/relo_core.c
@@ -1028,6 +1028,12 @@ static int insn_bytes_to_bpf_size(__u32 sz)
 	}
 }
 
+/* rX = typed_arena_cast(rY, <local type ID in imm>) */
+static bool is_typed_arena_cast_insn(struct bpf_insn *insn)
+{
+	return insn->code == (BPF_ALU64 | BPF_MOV | BPF_X) && insn->off == BPF_TYPED_ARENA_CAST;
+}
+
 /*
  * Patch relocatable BPF instruction.
  *
@@ -1043,7 +1049,8 @@ static int insn_bytes_to_bpf_size(__u32 sz)
  * 3. rX = <imm64> (load with 64-bit immediate value);
  * 4. rX = *(T *)(rY + <off>), where T is one of {u8, u16, u32, u64};
  * 5. *(T *)(rX + <off>) = rY, where T is one of {u8, u16, u32, u64};
- * 6. *(T *)(rX + <off>) = <imm>, where T is one of {u8, u16, u32, u64}.
+ * 6. *(T *)(rX + <off>) = <imm>, where T is one of {u8, u16, u32, u64};
+ * 7. rX = typed_arena_cast(rY, <imm>), for a local type ID relocation only.
  */
 int bpf_core_patch_insn(const char *prog_name, struct bpf_insn *insn,
 			int insn_idx, const struct bpf_core_relo *relo,
@@ -1060,8 +1067,16 @@ int bpf_core_patch_insn(const char *prog_name, struct bpf_insn *insn,
 	switch (class) {
 	case BPF_ALU:
 	case BPF_ALU64:
-		if (BPF_SRC(insn->code) != BPF_K)
+		/*
+		 * The typed arena cast is a register move whose imm names the
+		 * struct by its local type ID; the kernel takes the ID from there.
+		 */
+		if (is_typed_arena_cast_insn(insn)) {
+			if (relo->kind != BPF_CORE_TYPE_ID_LOCAL)
+				goto bad_insn;
+		} else if (BPF_SRC(insn->code) != BPF_K) {
 			goto bad_insn;
+		}
 		if (res->poison)
 			return bpf_core_poison_insn(prog_name, relo_idx, insn, insn_idx);
 		if (res->validate && insn->imm != orig_val) {
-- 
2.53.0