bpf: Derive the atomic load register in one place

check_atomic_rmw() open codes the mapping from a BPF_ATOMIC to the register
it reads the old value into, the BPF_STX case of insn_def_regno() open codes
the very same mapping a second time, the const folding and the liveness
transfer functions a third and a fourth time, and BPF JITs need it as well
to know which register a faulting BPF_PROBE_ATOMIC has to clear.

Add a small helper so that all of them can share it. No functional change.
The BPF_LOAD_ACQ case is there for the JITs, which do walk all instruction
classes. const_reg_xfer() loses its explicit BPF_ATOMIC mode test since the
helper checks class and mode itself; the BPF_PROBE_ATOMIC it additionally
accepts cannot be seen there as it is only set from bpf_do_misc_fixups(),
that is, after const folding has run. arg_track_xfer() keeps its mode test
since that also guards the stack clearing next to it.

Signed-off-by: Daniel Borkmann <daniel@iogearbox.net>
Acked-by: Eduard Zingerman <eddyz87@gmail.com>
Link: https://patch.msgid.link/20260811131600.506721-1-daniel@iogearbox.net
Signed-off-by: Eduard Zingerman <eddyz87@gmail.com>

1 files changed, 24 insertions, 0 deletions

diff --git a/include/linux/filter.h b/include/linux/filter.h
index 4edba8182db1..15d83684c6e9 100644
--- a/include/linux/filter.h
+++ b/include/linux/filter.h
@@ -414,6 +414,30 @@ static inline bool bpf_atomic_is_load_acq(const struct bpf_insn *insn)
 	       insn->imm == BPF_LOAD_ACQ;
 }
 
+/*
+ * Given an instruction @insn, return the number of the BPF register that a
+ * BPF_ATOMIC reads the value at its memory operand into, or -1 if there is
+ * no such register. That is the register a BPF_PROBE_ATOMIC has to clear when
+ * the access faults. Like bpf_atomic_is_load_acq(), @insn is not assumed to
+ * be a BPF_ATOMIC here.
+ */
+static inline int bpf_atomic_load_reg(const struct bpf_insn *insn)
+{
+	if (BPF_CLASS(insn->code) != BPF_STX ||
+	    (BPF_MODE(insn->code) != BPF_ATOMIC &&
+	     BPF_MODE(insn->code) != BPF_PROBE_ATOMIC))
+		return -1;
+
+	switch (insn->imm) {
+	case BPF_LOAD_ACQ:
+		return insn->dst_reg;
+	case BPF_CMPXCHG:
+		return BPF_REG_0;
+	default:
+		return (insn->imm & BPF_FETCH) ? insn->src_reg : -1;
+	}
+}
+
 /* Memory store, *(uint *) (dst_reg + off16) = imm32 */
 
 #define BPF_ST_MEM(SIZE, DST, OFF, IMM)				\