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

[Josef Bacik <josef@xxxxxxxxxxxxxx>]

Emit an increment of current->rcu_tramp_nesting once the trampoline's
frame is set up and a decrement before the final register restore, so
that a task preempted while running fentry/fexit/fmod_ret/LSM programs
or the __bpf_tramp_enter()/__bpf_tramp_exit() glue is not treated as
Tasks-RCU quiescent.  Drop the count around the call to the original
function: that may run arbitrarily long without sleeping and must not pin
a Tasks RCU grace period, and the trampoline frame above it is held by
im->pcref rather than by Tasks RCU (see bpf_tramp_image_put()).  The
fmod_ret early-exit branch and the ip_after_call -> ip_epilogue poke both
skip the decrement/increment pair around the original call, so the count
stays balanced on every path.

The sequence is "mov r11, gs:[current_task]; inc/dec dword [r11 + off]";
r11 is scratch at every emission point and (u32)&current_task is a valid
sign-extended %gs-absolute with the current per-CPU layout, the same form
the JIT already uses for this_cpu_off.  The image is dynamically
allocated text, so the instructions outside the bracketed region are
covered by the irq-exit IP check.

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

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index 2853e87797a7..a375c1b7bd50 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -722,6 +722,31 @@ static void emit_indirect_jump(u8 **pprog, int bpf_reg, u8 
*ip)
        *pprog = prog;
 }
 
+/*
+ * Tasks RCU trampoline nesting, see rcu_tasks_trampoline_enter().
+ *
+ *   mov r11, QWORD PTR gs:[current_task]
+ *   inc/dec DWORD PTR [r11 + offsetof(struct task_struct, rcu_tramp_nesting)]
+ *
+ * r11 (AUX_REG) is scratch in the trampoline at every point this is emitted.
+ */
+static void emit_rcu_tasks_tramp_nesting(u8 **pprog, bool enter)
+{
+#ifdef CONFIG_TASKS_RCU
+       u8 *prog = *pprog;
+
+       /* mov r11, gs:[abs32] */
+       EMIT2(0x65, 0x4C);
+       EMIT3(0x8B, 0x1C, 0x25);
+       EMIT((u32)(unsigned long)&current_task, 4);
+       /* inc/dec dword ptr [r11 + disp32] */
+       EMIT3(0x41, 0xFF, enter ? 0x83 : 0x8B);
+       EMIT(offsetof(struct task_struct, rcu_tramp_nesting), 4);
+
+       *pprog = prog;
+#endif
+}
+
 static void emit_return(u8 **pprog, u8 *ip)
 {
        u8 *prog = *pprog;
@@ -3610,6 +3635,13 @@ static int __arch_prepare_bpf_trampoline(struct 
bpf_tramp_image *im, void *rw_im
        /* mov QWORD PTR [rbp - rbx_off], rbx */
        emit_stx(&prog, BPF_DW, BPF_REG_FP, BPF_REG_6, -rbx_off);
 
+       /*
+        * From here until the matching decrement before the final return, 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(&prog, true);
+
        func_meta = nr_regs;
        /* Store number of argument registers of the traced function */
        emit_store_stack_imm64(&prog, BPF_REG_0, -func_meta_off, func_meta);
@@ -3670,6 +3702,13 @@ static int __arch_prepare_bpf_trampoline(struct 
bpf_tramp_image *im, void *rw_im
                        LOAD_TRAMP_TAIL_CALL_CNT_PTR(stack_size);
                }
 
+               /*
+                * 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(&prog, false);
                if (flags & BPF_TRAMP_F_ORIG_STACK) {
                        emit_ldx(&prog, BPF_DW, BPF_REG_6, BPF_REG_FP, 8);
                        EMIT2(0xff, 0xd3); /* call *rbx */
@@ -3680,6 +3719,7 @@ static int __arch_prepare_bpf_trampoline(struct 
bpf_tramp_image *im, void *rw_im
                                goto cleanup;
                        }
                }
+               emit_rcu_tasks_tramp_nesting(&prog, true);
                /* remember return value in a stack for bpf prog to access */
                emit_stx(&prog, BPF_DW, BPF_REG_FP, BPF_REG_0, -8);
                im->ip_after_call = image + (prog - (u8 *)rw_image);
@@ -3741,6 +3781,9 @@ static int __arch_prepare_bpf_trampoline(struct 
bpf_tramp_image *im, void *rw_im
        if (save_ret)
                emit_ldx(&prog, BPF_DW, BPF_REG_0, BPF_REG_FP, -8);
 
+       /* Remaining instructions are covered by the irq-exit IP check. */
+       emit_rcu_tasks_tramp_nesting(&prog, false);
+
        emit_ldx(&prog, BPF_DW, BPF_REG_6, BPF_REG_FP, -rbx_off);
 
        EMIT1(0xC9); /* leave */

-- 
2.55.0