[RFC PATCH bpf-next 4/7] bpf: Add kfuncs with bodies for common operations

[Yusheng Zheng <yunwei356@gmail.com>]

Add kfuncs with bodies for the operation families evaluated in [1]:
bpf_rol64() (rotate), bpf_select64() (conditional select),
bpf_extract64() (bit field extract), bpf_load_be64() (big-endian load),
bpf_prefetch(), bpf_copy16() (16-byte copy) and bpf_lea64() (address
computation). They are in kernel/bpf/insn_kfuncs/, apart from the
verifier and the JITs, under the new CONFIG_BPF_INSN_KFUNCS, which can
be built in or as a module. The x86-64 JIT inlines copies of their
compiled code.

A prefetch is a load whose value is not used, so the verifier checks the
address like that of any load. bpf_copy16() loads both halves before it
stores them, like its body, so that a copy of it and the body agree
when the buffers overlap.

[1] https://arxiv.org/abs/2606.24213

Assisted-by: LLM
Signed-off-by: Yusheng Zheng <yunwei356@gmail.com>
---
 kernel/bpf/Kconfig                       |   1 +
 kernel/bpf/Makefile                      |   1 +
 kernel/bpf/insn_kfuncs/Kconfig           |  11 ++
 kernel/bpf/insn_kfuncs/Makefile          |   2 +
 kernel/bpf/insn_kfuncs/bpf_insn_kfuncs.c | 173 +++++++++++++++++++++++
 5 files changed, 188 insertions(+)
 create mode 100644 kernel/bpf/insn_kfuncs/Kconfig
 create mode 100644 kernel/bpf/insn_kfuncs/Makefile
 create mode 100644 kernel/bpf/insn_kfuncs/bpf_insn_kfuncs.c

diff --git a/kernel/bpf/Kconfig b/kernel/bpf/Kconfig
index 98493e32db2ad..ca0dbfa1f81c5 100644
--- a/kernel/bpf/Kconfig
+++ b/kernel/bpf/Kconfig
@@ -101,6 +101,7 @@ config BPF_CRYPTO
 	  data. The supported algorithms are AES-CBC and AES-ECB.
 
 source "kernel/bpf/preload/Kconfig"
+source "kernel/bpf/insn_kfuncs/Kconfig"
 
 config BPF_LSM
 	bool "Enable BPF LSM Instrumentation"
diff --git a/kernel/bpf/Makefile b/kernel/bpf/Makefile
index ae3d04dae2d33..94c81491a8e65 100644
--- a/kernel/bpf/Makefile
+++ b/kernel/bpf/Makefile
@@ -60,6 +60,7 @@ obj-${CONFIG_BPF_LSM} += bpf_lsm_proto.o bpf_lsm.o
 endif
 obj-$(CONFIG_BPF_CRYPTO) += crypto.o
 obj-$(CONFIG_BPF_PRELOAD) += preload/
+obj-$(CONFIG_BPF_INSN_KFUNCS) += insn_kfuncs/
 
 obj-$(CONFIG_BPF_SYSCALL) += relo_core.o
 obj-$(CONFIG_BPF_SYSCALL) += btf_iter.o
