// SPDX-License-Identifier: GPL-2.0
#include <vmlinux.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>
#include "bpf_misc.h"
#include "bpf_experimental.h"

extern int bpf_copy_from_user_str(void *dst, u32 dst__sz, const void *unsafe_ptr__ign, u64 flags) __weak __ksym;
extern void bpf_rcu_read_lock(void) __ksym;
extern void bpf_rcu_read_unlock(void) __ksym;

SEC("?tc")
__failure __msg("BPF_EXIT instruction in main prog cannot be used inside bpf_preempt_disable-ed region")
int preempt_lock_missing_1(struct __sk_buff *ctx)
{
	bpf_preempt_disable();
	return 0;
}

SEC("?tc")
__failure __msg("BPF_EXIT instruction in main prog cannot be used inside bpf_preempt_disable-ed region")
int preempt_lock_missing_2(struct __sk_buff *ctx)
{
	bpf_preempt_disable();
	bpf_preempt_disable();
	return 0;
}

SEC("?tc")
__failure __msg("BPF_EXIT instruction in main prog cannot be used inside bpf_preempt_disable-ed region")
int preempt_lock_missing_3(struct __sk_buff *ctx)
{
	bpf_preempt_disable();
	bpf_preempt_disable();
	bpf_preempt_disable();
	return 0;
}

SEC("?tc")
__failure __msg("BPF_EXIT instruction in main prog cannot be used inside bpf_preempt_disable-ed region")
int preempt_lock_missing_3_minus_2(struct __sk_buff *ctx)
{
	bpf_preempt_disable();
	bpf_preempt_disable();
	bpf_preempt_disable();
	bpf_preempt_enable();
	bpf_preempt_enable();
	return 0;
}

static __noinline void preempt_disable(void)
{
	bpf_preempt_disable();
}

static __noinline void preempt_enable(void)
{
	bpf_preempt_enable();
}

SEC("?tc")
__failure __msg("BPF_EXIT instruction in main prog cannot be used inside bpf_preempt_disable-ed region")
int preempt_lock_missing_1_subprog(struct __sk_buff *ctx)
{
	preempt_disable();
	return 0;
}

SEC("?tc")
__failure __msg("BPF_EXIT instruction in main prog cannot be used inside bpf_preempt_disable-ed region")
int preempt_lock_missing_2_subprog(struct __sk_buff *ctx)
{
	preempt_disable();
	preempt_disable();
	return 0;
}

SEC("?tc")
__failure __msg("BPF_EXIT instruction in main prog cannot be used inside bpf_preempt_disable-ed region")
int preempt_lock_missing_2_minus_1_subprog(struct __sk_buff *ctx)
{
	preempt_disable();
	preempt_disable();
	preempt_enable();
	return 0;
}

static __noinline void preempt_balance_subprog(void)
{
	preempt_disable();
	preempt_enable();
}

SEC("?tc")
__success int preempt_balance(struct __sk_buff *ctx)
{
	bpf_guard_preempt();
	return 0;
}

SEC("?tc")
__success int preempt_balance_subprog_test(struct __sk_buff *ctx)
{
	preempt_balance_subprog();
	return 0;
}

SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
__failure __msg("sleepable helper bpf_copy_from_user#")
int preempt_sleepable_helper(void *ctx)
{
	u32 data;

	bpf_preempt_disable();
	bpf_copy_from_user(&data, sizeof(data), NULL);
	bpf_preempt_enable();
	return 0;
}

SEC("?uprobe.s")
__failure __msg("sleepable helper bpf_get_stack#")
int preempt_sleepable_get_stack(struct pt_regs *ctx)
{
	struct bpf_stack_build_id stack;

	bpf_preempt_disable();
	bpf_get_stack(ctx, &stack, sizeof(stack),
		      BPF_F_USER_STACK | BPF_F_USER_BUILD_ID);
	bpf_preempt_enable();
	return 0;
}

