When bpf_throw() unwinds to the exception callback of the main program,
the frames between the throw site and the main program are discarded
without their epilogues running, so the callee-saved registers they
clobbered have to be restored some other way. Doing that from the frames
themselves needs, for each JITed program, where its prologue saved the
caller's R6..R9: walking from the throw site outwards, a register is
either still in the physical register, when no deeper frame saved it, or
sits in the prologue spill area of the nearest deeper frame that did.

Add callee_reg_off[4] to struct bpf_prog_aux, holding the offset of each
saved register from the native frame pointer that arch_bpf_stack_walk()
reports for the frame, or 0 when the register is not saved, and fill it
in from the x86-64 and arm64 JITs. Nothing consumes it yet.

On x86-64 the frame below rbp holds, from high to low addresses, the
program stack of round_up(stack_depth, 8) bytes (none with a private
stack), the tail call count and its pointer when tail_call_reachable,
r12 when the program is an exception boundary or uses an arena, and then
rbx, r13, r14 and r15 in push order, each only when it is pushed. An
exception boundary pushes all four.

On arm64 the frame record sits at x29, the tail call count pair right
below it, and the callee-saved registers follow as 16 byte pairs in the
order find_used_callee_regs() lists them, with the first register of a
pair at the lower address. The i-th pushed register is thus at
-32 - 16 * (i / 2) + 8 * (i % 2). An exception boundary pushes x19..x28
in order, so R6..R9 take the first four slots.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 arch/arm64/net/bpf_jit_comp.c | 37 +++++++++++++++++++++++++++++++++++
 arch/x86/net/bpf_jit_comp.c   | 33 +++++++++++++++++++++++++++++++
 include/linux/bpf.h           |  6 ++++++
 3 files changed, 76 insertions(+)

diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
index 1d5f4cf88d33..6a99f883a091 100644
--- a/arch/arm64/net/bpf_jit_comp.c
+++ b/arch/arm64/net/bpf_jit_comp.c
@@ -421,6 +421,42 @@ static void find_used_callee_regs(struct jit_ctx *ctx)
 	ctx->nr_used_callee_reg = i;
 }
 
+/*
+ * Offset from the frame pointer of the @idx-th register pushed by push_callee_regs(). The
+ * frame record sits at the frame pointer, the tail call count pair right below it, and
+ * then the pushed pairs, each 16 bytes with its first register at the lower address.
+ */
+static int callee_reg_slot_off(int idx)
+{
+	return -32 - 16 * (idx / 2) + 8 * (idx % 2);
+}
+
+/*
+ * Record where push_callee_regs() saved the caller's R6..R9, relative to the frame
+ * pointer that arch_bpf_stack_walk() reports for this frame, for the exception unwinder.
+ * An exception boundary pushes x19..x28 in order, so R6..R9 take the first four slots.
+ */
+static void record_callee_reg_off(struct jit_ctx *ctx)
+{
+	struct bpf_prog_aux *aux = ctx->prog->aux;
+	int i, j;
+
+	if (aux->exception_boundary) {
+		for (i = 0; i < 4; i++)
+			aux->callee_reg_off[i] = callee_reg_slot_off(i);
+		return;
+	}
+	for (i = 0; i < 4; i++) {
+		aux->callee_reg_off[i] = 0;
+		for (j = 0; j < ctx->nr_used_callee_reg; j++) {
+			if (ctx->used_callee_reg[j] != bpf2a64[BPF_REG_6 + i])
+				continue;
+			aux->callee_reg_off[i] = callee_reg_slot_off(j);
+			break;
+		}
+	}
+}
+
 /* Save callee-saved registers */
 static void push_callee_regs(struct jit_ctx *ctx)
 {
@@ -451,6 +487,7 @@ static void push_callee_regs(struct jit_ctx *ctx)
 			emit(A64_PUSH(reg1, A64_ZR, A64_SP), ctx);
 		}
 	}
+	record_callee_reg_off(ctx);
 }
 
 /* Restore callee-saved registers */
diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index ea97e068efda..2d7495438d18 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -427,6 +427,24 @@ static void pop_callee_regs(u8 **pprog, bool *callee_regs_used)
 	*pprog = prog;
 }
 
+/*
+ * Record where push_callee_regs() saved the caller's R6..R9 relative to rbp, for the
+ * exception unwinder. @off is the offset of rsp from rbp right before the first push.
+ */
+static void record_callee_reg_off(struct bpf_prog *prog, s32 off, bool *callee_regs_used)
+{
+	int i;
+
+	for (i = 0; i < 4; i++) {
+		if (!callee_regs_used[i]) {
+			prog->aux->callee_reg_off[i] = 0;
+			continue;
+		}
+		off -= 8;
+		prog->aux->callee_reg_off[i] = off;
+	}
+}
+
 /* add rsp, depth */
 static void emit_add_rsp(u8 **pprog, u32 depth)
 {
@@ -2530,6 +2548,7 @@ static int do_jit(struct bpf_verifier_env *env, struct bpf_prog *bpf_prog, int *
 	u32 stack_depth;
 	int callee_saved_size;
 	s32 outgoing_arg_base;
+	s32 callee_regs_off;
 	int err;
 
 	stack_depth = bpf_prog->aux->stack_depth;
@@ -2591,6 +2610,20 @@ static int do_jit(struct bpf_verifier_env *env, struct bpf_prog *bpf_prog, int *
 		push_callee_regs(&prog, callee_regs_used);
 	}
 
+	/*
+	 * Record where the pushes above saved the caller's R6..R9, relative to the
+	 * rbp that arch_bpf_stack_walk() reports for this frame. Below rbp lie, from
+	 * high to low addresses: round_up(stack_depth, 8) bytes of program stack
+	 * (none with a private stack), the tail call count and its pointer (16 bytes)
+	 * when tail_call_reachable, r12 when exception_boundary or arena_vm_start,
+	 * and then rbx, r13, r14 and r15 in push order, each only when pushed.
+	 */
+	callee_regs_off = -(s32)round_up(stack_depth, 8) - (tail_call_reachable ? 16 : 0);
+	if (bpf_prog->aux->exception_boundary || arena_vm_start)
+		callee_regs_off -= 8;
+	record_callee_reg_off(bpf_prog, callee_regs_off,
+			      bpf_prog->aux->exception_boundary ? all_callee_regs_used : callee_regs_used);
+
 	/* Compute callee-saved register area size. */
 	callee_saved_size = 0;
 	if (bpf_prog->aux->exception_boundary || arena_vm_start)
diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 7750df973494..010f14bb6c77 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1939,6 +1939,12 @@ struct bpf_prog_aux {
 	u64 (*bpf_exception_cb)(u64 cookie, u64 sp, u64 bp, u64, u64);
 	u16 stack_arg_sp_adjust;
 	u16 freplace_link_cnt; /* counts freplace links extending this prog */
+	/*
+	 * Offsets from the native frame pointer of this JITed program at which
+	 * its prologue saved the caller's R6..R9, or 0 when that register is not
+	 * saved. Filled by the JIT; consumed by the exception unwinder.
+	 */
+	s16 callee_reg_off[4];
 #ifdef CONFIG_SECURITY
 	void *security;
 #endif
-- 
2.53.0