[PATCH bpf-next v1 3/6] bpf: Add bpf_throw_unwind() and a hook for bpf_throw() entry

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

bpf_throw() walks the stack to the frame of the main program and calls
its exception callback, which resets the stack to that frame and returns
from the program. Restoring the callee-saved registers of the frames the
walk discards needs their values at the throw site, and by the time C
code runs, the compiler may have clobbered them.

Move the body of bpf_throw() into bpf_throw_unwind(), which takes the
callee-saved registers of the throw site, R6..R9, along with the cookie.
An architecture that provides a native entry for the kfunc captures them
before calling it. bpf_arch_throw_entry() returns that entry, and the
verifier retargets the call to bpf_throw() there when it specializes the
kfunc, so the JIT keeps finding the function model of the call through
the specialized descriptor. Without an entry, bpf_throw() calls the
unwinder with NULL, as before. Nothing consumes the registers yet.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 include/linux/bpf.h       |  1 +
 include/linux/bpf_throw.h | 33 +++++++++++++++++++++++++++++++++
 kernel/bpf/core.c         |  9 +++++++++
 kernel/bpf/helpers.c      | 16 +++++++++++++---
 kernel/bpf/verifier.c     |  5 +++++
 5 files changed, 61 insertions(+), 3 deletions(-)
 create mode 100644 include/linux/bpf_throw.h

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 010f14bb6c77..839c28d3743f 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -7,6 +7,7 @@
 #include <uapi/linux/bpf.h>
 #include <uapi/linux/filter.h>
 #include <linux/bpf_defs.h>
+#include <linux/bpf_throw.h>
 
 #include <crypto/sha2.h>
 #include <linux/workqueue.h>
diff --git a/include/linux/bpf_throw.h b/include/linux/bpf_throw.h
new file mode 100644
index 000000000000..663776715b33
--- /dev/null
+++ b/include/linux/bpf_throw.h
@@ -0,0 +1,33 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * The interface between the architecture entries of bpf_throw() and the
+ * generic unwinder.
+ */
+#ifndef _LINUX_BPF_THROW_H
+#define _LINUX_BPF_THROW_H
+
+/* The callee-saved registers the unwinder tracks: R6..R9. */
+#define BPF_THROW_NR_CALLEE_SAVED	4
+/* The size of struct bpf_throw_regs, for the entries. */
+#define BPF_THROW_REGS_SIZE		(8 * BPF_THROW_NR_CALLEE_SAVED)
+
+#ifndef __ASSEMBLER__
+#include <linux/build_bug.h>
+#include <linux/types.h>
+
+/*
+ * The callee-saved registers captured at a throw site, indexed like
+ * bpf_prog_aux::callee_reg_off.
+ */
+struct bpf_throw_regs {
+	u64 saved[BPF_THROW_NR_CALLEE_SAVED];
+};
+static_assert(sizeof(struct bpf_throw_regs) == BPF_THROW_REGS_SIZE);
+
+void bpf_throw_unwind(struct bpf_throw_regs *regs, u64 cookie);
+/* The native entry of bpf_throw(), on an architecture that has one. */
+void bpf_throw_entry(void);
+void *bpf_arch_throw_entry(void);
+#endif /* __ASSEMBLER__ */
+
+#endif /* _LINUX_BPF_THROW_H */
diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index fc22d66d194a..0205a687e1b5 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -3601,6 +3601,15 @@ void __weak arch_bpf_stack_walk(bool (*consume_fn)(void *cookie, u64 ip, u64 sp,
 {
 }
 
+/*
+ * The native entry of bpf_throw(), which captures the BPF register file of
+ * the throw site, or NULL when the architecture calls the kfunc directly.
+ */
+void * __weak bpf_arch_throw_entry(void)
+{
+	return NULL;
+}
+
 bool __weak bpf_jit_supports_timed_may_goto(void)
 {
 	return false;
diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index 648abe1ba63a..7c78630e4dd3 100644
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -3427,7 +3427,12 @@ static bool bpf_stack_walker(void *cookie, u64 ip, u64 sp, u64 bp)
 	return false;
 }
 
-__bpf_kfunc void bpf_throw(u64 cookie)
+/*
+ * Unwind the stack to the exception callback of the main program. @regs holds
+ * the callee-saved registers of the throw site, or is NULL when the
+ * architecture has no native entry and the kfunc was called directly.
+ */
+void bpf_throw_unwind(struct bpf_throw_regs *regs, u64 cookie)
 {
 	struct bpf_throw_ctx ctx = {};
 
@@ -3439,8 +3444,8 @@ __bpf_kfunc void bpf_throw(u64 cookie)
 	WARN_ON_ONCE(!ctx.cnt);
 	/*
 	 * Prevent KASAN false positives for CONFIG_KASAN_STACK by unpoisoning
-	 * deeper stack depths than ctx.sp as we do not return from bpf_throw,
-	 * which skips compiler generated instrumentation to do the same. Some
+	 * deeper stack depths than ctx.sp as we do not return from here, which
+	 * skips compiler generated instrumentation to do the same. Some
 	 * architectures cannot recover sp while unwinding, so fall back to bp.
 	 */
 	kasan_unpoison_task_stack_below((void *)(long)(ctx.sp ?: ctx.bp));
@@ -3448,6 +3453,11 @@ __bpf_kfunc void bpf_throw(u64 cookie)
 	WARN(1, "A call to BPF exception callback should never return\n");
 }
 
+__bpf_kfunc void bpf_throw(u64 cookie)
+{
+	bpf_throw_unwind(NULL, cookie);
+}
+
 __bpf_kfunc int bpf_wq_init(struct bpf_wq *wq, void *p__const_map, unsigned int flags)
 {
 	struct bpf_async_kern *async = (struct bpf_async_kern *)wq;
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 01434e1d23b5..d980530e58bb 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -23162,6 +23162,11 @@ static int specialize_kfunc(struct bpf_verifier_env *env, struct bpf_kfunc_desc
 	} else if (func_id == special_kfunc_list[KF_bpf_arena_free_pages]) {
 		if (env->insn_aux_data[insn_idx].non_sleepable)
 			addr = (unsigned long)bpf_arena_free_pages_non_sleepable;
+	} else if (func_id == special_kfunc_list[KF_bpf_throw]) {
+		void *entry = bpf_arch_throw_entry();
+
+		if (entry)
+			addr = (unsigned long)entry;
 	}
 	desc->addr = addr;
 	return 0;
-- 
2.53.0