// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2025 Meta Platforms, Inc. and affiliates. */
#include <vmlinux.h>
#include <bpf/bpf_tracing.h>
#include <bpf/bpf_helpers.h>
#include "bpf_misc.h"
#include "bpf_experimental.h"
#include <bpf_arena_common.h>

struct arr_elem {
	struct bpf_res_spin_lock lock;
};

struct {
	__uint(type, BPF_MAP_TYPE_ARRAY);
	__uint(max_entries, 1);
	__type(key, int);
	__type(value, struct arr_elem);
} arrmap SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_ARENA);
	__uint(map_flags, BPF_F_MMAPABLE);
	__uint(max_entries, 1); /* number of pages */
} arena SEC(".maps");

struct elem {
	struct bpf_timer timer;
};

struct {
	__uint(type, BPF_MAP_TYPE_ARRAY);
	__uint(max_entries, 1);
	__type(key, int);
	__type(value, struct elem);
} array SEC(".maps");

#define ENOSPC 28
#define _STR "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"

int size;
u64 fault_addr;
void *arena_ptr;

#define private(name) SEC(".bss." #name) __hidden __attribute__((aligned(8)))

private(STREAM) struct bpf_spin_lock block;

SEC("syscall")
__success __retval(0)
int stream_exhaust(void *ctx)
{
	/* Use global variable for loop convergence. */
	size = 0;
	bpf_repeat(BPF_MAX_LOOPS) {
		if (bpf_stream_printk(BPF_STDOUT, _STR) == -ENOSPC && size == 99954)
			return 0;
		size += sizeof(_STR) - 1;
	}
	return 1;
}

SEC("syscall")
__arch_x86_64
__arch_arm64
__arch_s390x
__arch_riscv64
__arch_loongarch
__success __retval(0)
__stderr("ERROR: Timeout detected for may_goto instruction")
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
int stream_cond_break(void *ctx)
{
	while (can_loop)
		;
	return 0;
}

SEC("syscall")
__success __retval(0)
__stderr("ERROR: AA or ABBA deadlock detected for bpf_res_spin_lock")
__stderr("{{Attempted lock   = (0x[0-9a-fA-F]+)\n"
"Total held locks = 1\n"
"Held lock\\[ 0\\] = \\1}}")
__stderr("...")
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
int stream_deadlock(void *ctx)
{
	struct bpf_res_spin_lock *lock, *nlock;

	lock = bpf_map_lookup_elem(&arrmap, &(int){0});
	if (!lock)
		return 1;
	nlock = bpf_map_lookup_elem(&arrmap, &(int){0});
	if (!nlock)
		return 1;
	if (bpf_res_spin_lock(lock))
		return 1;
	if (bpf_res_spin_lock(nlock)) {
		bpf_res_spin_unlock(lock);
		return 0;
	}
	bpf_res_spin_unlock(nlock);
	bpf_res_spin_unlock(lock);
	return 1;
}

SEC("syscall")
__success __retval(0)
int stream_syscall(void *ctx)
{
	bpf_stream_printk(BPF_STDOUT, "foo");
	return 0;
}

SEC("syscall")
__arch_x86_64
__arch_arm64
__success __retval(0)
__stderr("ERROR: Arena WRITE access at unmapped address 0x{{.*}}")
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
int stream_arena_write_fault(void *ctx)
{
	struct bpf_arena *ptr = (void *)&arena;
	u64 user_vm_start;

	/* Prevent GCC bounds warning: casting &arena to struct bpf_arena *
	 * triggers bounds checking since the map definition is smaller than struct
	 * bpf_arena. barrier_var() makes the pointer opaque to GCC, preventing the
	 * bounds analysis
	 */
	barrier_var(ptr);
	user_vm_start = ptr->user_vm_start;
	fault_addr = user_vm_start + 0x7fff;
	bpf_addr_space_cast(user_vm_start, 0, 1);
	asm volatile (
		"r1 = %0;"
		"r2 = 1;"
		"*(u32 *)(r1 + 0x7fff) = r2;"
		:
		: "r" (user_vm_start)
		: "r1", "r2"
	);
	return 0;
}