diff --git a/kernel/bpf/insn_kfuncs/Kconfig b/kernel/bpf/insn_kfuncs/Kconfig
new file mode 100644
index 0000000000000..b2469853899dc
--- /dev/null
+++ b/kernel/bpf/insn_kfuncs/Kconfig
@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: GPL-2.0-only
+config BPF_INSN_KFUNCS
+	tristate "Kfuncs for CPU instructions that BPF lacks"
+	depends on BPF_SYSCALL && BPF_JIT && DEBUG_INFO_BTF
+	help
+	  Provide kfuncs for rotate, conditional select, bit field extract,
+	  big-endian load, prefetch, 16-byte copy and address computation.
+	  The verifier checks each call as the kfunc's BPF body, and the JIT
+	  puts native code in its place where it can.
+
+	  If unsure, say N.
diff --git a/kernel/bpf/insn_kfuncs/Makefile b/kernel/bpf/insn_kfuncs/Makefile
new file mode 100644
index 0000000000000..e397f626a17a8
--- /dev/null
+++ b/kernel/bpf/insn_kfuncs/Makefile
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: GPL-2.0
+obj-$(CONFIG_BPF_INSN_KFUNCS) += bpf_insn_kfuncs.o
diff --git a/kernel/bpf/insn_kfuncs/bpf_insn_kfuncs.c b/kernel/bpf/insn_kfuncs/bpf_insn_kfuncs.c
new file mode 100644
index 0000000000000..1bcc049a5603f
--- /dev/null
+++ b/kernel/bpf/insn_kfuncs/bpf_insn_kfuncs.c
@@ -0,0 +1,173 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Kfuncs for CPU instructions that BPF lacks, with BPF bodies, see struct
+ * bpf_kfunc_body. Arguments named __k must be known constants.
+ */
+#include <linux/bitops.h>
+#include <linux/btf.h>
+#include <linux/btf_ids.h>
+#include <linux/filter.h>
+#include <linux/module.h>
+#include <linux/unaligned.h>
+
+__bpf_kfunc_start_defs();
+
+__bpf_kfunc u64 bpf_rol64(u64 x, u32 n__k)
+{
+	return rol64(x, n__k);
+}
+
+__bpf_kfunc u64 bpf_select64(u64 cond, u64 a, u64 b)
+{
+	return cond ? a : b;
+}
+
+__bpf_kfunc u64 bpf_extract64(u64 x, u32 start__k, u32 len__k)
+{
+	return x << (64 - start__k - len__k) >> (64 - len__k);
+}
+
+__bpf_kfunc u64 bpf_load_be64(const void *p, s32 off__k)
+{
+	return get_unaligned_be64(p + off__k);
+}
+
+__bpf_kfunc void bpf_prefetch(const void *p)
+{
+	/* not prefetch(), which boot-time alternatives may rewrite */
+	__builtin_prefetch(p);
+}
+
+__bpf_kfunc void bpf_copy16(void *dst, const void *src)
+{
+	/* both loads first, as in the body, in case the buffers overlap */
+	u64 lo = get_unaligned((const u64 *)src);
+	u64 hi = get_unaligned((const u64 *)(src + 8));
+
+	put_unaligned(lo, (u64 *)dst);
+	put_unaligned(hi, (u64 *)(dst + 8));
+}
+
+__bpf_kfunc u64 bpf_lea64(u64 base, u64 index, u32 scale__k, s32 disp__k)
+{
+	return base + index * scale__k + disp__k;
+}
+
+__bpf_kfunc_end_defs();
+
+/* x << n | x >> (-n & 63) */
+static const struct bpf_insn rol64_body[] = {
+	BPF_ALU64_IMM(BPF_AND, BPF_REG_2, 63),
+	BPF_MOV64_REG(BPF_REG_0, BPF_REG_1),
+	BPF_ALU64_REG(BPF_LSH, BPF_REG_0, BPF_REG_2),
+	BPF_ALU64_IMM(BPF_NEG, BPF_REG_2, 0),
+	BPF_ALU64_IMM(BPF_AND, BPF_REG_2, 63),
+	BPF_ALU64_REG(BPF_RSH, BPF_REG_1, BPF_REG_2),
+	BPF_ALU64_REG(BPF_OR, BPF_REG_0, BPF_REG_1),
+};
+
+/* the jump lands within the body, here on its last instruction */
+static const struct bpf_insn select64_body[] = {
+	BPF_JMP_IMM(BPF_JNE, BPF_REG_1, 0, 1),
+	BPF_MOV64_REG(BPF_REG_2, BPF_REG_3),
+	BPF_MOV64_REG(BPF_REG_0, BPF_REG_2),
+};
+
+/* x << (64 - start - len) >> (64 - len) */
+static const struct bpf_insn extract64_body[] = {
+	BPF_MOV32_IMM(BPF_REG_4, 64),
+	BPF_ALU32_REG(BPF_SUB, BPF_REG_4, BPF_REG_2),
+	BPF_ALU32_REG(BPF_SUB, BPF_REG_4, BPF_REG_3),
+	BPF_MOV64_REG(BPF_REG_0, BPF_REG_1),
+	BPF_ALU64_REG(BPF_LSH, BPF_REG_0, BPF_REG_4),
+	BPF_MOV32_IMM(BPF_REG_4, 64),
+	BPF_ALU32_REG(BPF_SUB, BPF_REG_4, BPF_REG_3),
+	BPF_ALU64_REG(BPF_RSH, BPF_REG_0, BPF_REG_4),
+};
+
+/* the offset is an s32 */
+static const struct bpf_insn load_be64_body[] = {
+	BPF_ALU64_IMM(BPF_LSH, BPF_REG_2, 32),
+	BPF_ALU64_IMM(BPF_ARSH, BPF_REG_2, 32),
+	BPF_ALU64_REG(BPF_ADD, BPF_REG_1, BPF_REG_2),
+	BPF_LDX_MEM(BPF_DW, BPF_REG_0, BPF_REG_1, 0),
+	BPF_ENDIAN(BPF_TO_BE, BPF_REG_0, 64),
+};
+
+/* a load whose value is not used, of memory that the program may read */
+static const struct bpf_insn prefetch_body[] = {
+	BPF_LDX_MEM(BPF_B, BPF_REG_1, BPF_REG_1, 0),
+};
+
+/* two loads, then two stores */
+static const struct bpf_insn copy16_body[] = {
+	BPF_LDX_MEM(BPF_DW, BPF_REG_4, BPF_REG_2, 0),
+	BPF_LDX_MEM(BPF_DW, BPF_REG_5, BPF_REG_2, 8),
+	BPF_STX_MEM(BPF_DW, BPF_REG_1, BPF_REG_4, 0),
+	BPF_STX_MEM(BPF_DW, BPF_REG_1, BPF_REG_5, 8),
+};
+
+/* base + index * scale + disp, scale a u32 and disp an s32 */
+static const struct bpf_insn lea64_body[] = {
+	BPF_MOV32_REG(BPF_REG_3, BPF_REG_3),
+	BPF_MOV64_REG(BPF_REG_0, BPF_REG_2),
+	BPF_ALU64_REG(BPF_MUL, BPF_REG_0, BPF_REG_3),
+	BPF_ALU64_REG(BPF_ADD, BPF_REG_0, BPF_REG_1),
+	BPF_ALU64_IMM(BPF_LSH, BPF_REG_4, 32),
+	BPF_ALU64_IMM(BPF_ARSH, BPF_REG_4, 32),
+	BPF_ALU64_REG(BPF_ADD, BPF_REG_0, BPF_REG_4),
+};
+
+BTF_KFUNCS_START(insn_kfunc_ids)
+BTF_ID_FLAGS(func, bpf_rol64)
+BTF_ID_FLAGS(func, bpf_select64)
+BTF_ID_FLAGS(func, bpf_extract64)
+BTF_ID_FLAGS(func, bpf_load_be64)
+BTF_ID_FLAGS(func, bpf_prefetch)
+BTF_ID_FLAGS(func, bpf_copy16)
+BTF_ID_FLAGS(func, bpf_lea64)
+BTF_KFUNCS_END(insn_kfunc_ids)
+
+BTF_ID_LIST(body_ids)
+BTF_ID(func, bpf_rol64)
+BTF_ID(func, bpf_select64)
+BTF_ID(func, bpf_extract64)
+BTF_ID(func, bpf_load_be64)
+BTF_ID(func, bpf_prefetch)
+BTF_ID(func, bpf_copy16)
+BTF_ID(func, bpf_lea64)
+
+#define BODY(i, op, emit)	{ &body_ids[i], op##_body, ARRAY_SIZE(op##_body), emit }
+
+static const struct bpf_kfunc_body bodies[] = {
+	BODY(0, rol64, NULL),
+	BODY(1, select64, NULL),
+	BODY(2, extract64, NULL),
+	BODY(3, load_be64, NULL),
+	BODY(4, prefetch, NULL),
+	BODY(5, copy16, NULL),
+	BODY(6, lea64, NULL),
+};
+
+static const struct btf_kfunc_id_set insn_kfunc_set = {
+	.owner    = THIS_MODULE,
+	.set      = &insn_kfunc_ids,
+	.bodies   = bodies,
+	.body_cnt = ARRAY_SIZE(bodies),
+};
+
+static int __init insn_kfuncs_init(void)
+{
+	return register_btf_kfunc_id_set(BPF_PROG_TYPE_UNSPEC, &insn_kfunc_set);
+}
+
+/* the kfunc set goes away with the module's BTF */
+static void __exit insn_kfuncs_exit(void)
+{
+}
+
+module_init(insn_kfuncs_init);
+module_exit(insn_kfuncs_exit);
+
+MODULE_DESCRIPTION("Kfuncs for CPU instructions that BPF lacks");
+MODULE_LICENSE("GPL");
-- 
2.51.1