SEC("?uprobe.s")
__failure __msg("sleepable helper bpf_get_task_stack#")
int preempt_sleepable_get_task_stack(void *ctx)
{
	struct bpf_stack_build_id stack;
	struct task_struct *task;

	task = bpf_get_current_task_btf();
	bpf_preempt_disable();
	bpf_get_task_stack(task, &stack, sizeof(stack),
			   BPF_F_USER_STACK | BPF_F_USER_BUILD_ID);
	bpf_preempt_enable();
	return 0;
}

SEC("?uprobe.s")
__success
int sleepable_get_stack(struct pt_regs *ctx)
{
	struct bpf_stack_build_id stack;

	bpf_get_stack(ctx, &stack, sizeof(stack),
		      BPF_F_USER_STACK | BPF_F_USER_BUILD_ID);
	return 0;
}

SEC("?uprobe.s")
__success
int sleepable_get_task_stack(void *ctx)
{
	struct bpf_stack_build_id stack;
	struct task_struct *task;

	task = bpf_get_current_task_btf();
	bpf_get_task_stack(task, &stack, sizeof(stack),
			   BPF_F_USER_STACK | BPF_F_USER_BUILD_ID);
	return 0;
}

SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
__failure __msg("kernel func bpf_copy_from_user_str is sleepable within non-preemptible region")
int preempt_sleepable_kfunc(void *ctx)
{
	u32 data;

	bpf_preempt_disable();
	bpf_copy_from_user_str(&data, sizeof(data), NULL, 0);
	bpf_preempt_enable();
	return 0;
}

SEC("?fentry/" SYS_PREFIX "sys_getpgid")
__failure __msg("program must be sleepable to call sleepable kfunc bpf_copy_from_user_str")
__msg("cannot be used in non-sleepable program")
int non_sleepable_kfunc(void *ctx)
{
	u32 data;

	bpf_copy_from_user_str(&data, sizeof(data), NULL, 0);
	return 0;
}

SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
__failure __msg("kernel func bpf_copy_from_user_str is sleepable within rcu_read_lock region")
__msg("cannot be used in RCU read lock region")
int sleepable_kfunc_in_rcu(void *ctx)
{
	u32 data;

	bpf_rcu_read_lock();
	bpf_copy_from_user_str(&data, sizeof(data), NULL, 0);
	bpf_rcu_read_unlock();
	return 0;
}

int __noinline preempt_global_subprog(void)
{
	preempt_balance_subprog();
	return 0;
}

SEC("?tc")
__success
int preempt_global_subprog_test(struct __sk_buff *ctx)
{
	preempt_disable();
	preempt_global_subprog();
	preempt_enable();
	return 0;
}

int __noinline
global_subprog(int i)
{
	if (i)
		bpf_printk("%p", &i);
	return i;
}

int __noinline
global_sleepable_helper_subprog(int i)
{
	if (i)
		bpf_copy_from_user(&i, sizeof(i), NULL);
	return i;
}

int __noinline
global_sleepable_kfunc_subprog(int i)
{
	if (i)
		bpf_copy_from_user_str(&i, sizeof(i), NULL, 0);
	global_subprog(i);
	return i;
}

int __noinline
global_subprog_calling_sleepable_global(int i)
{
	if (!i)
		global_sleepable_kfunc_subprog(i);
	return i;
}

SEC("?syscall")
__failure __msg("sleepable global function")
int preempt_global_sleepable_helper_subprog(struct __sk_buff *ctx)
{
	preempt_disable();
	if (ctx->mark)
		global_sleepable_helper_subprog(ctx->mark);
	preempt_enable();
	return 0;
}

SEC("?syscall")
__failure __msg("sleepable global function")
int preempt_global_sleepable_kfunc_subprog(struct __sk_buff *ctx)
{
	preempt_disable();
	if (ctx->mark)
		global_sleepable_kfunc_subprog(ctx->mark);
	preempt_enable();
	return 0;
}

SEC("?syscall")
__failure __msg("sleepable global function")
int preempt_global_sleepable_subprog_indirect(struct __sk_buff *ctx)
{
	preempt_disable();
	if (ctx->mark)
		global_subprog_calling_sleepable_global(ctx->mark);
	preempt_enable();
	return 0;
}

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