SEC("syscall")
__arch_x86_64
__arch_arm64
__success __retval(0)
__stderr("ERROR: Arena READ access at unmapped address 0x{{.*}}")
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
int stream_arena_read_fault(void *ctx)
{
	struct bpf_arena *ptr = (void *)&arena;
	u64 user_vm_start;

	/* Prevent GCC bounds warning: casting &arena to struct bpf_arena *
	 * triggers bounds checking since the map definition is smaller than struct
	 * bpf_arena. barrier_var() makes the pointer opaque to GCC, preventing the
	 * bounds analysis
	 */
	barrier_var(ptr);
	user_vm_start = ptr->user_vm_start;
	fault_addr = user_vm_start + 0x7fff;
	bpf_addr_space_cast(user_vm_start, 0, 1);
	asm volatile (
		"r1 = %0;"
		"r1 = *(u32 *)(r1 + 0x7fff);"
		:
		: "r" (user_vm_start)
		: "r1"
	);
	return 0;
}

SEC("syscall")
__arch_x86_64
__arch_arm64
__success __retval(0)
__stderr("ERROR: Arena READ access at unmapped address 0x{{.*}}")
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
int stream_arena_load_acquire_fault(void *ctx)
{
	static const struct bpf_insn load_acquire_insn = {
		.code	 = 0xc3,	/* BPF_STX | BPF_ATOMIC | BPF_W */
		.dst_reg = 0,		/* BPF_REG_0 */
		.src_reg = 1,		/* BPF_REG_1 */
		.off	 = 0x7fff,
		.imm	 = 0x100,	/* BPF_LOAD_ACQ */
	};
	struct bpf_arena *ptr = (void *)&arena;
	u64 user_vm_start, val;

	/*
	 * Prevent GCC bounds warning: casting &arena to struct bpf_arena *
	 * triggers bounds checking since the map definition is smaller than
	 * struct bpf_arena. barrier_var() makes the pointer opaque to GCC,
	 * preventing the bounds analysis.
	 */
	barrier_var(ptr);
	user_vm_start = ptr->user_vm_start;
	fault_addr = user_vm_start + 0x7fff;
	bpf_addr_space_cast(user_vm_start, 0, 1);
	asm volatile (
		"r1 = %[user_vm_start];"
		"r0 = 1;"
		".8byte %[load_acquire_insn];" /* r0 = load_acquire((u32 *)(r1 + 0x7fff)) */
		"%[val] = r0;"
		: [val] "=r" (val)
		: [user_vm_start] "r" (user_vm_start),
		  __imm_insn(load_acquire_insn, load_acquire_insn)
		: "r0", "r1"
	);
	return val;
}

SEC("syscall")
__arch_x86_64
__arch_arm64
__success __retval(0)
__stderr("ERROR: Arena WRITE access at unmapped address 0x{{.*}}")
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
int stream_arena_xchg_fault(void *ctx)
{
	static const struct bpf_insn xchg_insn = {
		.code	 = 0xc3,	/* BPF_STX | BPF_ATOMIC | BPF_W */
		.dst_reg = 1,		/* BPF_REG_1 */
		.src_reg = 2,		/* BPF_REG_2 */
		.off	 = 0x7fff,
		.imm	 = 0xe1,	/* BPF_XCHG */
	};
	struct bpf_arena *ptr = (void *)&arena;
	u64 user_vm_start, val;

	/*
	 * Prevent GCC bounds warning: casting &arena to struct bpf_arena *
	 * triggers bounds checking since the map definition is smaller than
	 * struct bpf_arena. barrier_var() makes the pointer opaque to GCC,
	 * preventing the bounds analysis.
	 */
	barrier_var(ptr);
	user_vm_start = ptr->user_vm_start;
	fault_addr = user_vm_start + 0x7fff;
	bpf_addr_space_cast(user_vm_start, 0, 1);
	/*
	 * A read-modify-write carrying BPF_FETCH writes to memory, so the fault
	 * has to be reported as a WRITE from the dst_reg address, but it also
	 * reads the old value into src_reg, so the exception handler has to
	 * clear src_reg. Poison it up front, the returned value must be 0.
	 */
	asm volatile (
		"r1 = %[user_vm_start];"
		"r2 = 1;"
		".8byte %[xchg_insn];" /* r2 = xchg((u32 *)(r1 + 0x7fff), r2) */
		"%[val] = r2;"
		: [val] "=r" (val)
		: [user_vm_start] "r" (user_vm_start),
		  __imm_insn(xchg_insn, xchg_insn)
		: "r1", "r2"
	);
	return val;
}

