[PATCH RFC v2 10/15] bpf, arm64: Maintain Tasks RCU trampoline nesting in the BPF trampoline

[Josef Bacik <josef@xxxxxxxxxxxxxx>]

Same scheme as x86: emit "mrs x10, sp_el0; ldr/add|sub/str w11" to bump
current->rcu_tramp_nesting after the callee-saved registers are stored
and to drop it before they are restored, and release it around the call
to the original function, which im->pcref protects and which must not
pin a Tasks RCU grace period.  x10/x11 are scratch at every emission
point; the fmod_ret cbnz target lies after the decrement/increment pair
around the original call, and the ip_after_call nop follows the
re-increment, so the count is balanced on every path.  BUILD_BUG_ON
guards the LDR/STR immediate range for the task_struct offset.

Assisted-by: LLM
Signed-off-by: Josef Bacik <josef@xxxxxxxxxxxxxx>
---
 arch/arm64/net/bpf_jit_comp.c | 46 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
index c18e005a41db..5c9a7bde5cc9 100644
--- a/arch/arm64/net/bpf_jit_comp.c
+++ b/arch/arm64/net/bpf_jit_comp.c
@@ -2591,6 +2591,34 @@ static void emit_arena_arg_conv(struct jit_ctx *ctx, u8 
dst, u8 src, bool nullab
        emit(A64_SUB(0, dst, src, base_lo), ctx);
 }
 
+/*
+ * Tasks RCU trampoline nesting, see rcu_tasks_trampoline_enter().
+ *
+ *   mrs  x10, sp_el0
+ *   ldr  w11, [x10, #offsetof(struct task_struct, rcu_tramp_nesting)]
+ *   add/sub w11, w11, #1
+ *   str  w11, [x10, #...]
+ *
+ * x10/x11 are scratch in the trampoline at every point this is emitted.
+ */
+static void emit_rcu_tasks_tramp_nesting(struct jit_ctx *ctx, bool enter)
+{
+#ifdef CONFIG_TASKS_RCU
+       const int off = offsetof(struct task_struct, rcu_tramp_nesting);
+       const u8 tsk = A64_R(10), cnt = A64_R(11);
+
+       BUILD_BUG_ON(off & 3 || off >= SZ_16K); /* LDR/STR (imm12, scaled) */
+
+       emit(A64_MRS_SP_EL0(tsk), ctx);
+       emit(A64_LDR32I(cnt, tsk, off), ctx);
+       if (enter)
+               emit(A64_ADD_I(0, cnt, cnt, 1), ctx);
+       else
+               emit(A64_SUB_I(0, cnt, cnt, 1), ctx);
+       emit(A64_STR32I(cnt, tsk, off), ctx);
+#endif
+}
+
 static void save_args(struct jit_ctx *ctx, int bargs_off, int oargs_off,
                      const struct btf_func_model *m, const struct arg_aux *a,
                      bool for_call_origin, bool is_struct_ops, u64 arena_base)
@@ -2854,6 +2882,13 @@ static int prepare_trampoline(struct jit_ctx *ctx, 
struct bpf_tramp_image *im,
        emit(A64_STR64I(A64_R(19), A64_SP, regs_off), ctx);
        emit(A64_STR64I(A64_R(20), A64_SP, regs_off + 8), ctx);
 
+       /*
+        * From here until the matching decrement in the epilogue, a preemption
+        * of this task is not a Tasks RCU quiescent state.  The instructions
+        * above this point are covered by the irq-exit IP check.
+        */
+       emit_rcu_tasks_tramp_nesting(ctx, true);
+
        if (flags & BPF_TRAMP_F_CALL_ORIG) {
                /* for the first pass, assume the worst case */
                if (!ctx->image)
@@ -2898,12 +2933,20 @@ static int prepare_trampoline(struct jit_ctx *ctx, 
struct bpf_tramp_image *im,
        if (flags & BPF_TRAMP_F_CALL_ORIG) {
                /* the original func takes kernel addresses, never converted 
ones */
                save_args(ctx, bargs_off, oargs_off, m, a, true, is_struct_ops, 
0);
+               /*
+                * The original function may run for a long time without
+                * sleeping; do not let it pin a Tasks RCU grace period.  The
+                * trampoline frame above it is held by im->pcref
+                * (__bpf_tramp_enter()), not by Tasks RCU, across the call.
+                */
+               emit_rcu_tasks_tramp_nesting(ctx, false);
                /* call original func */
                emit(A64_LDR64I(A64_R(10), A64_SP, retaddr_off), ctx);
                emit(A64_ADR(A64_LR, AARCH64_INSN_SIZE * 2), ctx);
                emit(A64_RET(A64_R(10)), ctx);
                /* store return value */
                emit(A64_STR64I(A64_R(0), A64_SP, retval_off), ctx);
+               emit_rcu_tasks_tramp_nesting(ctx, true);
                /* reserve a nop for bpf_tramp_image_put */
                im->ip_after_call = ctx->ro_image + ctx->idx;
                emit(A64_NOP, ctx);
@@ -2945,6 +2988,9 @@ static int prepare_trampoline(struct jit_ctx *ctx, struct 
bpf_tramp_image *im,
        if (flags & BPF_TRAMP_F_RESTORE_REGS)
                restore_args(ctx, bargs_off, a->regs_for_args);
 
+       /* Remaining instructions are covered by the irq-exit IP check. */
+       emit_rcu_tasks_tramp_nesting(ctx, false);
+
        /* restore callee saved register x19 and x20 */
        emit(A64_LDR64I(A64_R(19), A64_SP, regs_off), ctx);
        emit(A64_LDR64I(A64_R(20), A64_SP, regs_off + 8), ctx);

-- 
2.55.0