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

#include <vmlinux.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>
#include "bpf_misc.h"
#include "bpf_experimental.h"

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

struct task_struct *bpf_task_acquire(struct task_struct *p) __ksym;
void bpf_task_release(struct task_struct *p) __ksym;
void bpf_rcu_read_lock(void) __ksym;
void bpf_rcu_read_unlock(void) __ksym;

/* Timer tests */

struct timer_elem {
	struct bpf_timer t;
};

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

static int timer_cb(void *map, int *key, struct bpf_timer *timer)
{
	u32 data;
	/* Timer callbacks are never sleepable, even from non-sleepable programs */
	bpf_copy_from_user(&data, sizeof(data), NULL);
	return 0;
}

SEC("fentry/bpf_fentry_test1")
__failure __msg("sleepable helper bpf_copy_from_user#{{[0-9]+}} in non-sleepable prog")
int timer_non_sleepable_prog(void *ctx)
{
	struct timer_elem *val;
	int key = 0;

	val = bpf_map_lookup_elem(&timer_map, &key);
	if (!val)
		return 0;

	bpf_timer_init(&val->t, &timer_map, 0);
	bpf_timer_set_callback(&val->t, timer_cb);
	return 0;
}

SEC("lsm.s/file_open")
__failure __msg("sleepable helper bpf_copy_from_user#{{[0-9]+}} in non-sleepable prog")
int timer_sleepable_prog(void *ctx)
{
	struct timer_elem *val;
	int key = 0;

	val = bpf_map_lookup_elem(&timer_map, &key);
	if (!val)
		return 0;

	bpf_timer_init(&val->t, &timer_map, 0);
	bpf_timer_set_callback(&val->t, timer_cb);
	return 0;
}

static int timer_sys_bpf_cb(void *map, int *key, struct bpf_timer *timer)
{
	__u64 attr = 0;

	bpf_sys_bpf(BPF_MAP_FREEZE, &attr, sizeof(attr));
	return 0;
}

SEC("syscall")
__failure __msg("sleepable helper bpf_sys_bpf#{{[0-9]+}} in non-sleepable prog")
int timer_sys_bpf_prog(void *ctx)
{
	struct timer_elem *val;
	int key = 0;

	val = bpf_map_lookup_elem(&timer_map, &key);
	if (!val)
		return 0;

	bpf_timer_init(&val->t, &timer_map, 0);
	bpf_timer_set_callback(&val->t, timer_sys_bpf_cb);
	return 0;
}

static int timer_sys_close_cb(void *map, int *key, struct bpf_timer *timer)
{
	bpf_sys_close(0);
	return 0;
}

SEC("syscall")
__failure __msg("sleepable helper bpf_sys_close#{{[0-9]+}} in non-sleepable prog")
int timer_sys_close_prog(void *ctx)
{
	struct timer_elem *val;
	int key = 0;

	val = bpf_map_lookup_elem(&timer_map, &key);
	if (!val)
		return 0;

	bpf_timer_init(&val->t, &timer_map, 0);
	bpf_timer_set_callback(&val->t, timer_sys_close_cb);
	return 0;
}

static int timer_btf_find_cb(void *map, int *key, struct bpf_timer *timer)
{
	char name[] = "task_struct";

	bpf_btf_find_by_name_kind(name, sizeof(name), BTF_KIND_STRUCT, 0);
	return 0;
}

SEC("syscall")
__failure __msg("sleepable helper bpf_btf_find_by_name_kind#{{[0-9]+}} in non-sleepable prog")
int timer_btf_find_prog(void *ctx)
{
	struct timer_elem *val;
	int key = 0;

	val = bpf_map_lookup_elem(&timer_map, &key);
	if (!val)
		return 0;

	bpf_timer_init(&val->t, &timer_map, 0);
	bpf_timer_set_callback(&val->t, timer_btf_find_cb);
	return 0;
}

SEC("syscall")
__success
int syscall_sys_bpf_prog(void *ctx)
{
	__u64 attr = 0;

	bpf_sys_bpf(BPF_MAP_FREEZE, &attr, sizeof(attr));
	return 0;
}

SEC("syscall")
__success
int syscall_sys_close_prog(void *ctx)
{
	bpf_sys_close(0);
	return 0;
}

SEC("syscall")
__success
int syscall_btf_find_prog(void *ctx)
{
	char name[] = "task_struct";

	bpf_btf_find_by_name_kind(name, sizeof(name), BTF_KIND_STRUCT, 0);
	return 0;
}

/* Workqueue tests */

struct wq_elem {
	struct bpf_wq w;
	struct task_struct __kptr *task;
};

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

static int wq_cb(void *map, int *key, void *value)
{
	u32 data;
	/* Workqueue callbacks are always sleepable, even from non-sleepable programs */
	bpf_copy_from_user(&data, sizeof(data), NULL);
	return 0;
}

