bpf: Treat load and store through a number as arena access

There are no address spaces in Rust. LLVM emits plain loads and stores
for arena memory, without cast_kern:

  r3 = 0x20 ll            // R_BPF_64_64 .bss, libbpf puts it into arena
  r2 = *(u64 *)(r3 + 0)   // R3 invalid mem access 'scalar'
  lock *(u64 *)(r2 + 8) += r1

Add BPF_F_ARENA_SCALAR flag of BPF_PROG_LOAD. When the program is loaded
with it and has an arena treat ldx, stx, st and atomics through a number
as arena access. It's as safe as access through PTR_TO_ARENA: JIT adds
the base of arena to the low 32 bits of the address. The program has
an arena only with CAP_BPF and CAP_PERFMON.

Registers of the program are not changed, since Rust compares and
stores the address after the access. bpf_do_misc_fixups() copies
the low 32 bits into BPF_REG_AX and the access goes through it:

  r2 = *(u64 *)(r3 + 0)  ->  w12 = w3
                             r2 = *(u64 *)(r12 + 0)

Constant blinding needs BPF_REG_AX for immediates and leaves insns that
use BPF_REG_AX alone. So the immediate of st through a number is not
blinded, the rest of the program is:

  *(u64 *)(r3 + 0) = 1   ->  w12 = w3
                             *(u64 *)(r12 + 0) = 1

The same insn may see PTR_TO_ARENA on another path. It works for both.

It's a flag, since a number is also what the verifier makes of a pointer
that went away when the program has CAP_PERFMON: a ringbuf record after
bpf_ringbuf_submit(), a pointer to the packet after bpf_skb_pull_data().
Programs in C that have an arena keep "invalid mem access 'scalar'" for
such bugs.

JIT tells with bpf_jit_supports_arena_scalar() that it takes BPF_REG_AX
as the address of arena access. With other JITs the flag is rejected
with -EOPNOTSUPP. x86 and arm64 do:
- x86: the fault handler finds the register that holds the address
  through reg2pt_regs[]. Add BPF_REG_AX, r10 of x86, there.
- arm64: nothing else is needed. The JIT takes any register as the
  address and the fault handler reads it by its number. Compile tested
  only.

Any number is an address of arena, NULL and small numbers included, like
it is for PTR_TO_ARENA: arena code in C dereferences NULL and relies on
the fault being handled.

Still rejected:
 - insn that sees a number on one path and a pointer that is not
   PTR_TO_ARENA on another.
 - numbers passed to helpers and kfuncs, except __arena arguments.

Signed-off-by: Alexei Starovoitov <ast@kernel.org>
Link: https://lore.kernel.org/bpf/20261002124714.180012-4-alexei.starovoitov@gmail.com
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>

4 files changed, 98 insertions, 12 deletions

diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index a67c895ce734..92a87c73aa44 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -3435,6 +3435,12 @@ bool __weak bpf_jit_supports_arena(void)
 	return false;
 }
 
