// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */

#include <vmlinux.h>
#include <bpf/bpf_helpers.h>
#include "../test_kmods/bpf_testmod_kfunc.h"
#include "bpf_misc.h"

#if (defined(__TARGET_ARCH_x86) || defined(__TARGET_ARCH_arm64)) && \
        defined(__BPF_FEATURE_STACK_ARGUMENT)

/* Force kfunc extern BTF generation for inline asm call below.
 * Uses its own SEC so it's not included as a .text subprog.
 * The '?' prefix sets autoload=false so libbpf won't load it.
 */
SEC("?tc")
int __btf_kfunc_gen(struct __sk_buff *ctx)
{
        char buf[8] = {};

        return bpf_kfunc_call_stack_arg_mem(0, 0, 0, 0, 0, buf, sizeof(buf));
}

/*
 * Test precision backtracking across bpf-to-bpf call for kfunc stack arg.
 * subprog_call_mem_kfunc receives a size as incoming stack arg (arg6)
 * and forwards it as mem__sz (arg7) to bpf_kfunc_call_stack_arg_mem.
 */
__naked __noinline __used
static long subprog_call_mem_kfunc(long a, long b, long c, long d, long e, long size)
{
        asm volatile (
                "r1 = *(u64 *)(r11 + 8);"       /* r1 = incoming arg6 (size) */
                "r2 = 0x0807060504030201 ll;"   /* r2 = buf contents */
                "*(u64 *)(r10 - 8) = r2;"       /* store buf to stack */
                "r2 = r10;"
                "r2 += -8;"                     /* r2 = &buf */
                "*(u64 *)(r11 - 8) = r2;"       /* outgoing arg6 = buf */
                "*(u64 *)(r11 - 16) = r1;"      /* outgoing arg7 = size */
                "r1 = 1;"
                "r2 = 2;"
                "r3 = 3;"
                "r4 = 4;"
                "r5 = 5;"
                "call %[bpf_kfunc_call_stack_arg_mem];"
                "exit;"
                :
                : __imm(bpf_kfunc_call_stack_arg_mem)
                : __clobber_all
        );
}

SEC("tc")
__description("stack_arg: precision backtracking across bpf2bpf call for kfunc")
__success
__log_level(2)
__flag(BPF_F_TEST_STATE_FREQ)
__btf_func_path("btf__stack_arg_precision.bpf.o")
__msg("mark_precise: frame1: last_idx 26 first_idx 13 subseq_idx -1")
__msg("mark_precise: frame1: regs= stack= before 25: (b7) r5 = 5")
__msg("mark_precise: frame1: regs= stack= before 24: (b7) r4 = 4")
__msg("mark_precise: frame1: regs= stack= before 23: (b7) r3 = 3")
__msg("mark_precise: frame1: regs= stack= before 22: (b7) r2 = 2")
__msg("mark_precise: frame1: regs= stack= before 21: (b7) r1 = 1")
__msg("mark_precise: frame1: regs= stack= before 20: (7b) *(u64 *)(r11 -16) = r1")
__msg("mark_precise: frame1: regs=r1 stack= before 19: (7b) *(u64 *)(r11 -8) = r2")
__msg("mark_precise: frame1: regs=r1 stack= before 18: (07) r2 += -8")
__msg("mark_precise: frame1: regs=r1 stack= before 17: (bf) r2 = r10")
__msg("mark_precise: frame1: regs=r1 stack= before 16: (7b) *(u64 *)(r10 -8) = r2")
__msg("mark_precise: frame1: regs=r1 stack= before 14: (18) r2 = 0x807060504030201")
__msg("mark_precise: frame1: regs=r1 stack= before 13: (79) r1 = *(u64 *)(r11 +8)")
__msg("mark_precise: frame1: parent state regs= stack=:  frame1: R10=fp0")
__msg("mark_precise: frame0: parent state regs= stack=:  R10=fp0")
__msg("mark_precise: frame1: last_idx 11 first_idx 11 subseq_idx 13")
__msg("mark_precise: frame1: regs= stack= before 11: (85) call pc+1")
__msg("mark_precise: frame0: parent state regs= stack=:  R1=1 R2=2 R3=3 R4=4 R5=5 R10=fp0")
__msg("mark_precise: frame0: last_idx 9 first_idx 7 subseq_idx 11")
__msg("mark_precise: frame0: regs= stack= before 9: (05) goto pc+1")
__msg("mark_precise: frame0: regs= stack= before 8: (7a) *(u64 *)(r11 -8) = 4")
__msg("mark_precise: frame1: last_idx 26 first_idx 13 subseq_idx -1 ")
__msg("mark_precise: frame1: regs= stack= before 25: (b7) r5 = 5")
__msg("mark_precise: frame1: regs= stack= before 24: (b7) r4 = 4")
__msg("mark_precise: frame1: regs= stack= before 23: (b7) r3 = 3")
__msg("mark_precise: frame1: regs= stack= before 22: (b7) r2 = 2")
__msg("mark_precise: frame1: regs= stack= before 21: (b7) r1 = 1")
__msg("mark_precise: frame1: regs= stack= before 20: (7b) *(u64 *)(r11 -16) = r1")
__msg("mark_precise: frame1: regs=r1 stack= before 19: (7b) *(u64 *)(r11 -8) = r2")
__msg("mark_precise: frame1: regs=r1 stack= before 18: (07) r2 += -8")
__msg("mark_precise: frame1: regs=r1 stack= before 17: (bf) r2 = r10")
__msg("mark_precise: frame1: regs=r1 stack= before 16: (7b) *(u64 *)(r10 -8) = r2")
__msg("mark_precise: frame1: regs=r1 stack= before 14: (18) r2 = 0x807060504030201")
__msg("mark_precise: frame1: regs=r1 stack= before 13: (79) r1 = *(u64 *)(r11 +8)")
__msg("mark_precise: frame1: parent state regs= stack=:  frame1: R10=fp0")
__msg("mark_precise: frame0: parent state regs= stack=:  R10=fp0")
__msg("mark_precise: frame1: last_idx 11 first_idx 11 subseq_idx 13 ")
__msg("mark_precise: frame1: regs= stack= before 11: (85) call pc+1")
__msg("mark_precise: frame0: parent state regs= stack=:  R1=1 R2=2 R3=3 R4=4 R5=5 R10=fp0")
__msg("mark_precise: frame0: last_idx 10 first_idx 10 subseq_idx 11 ")
__msg("mark_precise: frame0: regs= stack= before 10: (7a) *(u64 *)(r11 -8) = 6")
__naked void stack_arg_precision_bpf2bpf(void)
{
        asm volatile (
                "call %[bpf_get_prandom_u32];"
                "r6 = r0;"
                "r1 = 1;"
                "r2 = 2;"
                "r3 = 3;"
                "r4 = 4;"
                "r5 = 5;"
                "if r6 < 2 goto l0_%=;"
                "*(u64 *)(r11 - 8) = 4;"
                "goto l1_%=;"
        "l0_%=:"
                "*(u64 *)(r11 - 8) = 6;"
        "l1_%=:"
                "call subprog_call_mem_kfunc;"
                "exit;"
                :: __imm(bpf_get_prandom_u32)
                : __clobber_all
        );
}

#else

SEC("socket")
__description("stack_arg_precision: not supported, dummy test")
__success
int dummy_test(void)
{
        return 0;
}

#endif

char _license[] SEC("license") = "GPL";