SEC("fentry/bpf_fentry_test1")
__success
int wq_non_sleepable_prog(void *ctx)
{
	struct wq_elem *val;
	int key = 0;

	val = bpf_map_lookup_elem(&wq_map, &key);
	if (!val)
		return 0;

	if (bpf_wq_init(&val->w, &wq_map, 0) != 0)
		return 0;
	if (bpf_wq_set_callback(&val->w, wq_cb, 0) != 0)
		return 0;
	return 0;
}

SEC("lsm.s/file_open")
__success
int wq_sleepable_prog(void *ctx)
{
	struct wq_elem *val;
	int key = 0;

	val = bpf_map_lookup_elem(&wq_map, &key);
	if (!val)
		return 0;

	if (bpf_wq_init(&val->w, &wq_map, 0) != 0)
		return 0;
	if (bpf_wq_set_callback(&val->w, wq_cb, 0) != 0)
		return 0;
	return 0;
}

__noinline int wq_global_acquire(void)
{
	struct task_struct *task, *acquired;
	struct wq_elem *val;
	int key = 0;

	val = bpf_map_lookup_elem(&wq_map, &key);
	if (!val)
		return 0;

	task = val->task;
	if (!task)
		return 0;

	acquired = bpf_task_acquire(task);
	if (acquired)
		bpf_task_release(acquired);
	return 0;
}

static int wq_global_rcu_cb(void *map, int *key, void *value)
{
	wq_global_acquire();
	return 0;
}

SEC("fentry/bpf_fentry_test1")
__failure __msg("R1 must be a rcu pointer")
int wq_global_rcu_prog(void *ctx)
{
	struct wq_elem *val;
	int key = 0;

	val = bpf_map_lookup_elem(&wq_map, &key);
	if (!val)
		return 0;

	bpf_wq_init(&val->w, &wq_map, 0);
	bpf_wq_set_callback(&val->w, wq_global_rcu_cb, 0);
	return 0;
}

static int wq_global_rcu_lock_cb(void *map, int *key, void *value)
{
	bpf_rcu_read_lock();
	wq_global_acquire();
	bpf_rcu_read_unlock();
	return 0;
}

SEC("fentry/bpf_fentry_test1")
__success
int wq_global_rcu_lock_prog(void *ctx)
{
	struct wq_elem *val;
	int key = 0;

	/* Verify the same global subprog in non-sleepable and protected contexts. */
	wq_global_acquire();

	val = bpf_map_lookup_elem(&wq_map, &key);
	if (!val)
		return 0;

	bpf_wq_init(&val->w, &wq_map, 0);
	bpf_wq_set_callback(&val->w, wq_global_rcu_lock_cb, 0);
	return 0;
}

__weak __noinline int wq_global_no_rcu(void)
{
	return 0;
}

static int wq_global_no_rcu_cb(void *map, int *key, void *value)
{
	wq_global_no_rcu();
	return 0;
}

SEC("fentry/bpf_fentry_test1")
__success __log_level(4)
__msg("subprog {{[0-9]+}} (wq_global_no_rcu) global insns_self 4 insns_total 4 stack 0")
int wq_global_no_rcu_prog(void *ctx)
{
	struct wq_elem *val;
	int key = 0;

	/* Verify the same global in non-sleepable and unprotected contexts. */
	wq_global_no_rcu();

	val = bpf_map_lookup_elem(&wq_map, &key);
	if (!val)
		return 0;

	bpf_wq_init(&val->w, &wq_map, 0);
	bpf_wq_set_callback(&val->w, wq_global_no_rcu_cb, 0);
	return 0;
}

/* Task work tests */

struct task_work_elem {
	struct bpf_task_work tw;
};

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

static int task_work_cb(struct bpf_map *map, void *key, void *value)
{
	u32 data;
	/* Task work callbacks are always sleepable, even from non-sleepable programs */
	bpf_copy_from_user(&data, sizeof(data), NULL);
	return 0;
}

SEC("fentry/bpf_fentry_test1")
__success
int task_work_non_sleepable_prog(void *ctx)
{
	struct task_work_elem *val;
	struct task_struct *task;
	int key = 0;

	val = bpf_map_lookup_elem(&task_work_map, &key);
	if (!val)
		return 0;

	task = bpf_get_current_task_btf();
	if (!task)
		return 0;

	bpf_task_work_schedule_resume(task, &val->tw, &task_work_map, task_work_cb);
	return 0;
}

SEC("lsm.s/file_open")
__success
int task_work_sleepable_prog(void *ctx)
{
	struct task_work_elem *val;
	struct task_struct *task;
	int key = 0;

	val = bpf_map_lookup_elem(&task_work_map, &key);
	if (!val)
		return 0;

	task = bpf_get_current_task_btf();
	if (!task)
		return 0;

	bpf_task_work_schedule_resume(task, &val->tw, &task_work_map, task_work_cb);
	return 0;
}