+/* Whether JIT takes BPF_REG_AX as the address of arena access */
+bool __weak bpf_jit_supports_arena_scalar(void)
+{
+	return false;
+}
+
 bool __weak bpf_jit_supports_insn(struct bpf_insn *insn, bool in_arena)
 {
 	return false;
diff --git a/kernel/bpf/fixups.c b/kernel/bpf/fixups.c
index cb7219ff68bd..206cc9a614a9 100644
--- a/kernel/bpf/fixups.c
+++ b/kernel/bpf/fixups.c
@@ -44,6 +44,31 @@ static bool is_addr_space_cast32(struct bpf_prog *prog, const struct bpf_insn *i
 	return false;
 }
 
+/*
+ * The insn accesses arena through a number. JITs add the base of arena
+ * to the register as it is, so the access goes through the low 32 bits of
+ * the number in BPF_REG_AX. Registers of the program are not changed: the
+ * number may be compared or stored later, and on another path the register
+ * may be PTR_TO_ARENA.
+ *
+ * Constant blinding leaves insns that use BPF_REG_AX alone, so the immediate
+ * of st through a number is not blinded.
+ */
+static int arena_scalar_access(const struct bpf_insn *insn, struct bpf_insn *buf)
+{
+	bool load = BPF_CLASS(insn->code) == BPF_LDX || bpf_atomic_is_load_acq(insn);
+	struct bpf_insn *patch = buf;
+
+	*patch++ = BPF_MOV32_REG(BPF_REG_AX, load ? insn->src_reg : insn->dst_reg);
+	*patch = *insn;
+	if (load)
+		patch->src_reg = BPF_REG_AX;
+	else
+		patch->dst_reg = BPF_REG_AX;
+	patch++;
+	return patch - buf;
+}
+
 /* Return the regno defined by the insn, or -1. */
 static int insn_def_regno(const struct bpf_insn *insn)
 {
@@ -1736,6 +1761,20 @@ int bpf_do_misc_fixups(struct bpf_verifier_env *env)
 			/* Convert BPF_CLASS(insn->code) == BPF_ALU64 to 32-bit ALU */
 			insn->code = BPF_ALU | BPF_OP(insn->code) | BPF_SRC(insn->code);
 
+		if (env->insn_aux_data[i + delta].arena_scalar) {
+			cnt = arena_scalar_access(insn, insn_buf);
+
+			new_prog = bpf_patch_insn_data(env, i + delta, insn_buf, cnt);
+			if (!new_prog)
+				return -ENOMEM;
+
+			delta += cnt - 1;
+			prog = new_prog;
+			env->prog = prog;
+			insn = prog->insnsi + i + delta;
+			goto next_insn;
+		}
+
 		/* Make sdiv/smod divide-by-minus-one exceptions impossible. */
 		if ((insn->code == (BPF_ALU64 | BPF_MOD | BPF_K) ||
 		     insn->code == (BPF_ALU64 | BPF_DIV | BPF_K) ||
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index 4e5377447e1b..654f896af265 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -2972,9 +2972,13 @@ static int bpf_prog_load(union bpf_attr *attr, bpfptr_t uattr, struct bpf_log_at
 				 BPF_F_XDP_HAS_FRAGS |
 				 BPF_F_XDP_DEV_BOUND_ONLY |
 				 BPF_F_TEST_REG_INVARIANTS |
+				 BPF_F_ARENA_SCALAR |
 				 BPF_F_TOKEN_FD))
 		return -EINVAL;
 
+	if ((attr->prog_flags & BPF_F_ARENA_SCALAR) && !bpf_jit_supports_arena_scalar())
+		return -EOPNOTSUPP;
+
 	bpf_prog_load_fixup_attach_type(attr);
 
 	if (attr->prog_flags & BPF_F_TOKEN_FD) {
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 866c0540c68a..fd3c0206bd67 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -5260,6 +5260,19 @@ static bool is_arena_reg(struct bpf_verifier_env *env, int regno)
 	return reg->type == PTR_TO_ARENA;
 }
 
+/*
+ * There are no address spaces in Rust, addresses of arena are plain numbers.
+ * When the program is loaded with BPF_F_ARENA_SCALAR and has an arena a load or
+ * a store through a number is an access to arena at the low 32 bits of it,
+ * like the access through PTR_TO_ARENA is. The register stays a number.
+ */
+static bool is_arena_scalar(struct bpf_verifier_env *env, int regno)
+{
+	const struct bpf_reg_state *reg = reg_state(env, regno);
+
+	return reg->type == SCALAR_VALUE && env->arena_scalar && env->prog->aux->arena;
+}
+
 static bool is_load_acq_unsafe(struct bpf_verifier_env *env, int regno,
 			       struct bpf_insn *insn)
 {
@@ -5290,7 +5303,7 @@ static bool atomic_ptr_type_ok(struct bpf_verifier_env *env, int regno,
 		return false;
 	if (is_sk_reg(env, regno))
 		return false;
-	if (is_arena_reg(env, regno))
+	if (is_arena_reg(env, regno) || is_arena_scalar(env, regno))
 		return bpf_jit_supports_insn(insn, true);
 	if (is_load_acq_unsafe(env, regno, insn))
 		return false;
@@ -7171,6 +7184,22 @@ static int check_mem_access(struct bpf_verifier_env *env, int insn_idx, struct b
 static int save_aux_ptr_type(struct bpf_verifier_env *env, enum bpf_reg_type type,
 			     bool allow_trust_mismatch);
 
+/*
+ * Returns the register to check the access of the current insn with.
+ * For a number in a program with an arena that is 'arena'.
+ */
+static struct bpf_reg_state *mem_access_reg(struct bpf_verifier_env *env, int regno,
+					    struct bpf_reg_state *arena)
+{
+	if (!is_arena_scalar(env, regno))
+		return cur_regs(env) + regno;
+
+	memset(arena, 0, sizeof(*arena));
+	arena->type = PTR_TO_ARENA;
+	env->insn_aux_data[env->insn_idx].arena_scalar = true;
+	return arena;
+}
+
 static int check_load_mem(struct bpf_verifier_env *env, struct bpf_insn *insn,
 			  bool strict_alignment_once, bool is_ldsx,
 			  bool allow_trust_mismatch, const char *ctx)
@@ -7178,6 +7207,7 @@ static int check_load_mem(struct bpf_verifier_env *env, struct bpf_insn *insn,
 	struct bpf_verifier_state *vstate = env->cur_state;
 	struct bpf_func_state *state = vstate->frame[vstate->curframe];
 	struct bpf_reg_state *regs = cur_regs(env);
+	struct bpf_reg_state arena, *src_reg;
 	enum bpf_reg_type src_reg_type;
 	int err;
 
@@ -7199,15 +7229,16 @@ static int check_load_mem(struct bpf_verifier_env *env, struct bpf_insn *insn,
 	if (err)
 		return err;
 
-	src_reg_type = regs[insn->src_reg].type;
+	src_reg = mem_access_reg(env, insn->src_reg, &arena);
+	src_reg_type = src_reg->type;
 
 	/*
 	 * check_stack_read_fixed_off() may refine the modification's origin to
 	 * the source stack slot.
 	 */
 	bpf_diag_mod_begin(env, &regs[insn->dst_reg], NULL, BPF_DIAG_MOD_WRITE);
-	err = check_mem_access(env, env->insn_idx, regs + insn->src_reg, argno_from_reg(insn->src_reg), insn->off,
-			       BPF_SIZE(insn->code), BPF_READ, insn->dst_reg,
+	err = check_mem_access(env, env->insn_idx, src_reg, argno_from_reg(insn->src_reg),
+			       insn->off, BPF_SIZE(insn->code), BPF_READ, insn->dst_reg,
 			       strict_alignment_once, is_ldsx);
 	err = err ?: save_aux_ptr_type(env, src_reg_type,
 				       allow_trust_mismatch);
@@ -7224,6 +7255,7 @@ static int check_store_reg(struct bpf_verifier_env *env, struct bpf_insn *insn,
 	struct bpf_verifier_state *vstate = env->cur_state;
 	struct bpf_func_state *state = vstate->frame[vstate->curframe];
 	struct bpf_reg_state *regs = cur_regs(env);
+	struct bpf_reg_state arena, *dst_reg;
 	enum bpf_reg_type dst_reg_type;
 	int err;
 
@@ -7245,11 +7277,12 @@ static int check_store_reg(struct bpf_verifier_env *env, struct bpf_insn *insn,
 	if (err)
 		return err;
 
-	dst_reg_type = regs[insn->dst_reg].type;
+	dst_reg = mem_access_reg(env, insn->dst_reg, &arena);
+	dst_reg_type = dst_reg->type;
 
 	/* Check if (dst_reg + off) is writeable. */
-	err = check_mem_access(env, env->insn_idx, regs + insn->dst_reg, argno_from_reg(insn->dst_reg), insn->off,
-			       BPF_SIZE(insn->code), BPF_WRITE, insn->src_reg,
+	err = check_mem_access(env, env->insn_idx, dst_reg, argno_from_reg(insn->dst_reg),
+			       insn->off, BPF_SIZE(insn->code), BPF_WRITE, insn->src_reg,
 			       strict_alignment_once, false);
 	err = err ?: save_aux_ptr_type(env, dst_reg_type, false);
 
@@ -7259,7 +7292,7 @@ static int check_store_reg(struct bpf_verifier_env *env, struct bpf_insn *insn,
 static int check_atomic_rmw(struct bpf_verifier_env *env,
 			    struct bpf_insn *insn)
 {
-	struct bpf_reg_state *dst_reg;
+	struct bpf_reg_state arena, *dst_reg;
 	int load_reg;
 	int err;
 
@@ -7304,6 +7337,9 @@ static int check_atomic_rmw(struct bpf_verifier_env *env,
 		return -EACCES;
 	}
 
+	/* load_reg may be dst_reg. Look at dst_reg before it's marked as unknown. */
+	dst_reg = mem_access_reg(env, insn->dst_reg, &arena);
+
 	load_reg = bpf_atomic_load_reg(insn);
 	if (load_reg >= 0) {
 		/* check and record load of old value */
@@ -7312,8 +7348,6 @@ static int check_atomic_rmw(struct bpf_verifier_env *env,
 			return err;
 	}
 
-	dst_reg = cur_regs(env) + insn->dst_reg;
-
 	/* Check whether we can read the memory, with second call for fetch
 	 * case to simulate the register fill.
 	 */
@@ -19369,15 +19403,17 @@ static int do_check_insn(struct bpf_verifier_env *env, bool *do_print_state)
 			return check_stack_arg_write(env, state, insn->off, NULL);
 		}
 
+		struct bpf_reg_state arena, *dst_reg;
 		enum bpf_reg_type dst_reg_type;
 
 		err = check_reg_arg(env, insn->dst_reg, SRC_OP);
 		if (err)
 			return err;
 
-		dst_reg_type = cur_regs(env)[insn->dst_reg].type;
+		dst_reg = mem_access_reg(env, insn->dst_reg, &arena);
+		dst_reg_type = dst_reg->type;
 
-		err = check_mem_access(env, env->insn_idx, cur_regs(env) + insn->dst_reg, argno_from_reg(insn->dst_reg),
+		err = check_mem_access(env, env->insn_idx, dst_reg, argno_from_reg(insn->dst_reg),
 				       insn->off, BPF_SIZE(insn->code),
 				       BPF_WRITE, -1, false, false);
 		if (err)
@@ -22617,6 +22653,7 @@ int bpf_check(struct bpf_prog **prog, union bpf_attr *attr, bpfptr_t uattr,
 	if (is_priv)
 		env->test_state_freq = attr->prog_flags & BPF_F_TEST_STATE_FREQ;
 	env->test_reg_invariants = attr->prog_flags & BPF_F_TEST_REG_INVARIANTS;
+	env->arena_scalar = attr->prog_flags & BPF_F_ARENA_SCALAR;
 
 	env->explored_states = kvzalloc_objs(struct list_head,
 					     state_htab_size(env),