Re: [PATCH bpf-next 2/4] bpf, arm64: Fix tailcall infinite loop caused by freplace

[Leon Hwang <leon.hwang@linux.dev>]

On 30/8/24 15:37, Xu Kuohai wrote:
> On 8/27/2024 10:23 AM, Leon Hwang wrote:
>>

[...]

> 
> I think the complexity arises from having to decide whether
> to initialize or keep the tail counter value in the prologue.
> 
> To get rid of this complexity, a straightforward idea is to
> move the tail call counter initialization to the entry of
> bpf world, and in the bpf world, we only increase and check
> the tail call counter, never save/restore or set it. The
> "entry of the bpf world" here refers to mechanisms like
> bpf_prog_run, bpf dispatcher, or bpf trampoline that
> allows bpf prog to be invoked from C function.
> 
> Below is a rough POC diff for arm64 that could pass all
> of your tests. The tail call counter is held in callee-saved
> register x26, and is set to 0 by arch_run_bpf.
> 
> diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
> index 8aa32cb140b9..2c0f7daf1655 100644
> --- a/arch/arm64/net/bpf_jit_comp.c
> +++ b/arch/arm64/net/bpf_jit_comp.c
> @@ -26,7 +26,7 @@
> 
>  #define TMP_REG_1 (MAX_BPF_JIT_REG + 0)
>  #define TMP_REG_2 (MAX_BPF_JIT_REG + 1)
> -#define TCCNT_PTR (MAX_BPF_JIT_REG + 2)
> +#define TCALL_CNT (MAX_BPF_JIT_REG + 2)
>  #define TMP_REG_3 (MAX_BPF_JIT_REG + 3)
>  #define ARENA_VM_START (MAX_BPF_JIT_REG + 5)
> 
> @@ -63,7 +63,7 @@ static const int bpf2a64[] = {
>      [TMP_REG_2] = A64_R(11),
>      [TMP_REG_3] = A64_R(12),
>      /* tail_call_cnt_ptr */
> -    [TCCNT_PTR] = A64_R(26),
> +    [TCALL_CNT] = A64_R(26), // x26 is used to hold tail call counter
>      /* temporary register for blinding constants */
>      [BPF_REG_AX] = A64_R(9),
>      /* callee saved register for kern_vm_start address */
> @@ -286,19 +286,6 @@ static bool is_lsi_offset(int offset, int scale)
>   *      // PROLOGUE_OFFSET
>   *    // save callee-saved registers
>   */
> -static void prepare_bpf_tail_call_cnt(struct jit_ctx *ctx)
> -{
> -    const bool is_main_prog = !bpf_is_subprog(ctx->prog);
> -    const u8 ptr = bpf2a64[TCCNT_PTR];
> -
> -    if (is_main_prog) {
> -        /* Initialize tail_call_cnt. */
> -        emit(A64_PUSH(A64_ZR, ptr, A64_SP), ctx);
> -        emit(A64_MOV(1, ptr, A64_SP), ctx);
> -    } else
> -        emit(A64_PUSH(ptr, ptr, A64_SP), ctx);
> -}
> -
>  static void find_used_callee_regs(struct jit_ctx *ctx)
>  {
>      int i;
> @@ -419,7 +406,7 @@ static void pop_callee_regs(struct jit_ctx *ctx)
>  #define POKE_OFFSET (BTI_INSNS + 1)
> 
>  /* Tail call offset to jump into */
> -#define PROLOGUE_OFFSET (BTI_INSNS + 2 + PAC_INSNS + 4)
> +#define PROLOGUE_OFFSET (BTI_INSNS + 2 + PAC_INSNS + 2)
> 
>  static int build_prologue(struct jit_ctx *ctx, bool ebpf_from_cbpf)
>  {
> @@ -473,8 +460,6 @@ static int build_prologue(struct jit_ctx *ctx, bool
> ebpf_from_cbpf)
>          emit(A64_PUSH(A64_FP, A64_LR, A64_SP), ctx);
>          emit(A64_MOV(1, A64_FP, A64_SP), ctx);
> 
> -        prepare_bpf_tail_call_cnt(ctx);
> -
>          if (!ebpf_from_cbpf && is_main_prog) {
>              cur_offset = ctx->idx - idx0;
>              if (cur_offset != PROLOGUE_OFFSET) {
> @@ -499,7 +484,7 @@ static int build_prologue(struct jit_ctx *ctx, bool
> ebpf_from_cbpf)
>           *
>           * 12 registers are on the stack
>           */
> -        emit(A64_SUB_I(1, A64_SP, A64_FP, 96), ctx);
> +        emit(A64_SUB_I(1, A64_SP, A64_FP, 80), ctx);
>      }
> 
>      if (ctx->fp_used)
> @@ -527,8 +512,7 @@ static int emit_bpf_tail_call(struct jit_ctx *ctx)
> 
>      const u8 tmp = bpf2a64[TMP_REG_1];
>      const u8 prg = bpf2a64[TMP_REG_2];
> -    const u8 tcc = bpf2a64[TMP_REG_3];
> -    const u8 ptr = bpf2a64[TCCNT_PTR];
> +    const u8 tcc = bpf2a64[TCALL_CNT];
>      size_t off;
>      __le32 *branch1 = NULL;
>      __le32 *branch2 = NULL;
> @@ -546,16 +530,15 @@ static int emit_bpf_tail_call(struct jit_ctx *ctx)
>      emit(A64_NOP, ctx);
> 
>      /*
> -     * if ((*tail_call_cnt_ptr) >= MAX_TAIL_CALL_CNT)
> +     * if (tail_call_cnt >= MAX_TAIL_CALL_CNT)
>       *     goto out;
>       */
>      emit_a64_mov_i64(tmp, MAX_TAIL_CALL_CNT, ctx);
> -    emit(A64_LDR64I(tcc, ptr, 0), ctx);
>      emit(A64_CMP(1, tcc, tmp), ctx);
>      branch2 = ctx->image + ctx->idx;
>      emit(A64_NOP, ctx);
> 
> -    /* (*tail_call_cnt_ptr)++; */
> +    /* tail_call_cnt++; */
>      emit(A64_ADD_I(1, tcc, tcc, 1), ctx);
> 
>      /* prog = array->ptrs[index];
> @@ -570,9 +553,6 @@ static int emit_bpf_tail_call(struct jit_ctx *ctx)
>      branch3 = ctx->image + ctx->idx;
>      emit(A64_NOP, ctx);
> 
> -    /* Update tail_call_cnt if the slot is populated. */
> -    emit(A64_STR64I(tcc, ptr, 0), ctx);
> -
>      /* restore SP */
>      if (ctx->stack_size)
>          emit(A64_ADD_I(1, A64_SP, A64_SP, ctx->stack_size), ctx);
> @@ -793,6 +773,27 @@ asm (
>  "    .popsection\n"
>  );
> 
> +unsigned int arch_run_bpf(const void *ctx, const struct bpf_insn
> *insnsi, bpf_func_t bpf_func);
> +asm (
> +"    .pushsection .text, \"ax\", @progbits\n"
> +"    .global arch_run_bpf\n"
> +"    .type arch_run_bpf, %function\n"
> +"arch_run_bpf:\n"
> +#if IS_ENABLED(CONFIG_ARM64_BTI_KERNEL)
> +"    bti j\n"
> +#endif
> +"    stp x29, x30, [sp, #-16]!\n"
> +"    stp xzr, x26, [sp, #-16]!\n"
> +"    mov x26, #0\n"
> +"    blr x2\n"
> +"    ldp xzr, x26, [sp], #16\n"
> +"    ldp x29, x30, [sp], #16\n"
> +"    ret x30\n"
> +"    .size arch_run_bpf, . - arch_run_bpf\n"
> +"    .popsection\n"
> +);
> +EXPORT_SYMBOL_GPL(arch_run_bpf);
> +
>  /* build a plt initialized like this:
>   *
>   * plt:
> @@ -826,7 +827,6 @@ static void build_plt(struct jit_ctx *ctx)
>  static void build_epilogue(struct jit_ctx *ctx)
>  {
>      const u8 r0 = bpf2a64[BPF_REG_0];
> -    const u8 ptr = bpf2a64[TCCNT_PTR];
> 
>      /* We're done with BPF stack */
>      if (ctx->stack_size)
> @@ -834,8 +834,6 @@ static void build_epilogue(struct jit_ctx *ctx)
> 
>      pop_callee_regs(ctx);
> 
> -    emit(A64_POP(A64_ZR, ptr, A64_SP), ctx);
> -
>      /* Restore FP/LR registers */
>      emit(A64_POP(A64_FP, A64_LR, A64_SP), ctx);
> 
> @@ -2066,6 +2064,8 @@ static int prepare_trampoline(struct jit_ctx *ctx,
> struct bpf_tramp_image *im,
>      bool save_ret;
>      __le32 **branches = NULL;
> 
> +    bool target_is_bpf = is_bpf_text_address((unsigned long)func_addr);
> +
>      /* trampoline stack layout:
>       *                  [ parent ip         ]
>       *                  [ FP                ]
> @@ -2133,6 +2133,11 @@ static int prepare_trampoline(struct jit_ctx
> *ctx, struct bpf_tramp_image *im,
>       */
>      emit_bti(A64_BTI_JC, ctx);
> 
> +    if (!target_is_bpf) {
> +        emit(A64_PUSH(A64_ZR, A64_R(26), A64_SP), ctx);
> +        emit(A64_MOVZ(1, A64_R(26), 0, 0), ctx);
> +    }
> +
>      /* frame for parent function */
>      emit(A64_PUSH(A64_FP, A64_R(9), A64_SP), ctx);
>      emit(A64_MOV(1, A64_FP, A64_SP), ctx);
> @@ -2226,6 +2231,8 @@ static int prepare_trampoline(struct jit_ctx *ctx,
> struct bpf_tramp_image *im,
>      /* pop frames  */
>      emit(A64_POP(A64_FP, A64_LR, A64_SP), ctx);
>      emit(A64_POP(A64_FP, A64_R(9), A64_SP), ctx);
> +    if (!target_is_bpf)
> +        emit(A64_POP(A64_ZR, A64_R(26), A64_SP), ctx);
> 
>      if (flags & BPF_TRAMP_F_SKIP_FRAME) {
>          /* skip patched function, return to parent */
> diff --git a/include/linux/bpf.h b/include/linux/bpf.h
> index dc63083f76b7..8660d15dd50c 100644
> --- a/include/linux/bpf.h
> +++ b/include/linux/bpf.h
> @@ -1244,12 +1244,14 @@ struct bpf_dispatcher {
>  #define __bpfcall __nocfi
>  #endif
> 
> +unsigned int arch_run_bpf(const void *ctx, const struct bpf_insn
> *insnsi, bpf_func_t bpf_func);
> +
>  static __always_inline __bpfcall unsigned int bpf_dispatcher_nop_func(
>      const void *ctx,
>      const struct bpf_insn *insnsi,
>      bpf_func_t bpf_func)
>  {
> -    return bpf_func(ctx, insnsi);
> +    return arch_run_bpf(ctx, insnsi, bpf_func);
>  }
> 
>  /* the implementation of the opaque uapi struct bpf_dynptr */
> @@ -1317,7 +1319,7 @@ int arch_prepare_bpf_dispatcher(void *image, void
> *buf, s64 *funcs, int num_func
>  #else
>  #define __BPF_DISPATCHER_SC_INIT(name)
>  #define __BPF_DISPATCHER_SC(name)
> -#define __BPF_DISPATCHER_CALL(name)        bpf_func(ctx, insnsi)
> +#define __BPF_DISPATCHER_CALL(name)        arch_run_bpf(ctx, insnsi,
> bpf_func);
>  #define __BPF_DISPATCHER_UPDATE(_d, _new)
>  #endif
> 

This approach is really cool!

I want an alike approach on x86. But I failed. Because, on x86, it's an
indirect call to "call *rdx", aka "bpf_func(ctx, insnsi)".

Let us imagine the arch_run_bpf() on x86:

unsigned int __naked arch_run_bpf(const void *ctx, const struct bpf_insn
*insnsi, bpf_func_t bpf_func)
{
	asm (
		"pushq %rbp\n\t"
		"movq %rsp, %rbp\n\t"
		"xor %rax, %rax\n\t"
		"pushq %rax\n\t"
		"movq %rsp, %rax\n\t"
		"callq *%rdx\n\t"
		"leave\n\t"
		"ret\n\t"
	);
}

If we can change "callq *%rdx" to a direct call, it'll be really
wonderful to resolve this tailcall issue on x86.

How to introduce arch_bpf_run() for all JIT backends?

Thanks,
Leon