162306a36Sopenharmony_ci// SPDX-License-Identifier: GPL-2.0
262306a36Sopenharmony_ci/* Copyright (c) 2023 Meta Platforms, Inc. and affiliates. */
362306a36Sopenharmony_ci
462306a36Sopenharmony_ci#include <vmlinux.h>
562306a36Sopenharmony_ci#include <bpf/bpf_tracing.h>
662306a36Sopenharmony_ci#include <bpf/bpf_helpers.h>
762306a36Sopenharmony_ci#include "bpf_misc.h"
862306a36Sopenharmony_ci
962306a36Sopenharmony_ci#include "cpumask_common.h"
1062306a36Sopenharmony_ci
1162306a36Sopenharmony_cichar _license[] SEC("license") = "GPL";
1262306a36Sopenharmony_ci
1362306a36Sopenharmony_ci/* Prototype for all of the program trace events below:
1462306a36Sopenharmony_ci *
1562306a36Sopenharmony_ci * TRACE_EVENT(task_newtask,
1662306a36Sopenharmony_ci *         TP_PROTO(struct task_struct *p, u64 clone_flags)
1762306a36Sopenharmony_ci */
1862306a36Sopenharmony_ci
1962306a36Sopenharmony_ciSEC("tp_btf/task_newtask")
2062306a36Sopenharmony_ci__failure __msg("Unreleased reference")
2162306a36Sopenharmony_ciint BPF_PROG(test_alloc_no_release, struct task_struct *task, u64 clone_flags)
2262306a36Sopenharmony_ci{
2362306a36Sopenharmony_ci	struct bpf_cpumask *cpumask;
2462306a36Sopenharmony_ci
2562306a36Sopenharmony_ci	cpumask = create_cpumask();
2662306a36Sopenharmony_ci	__sink(cpumask);
2762306a36Sopenharmony_ci
2862306a36Sopenharmony_ci	/* cpumask is never released. */
2962306a36Sopenharmony_ci	return 0;
3062306a36Sopenharmony_ci}
3162306a36Sopenharmony_ci
3262306a36Sopenharmony_ciSEC("tp_btf/task_newtask")
3362306a36Sopenharmony_ci__failure __msg("NULL pointer passed to trusted arg0")
3462306a36Sopenharmony_ciint BPF_PROG(test_alloc_double_release, struct task_struct *task, u64 clone_flags)
3562306a36Sopenharmony_ci{
3662306a36Sopenharmony_ci	struct bpf_cpumask *cpumask;
3762306a36Sopenharmony_ci
3862306a36Sopenharmony_ci	cpumask = create_cpumask();
3962306a36Sopenharmony_ci
4062306a36Sopenharmony_ci	/* cpumask is released twice. */
4162306a36Sopenharmony_ci	bpf_cpumask_release(cpumask);
4262306a36Sopenharmony_ci	bpf_cpumask_release(cpumask);
4362306a36Sopenharmony_ci
4462306a36Sopenharmony_ci	return 0;
4562306a36Sopenharmony_ci}
4662306a36Sopenharmony_ci
4762306a36Sopenharmony_ciSEC("tp_btf/task_newtask")
4862306a36Sopenharmony_ci__failure __msg("must be referenced")
4962306a36Sopenharmony_ciint BPF_PROG(test_acquire_wrong_cpumask, struct task_struct *task, u64 clone_flags)
5062306a36Sopenharmony_ci{
5162306a36Sopenharmony_ci	struct bpf_cpumask *cpumask;
5262306a36Sopenharmony_ci
5362306a36Sopenharmony_ci	/* Can't acquire a non-struct bpf_cpumask. */
5462306a36Sopenharmony_ci	cpumask = bpf_cpumask_acquire((struct bpf_cpumask *)task->cpus_ptr);
5562306a36Sopenharmony_ci	__sink(cpumask);
5662306a36Sopenharmony_ci
5762306a36Sopenharmony_ci	return 0;
5862306a36Sopenharmony_ci}
5962306a36Sopenharmony_ci
6062306a36Sopenharmony_ciSEC("tp_btf/task_newtask")
6162306a36Sopenharmony_ci__failure __msg("bpf_cpumask_set_cpu args#1 expected pointer to STRUCT bpf_cpumask")
6262306a36Sopenharmony_ciint BPF_PROG(test_mutate_cpumask, struct task_struct *task, u64 clone_flags)
6362306a36Sopenharmony_ci{
6462306a36Sopenharmony_ci	struct bpf_cpumask *cpumask;
6562306a36Sopenharmony_ci
6662306a36Sopenharmony_ci	/* Can't set the CPU of a non-struct bpf_cpumask. */
6762306a36Sopenharmony_ci	bpf_cpumask_set_cpu(0, (struct bpf_cpumask *)task->cpus_ptr);
6862306a36Sopenharmony_ci	__sink(cpumask);
6962306a36Sopenharmony_ci
7062306a36Sopenharmony_ci	return 0;
7162306a36Sopenharmony_ci}
7262306a36Sopenharmony_ci
7362306a36Sopenharmony_ciSEC("tp_btf/task_newtask")
7462306a36Sopenharmony_ci__failure __msg("Unreleased reference")
7562306a36Sopenharmony_ciint BPF_PROG(test_insert_remove_no_release, struct task_struct *task, u64 clone_flags)
7662306a36Sopenharmony_ci{
7762306a36Sopenharmony_ci	struct bpf_cpumask *cpumask;
7862306a36Sopenharmony_ci	struct __cpumask_map_value *v;
7962306a36Sopenharmony_ci
8062306a36Sopenharmony_ci	cpumask = create_cpumask();
8162306a36Sopenharmony_ci	if (!cpumask)
8262306a36Sopenharmony_ci		return 0;
8362306a36Sopenharmony_ci
8462306a36Sopenharmony_ci	if (cpumask_map_insert(cpumask))
8562306a36Sopenharmony_ci		return 0;
8662306a36Sopenharmony_ci
8762306a36Sopenharmony_ci	v = cpumask_map_value_lookup();
8862306a36Sopenharmony_ci	if (!v)
8962306a36Sopenharmony_ci		return 0;
9062306a36Sopenharmony_ci
9162306a36Sopenharmony_ci	cpumask = bpf_kptr_xchg(&v->cpumask, NULL);
9262306a36Sopenharmony_ci
9362306a36Sopenharmony_ci	/* cpumask is never released. */
9462306a36Sopenharmony_ci	return 0;
9562306a36Sopenharmony_ci}
9662306a36Sopenharmony_ci
9762306a36Sopenharmony_ciSEC("tp_btf/task_newtask")
9862306a36Sopenharmony_ci__failure __msg("NULL pointer passed to trusted arg0")
9962306a36Sopenharmony_ciint BPF_PROG(test_cpumask_null, struct task_struct *task, u64 clone_flags)
10062306a36Sopenharmony_ci{
10162306a36Sopenharmony_ci  /* NULL passed to KF_TRUSTED_ARGS kfunc. */
10262306a36Sopenharmony_ci	bpf_cpumask_empty(NULL);
10362306a36Sopenharmony_ci
10462306a36Sopenharmony_ci	return 0;
10562306a36Sopenharmony_ci}
10662306a36Sopenharmony_ci
10762306a36Sopenharmony_ciSEC("tp_btf/task_newtask")
10862306a36Sopenharmony_ci__failure __msg("R2 must be a rcu pointer")
10962306a36Sopenharmony_ciint BPF_PROG(test_global_mask_out_of_rcu, struct task_struct *task, u64 clone_flags)
11062306a36Sopenharmony_ci{
11162306a36Sopenharmony_ci	struct bpf_cpumask *local, *prev;
11262306a36Sopenharmony_ci
11362306a36Sopenharmony_ci	local = create_cpumask();
11462306a36Sopenharmony_ci	if (!local)
11562306a36Sopenharmony_ci		return 0;
11662306a36Sopenharmony_ci
11762306a36Sopenharmony_ci	prev = bpf_kptr_xchg(&global_mask, local);
11862306a36Sopenharmony_ci	if (prev) {
11962306a36Sopenharmony_ci		bpf_cpumask_release(prev);
12062306a36Sopenharmony_ci		err = 3;
12162306a36Sopenharmony_ci		return 0;
12262306a36Sopenharmony_ci	}
12362306a36Sopenharmony_ci
12462306a36Sopenharmony_ci	bpf_rcu_read_lock();
12562306a36Sopenharmony_ci	local = global_mask;
12662306a36Sopenharmony_ci	if (!local) {
12762306a36Sopenharmony_ci		err = 4;
12862306a36Sopenharmony_ci		bpf_rcu_read_unlock();
12962306a36Sopenharmony_ci		return 0;
13062306a36Sopenharmony_ci	}
13162306a36Sopenharmony_ci
13262306a36Sopenharmony_ci	bpf_rcu_read_unlock();
13362306a36Sopenharmony_ci
13462306a36Sopenharmony_ci	/* RCU region is exited before calling KF_RCU kfunc. */
13562306a36Sopenharmony_ci
13662306a36Sopenharmony_ci	bpf_cpumask_test_cpu(0, (const struct cpumask *)local);
13762306a36Sopenharmony_ci
13862306a36Sopenharmony_ci	return 0;
13962306a36Sopenharmony_ci}
14062306a36Sopenharmony_ci
14162306a36Sopenharmony_ciSEC("tp_btf/task_newtask")
14262306a36Sopenharmony_ci__failure __msg("NULL pointer passed to trusted arg1")
14362306a36Sopenharmony_ciint BPF_PROG(test_global_mask_no_null_check, struct task_struct *task, u64 clone_flags)
14462306a36Sopenharmony_ci{
14562306a36Sopenharmony_ci	struct bpf_cpumask *local, *prev;
14662306a36Sopenharmony_ci
14762306a36Sopenharmony_ci	local = create_cpumask();
14862306a36Sopenharmony_ci	if (!local)
14962306a36Sopenharmony_ci		return 0;
15062306a36Sopenharmony_ci
15162306a36Sopenharmony_ci	prev = bpf_kptr_xchg(&global_mask, local);
15262306a36Sopenharmony_ci	if (prev) {
15362306a36Sopenharmony_ci		bpf_cpumask_release(prev);
15462306a36Sopenharmony_ci		err = 3;
15562306a36Sopenharmony_ci		return 0;
15662306a36Sopenharmony_ci	}
15762306a36Sopenharmony_ci
15862306a36Sopenharmony_ci	bpf_rcu_read_lock();
15962306a36Sopenharmony_ci	local = global_mask;
16062306a36Sopenharmony_ci
16162306a36Sopenharmony_ci	/* No NULL check is performed on global cpumask kptr. */
16262306a36Sopenharmony_ci	bpf_cpumask_test_cpu(0, (const struct cpumask *)local);
16362306a36Sopenharmony_ci
16462306a36Sopenharmony_ci	bpf_rcu_read_unlock();
16562306a36Sopenharmony_ci
16662306a36Sopenharmony_ci	return 0;
16762306a36Sopenharmony_ci}
16862306a36Sopenharmony_ci
16962306a36Sopenharmony_ciSEC("tp_btf/task_newtask")
17062306a36Sopenharmony_ci__failure __msg("Possibly NULL pointer passed to helper arg2")
17162306a36Sopenharmony_ciint BPF_PROG(test_global_mask_rcu_no_null_check, struct task_struct *task, u64 clone_flags)
17262306a36Sopenharmony_ci{
17362306a36Sopenharmony_ci	struct bpf_cpumask *prev, *curr;
17462306a36Sopenharmony_ci
17562306a36Sopenharmony_ci	curr = bpf_cpumask_create();
17662306a36Sopenharmony_ci	if (!curr)
17762306a36Sopenharmony_ci		return 0;
17862306a36Sopenharmony_ci
17962306a36Sopenharmony_ci	prev = bpf_kptr_xchg(&global_mask, curr);
18062306a36Sopenharmony_ci	if (prev)
18162306a36Sopenharmony_ci		bpf_cpumask_release(prev);
18262306a36Sopenharmony_ci
18362306a36Sopenharmony_ci	bpf_rcu_read_lock();
18462306a36Sopenharmony_ci	curr = global_mask;
18562306a36Sopenharmony_ci	/* PTR_TO_BTF_ID | PTR_MAYBE_NULL | MEM_RCU passed to bpf_kptr_xchg() */
18662306a36Sopenharmony_ci	prev = bpf_kptr_xchg(&global_mask, curr);
18762306a36Sopenharmony_ci	bpf_rcu_read_unlock();
18862306a36Sopenharmony_ci	if (prev)
18962306a36Sopenharmony_ci		bpf_cpumask_release(prev);
19062306a36Sopenharmony_ci
19162306a36Sopenharmony_ci	return 0;
19262306a36Sopenharmony_ci}
193