SEC("syscall")
__arch_x86_64
__arch_arm64
__success __retval(0)
__stderr("ERROR: Arena WRITE access at unmapped address 0x{{.*}}")
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
int stream_arena_cmpxchg_fault(void *ctx)
{
	static const struct bpf_insn cmpxchg_insn = {
		.code	 = 0xc3,	/* BPF_STX | BPF_ATOMIC | BPF_W */
		.dst_reg = 1,		/* BPF_REG_1 */
		.src_reg = 2,		/* BPF_REG_2 */
		.off	 = 0x7fff,
		.imm	 = 0xf1,	/* BPF_CMPXCHG */
	};
	struct bpf_arena *ptr = (void *)&arena;
	u64 user_vm_start, val;

	/*
	 * Prevent GCC bounds warning: casting &arena to struct bpf_arena *
	 * triggers bounds checking since the map definition is smaller than
	 * struct bpf_arena. barrier_var() makes the pointer opaque to GCC,
	 * preventing the bounds analysis.
	 */
	barrier_var(ptr);
	user_vm_start = ptr->user_vm_start;
	fault_addr = user_vm_start + 0x7fff;
	bpf_addr_space_cast(user_vm_start, 0, 1);
	/*
	 * Same as the exchange above, except that a BPF_CMPXCHG reads the old
	 * value into r0 rather than into src_reg, so r0 is the register the
	 * exception handler has to clear. It doubles as the compare value, but
	 * the comparison never happens since the access faults first.
	 */
	asm volatile (
		"r1 = %[user_vm_start];"
		"r0 = 1;"
		"r2 = 2;"
		".8byte %[cmpxchg_insn];" /* r0 = cmpxchg((u32 *)(r1 + 0x7fff), r0, r2) */
		"%[val] = r0;"
		: [val] "=r" (val)
		: [user_vm_start] "r" (user_vm_start),
		  __imm_insn(cmpxchg_insn, cmpxchg_insn)
		: "r0", "r1", "r2"
	);
	return val;
}

static __noinline void subprog(void)
{
	int __arena *addr = (int __arena *)0xdeadbeef;

	arena_ptr = &arena;
	*addr = 1;
}

SEC("syscall")
__arch_x86_64
__arch_arm64
__success __retval(0)
__stderr("ERROR: Arena WRITE access at unmapped address 0x{{.*}}")
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
int stream_arena_subprog_fault(void *ctx)
{
	subprog();
	return 0;
}

static __noinline int timer_cb(void *map, int *key, struct bpf_timer *timer)
{
	int __arena *addr = (int __arena *)0xdeadbeef;

	arena_ptr = &arena;
	*addr = 1;
	return 0;
}

SEC("syscall")
__arch_x86_64
__arch_arm64
__success __retval(0)
__stderr("ERROR: Arena WRITE access at unmapped address 0x{{.*}}")
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
int stream_arena_callback_fault(void *ctx)
{
	struct bpf_timer *arr_timer;

	arr_timer = bpf_map_lookup_elem(&array, &(int){0});
	if (!arr_timer)
		return 0;
	bpf_timer_init(arr_timer, &array, 1);
	bpf_timer_set_callback(arr_timer, timer_cb);
	bpf_timer_start(arr_timer, 0, 0);
	return 0;
}

SEC("syscall")
__arch_x86_64
__arch_arm64
__success __retval(0)
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
int stream_print_stack_kfunc(void *ctx)
{
	return bpf_stream_print_stack(BPF_STDERR);
}

SEC("syscall")
__success __retval(-2)
int stream_print_stack_invalid_id(void *ctx)
{
	/* Try to pass an invalid stream ID. */
	return bpf_stream_print_stack((enum bpf_stream_id)0xbadcafe);
}

SEC("syscall")
__arch_x86_64
__arch_arm64
__success __retval(0)
__stdout(_STR)
__stderr("CPU: {{[0-9]+}} UID: 0 PID: {{[0-9]+}} Comm: {{.*}}")
__stderr("Call trace:\n"
"{{([a-zA-Z_][a-zA-Z0-9_]*\\+0x[0-9a-fA-F]+/0x[0-9a-fA-F]+\n"
"|[ \t]+[^\n]+\n)*}}")
int stream_print_kfuncs_locked(void *ctx)
{
	int ret;

	bpf_spin_lock(&block);

	ret = bpf_stream_printk(BPF_STDOUT, _STR);
	if (ret)
		goto out;

	ret = bpf_stream_print_stack(BPF_STDERR);

out:
	bpf_spin_unlock(&block);

	return ret;
}


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