1// SPDX-License-Identifier: GPL-2.0
2/* Copyright (c) 2024 Meta Platforms, Inc. and affiliates. */
3
4#include <vmlinux.h>
5#include <bpf/bpf_helpers.h>
6#include <bpf/bpf_tracing.h>
7#include <bpf/bpf_core_read.h>
8#include "bpf_misc.h"
9#include "xdp_metadata.h"
10#include "bpf_kfuncs.h"
11
12extern struct task_struct *bpf_task_acquire(struct task_struct *p) __ksym __weak;
13extern void bpf_task_release(struct task_struct *p) __ksym __weak;
14
15__weak int subprog_trusted_task_nullable(struct task_struct *task __arg_trusted __arg_nullable)
16{
17	if (!task)
18		return 0;
19	return task->pid + task->tgid;
20}
21
22__weak int subprog_trusted_task_nullable_extra_layer(struct task_struct *task __arg_trusted __arg_nullable)
23{
24	return subprog_trusted_task_nullable(task) + subprog_trusted_task_nullable(NULL);
25}
26
27SEC("?tp_btf/task_newtask")
28__success __log_level(2)
29__msg("Validating subprog_trusted_task_nullable() func#1...")
30__msg(": R1=trusted_ptr_or_null_task_struct(")
31int trusted_task_arg_nullable(void *ctx)
32{
33	struct task_struct *t1 = bpf_get_current_task_btf();
34	struct task_struct *t2 = bpf_task_acquire(t1);
35	int res = 0;
36
37	/* known NULL */
38	res += subprog_trusted_task_nullable(NULL);
39
40	/* known non-NULL */
41	res += subprog_trusted_task_nullable(t1);
42	res += subprog_trusted_task_nullable_extra_layer(t1);
43
44	/* unknown if NULL or not */
45	res += subprog_trusted_task_nullable(t2);
46	res += subprog_trusted_task_nullable_extra_layer(t2);
47
48	if (t2) {
49		/* known non-NULL after explicit NULL check, just in case */
50		res += subprog_trusted_task_nullable(t2);
51		res += subprog_trusted_task_nullable_extra_layer(t2);
52
53		bpf_task_release(t2);
54	}
55
56	return res;
57}
58
59__weak int subprog_trusted_task_nonnull(struct task_struct *task __arg_trusted)
60{
61	return task->pid + task->tgid;
62}
63
64SEC("?kprobe")
65__failure __log_level(2)
66__msg("R1 type=scalar expected=ptr_, trusted_ptr_, rcu_ptr_")
67__msg("Caller passes invalid args into func#1 ('subprog_trusted_task_nonnull')")
68int trusted_task_arg_nonnull_fail1(void *ctx)
69{
70	return subprog_trusted_task_nonnull(NULL);
71}
72
73SEC("?tp_btf/task_newtask")
74__failure __log_level(2)
75__msg("R1 type=ptr_or_null_ expected=ptr_, trusted_ptr_, rcu_ptr_")
76__msg("Caller passes invalid args into func#1 ('subprog_trusted_task_nonnull')")
77int trusted_task_arg_nonnull_fail2(void *ctx)
78{
79	struct task_struct *t = bpf_get_current_task_btf();
80	struct task_struct *nullable;
81	int res;
82
83	nullable = bpf_task_acquire(t);
84
85	 /* should fail, PTR_TO_BTF_ID_OR_NULL */
86	res = subprog_trusted_task_nonnull(nullable);
87
88	if (nullable)
89		bpf_task_release(nullable);
90
91	return res;
92}
93
94SEC("?kprobe")
95__success __log_level(2)
96__msg("Validating subprog_trusted_task_nonnull() func#1...")
97__msg(": R1=trusted_ptr_task_struct(")
98int trusted_task_arg_nonnull(void *ctx)
99{
100	struct task_struct *t = bpf_get_current_task_btf();
101
102	return subprog_trusted_task_nonnull(t);
103}
104
105struct task_struct___local {} __attribute__((preserve_access_index));
106
107__weak int subprog_nullable_task_flavor(
108	struct task_struct___local *task __arg_trusted __arg_nullable)
109{
110	char buf[16];
111
112	if (!task)
113		return 0;
114
115	return bpf_copy_from_user_task(&buf, sizeof(buf), NULL, (void *)task, 0);
116}
117
118SEC("?uprobe.s")
119__success __log_level(2)
120__msg("Validating subprog_nullable_task_flavor() func#1...")
121__msg(": R1=trusted_ptr_or_null_task_struct(")
122int flavor_ptr_nullable(void *ctx)
123{
124	struct task_struct___local *t = (void *)bpf_get_current_task_btf();
125
126	return subprog_nullable_task_flavor(t);
127}
128
129__weak int subprog_nonnull_task_flavor(struct task_struct___local *task __arg_trusted)
130{
131	char buf[16];
132
133	return bpf_copy_from_user_task(&buf, sizeof(buf), NULL, (void *)task, 0);
134}
135
136SEC("?uprobe.s")
137__success __log_level(2)
138__msg("Validating subprog_nonnull_task_flavor() func#1...")
139__msg(": R1=trusted_ptr_task_struct(")
140int flavor_ptr_nonnull(void *ctx)
141{
142	struct task_struct *t = bpf_get_current_task_btf();
143
144	return subprog_nonnull_task_flavor((void *)t);
145}
146
147__weak int subprog_trusted_destroy(struct task_struct *task __arg_trusted)
148{
149	bpf_task_release(task); /* should be rejected */
150
151	return 0;
152}
153
154SEC("?tp_btf/task_newtask")
155__failure __log_level(2)
156__msg("release kernel function bpf_task_release expects refcounted PTR_TO_BTF_ID")
157int BPF_PROG(trusted_destroy_fail, struct task_struct *task, u64 clone_flags)
158{
159	return subprog_trusted_destroy(task);
160}
161
162__weak int subprog_trusted_acq_rel(struct task_struct *task __arg_trusted)
163{
164	struct task_struct *owned;
165
166	owned = bpf_task_acquire(task);
167	if (!owned)
168		return 0;
169
170	bpf_task_release(owned); /* this one is OK, we acquired it locally */
171
172	return 0;
173}
174
175SEC("?tp_btf/task_newtask")
176__success __log_level(2)
177int BPF_PROG(trusted_acq_rel, struct task_struct *task, u64 clone_flags)
178{
179	return subprog_trusted_acq_rel(task);
180}
181
182char _license[] SEC("license") = "GPL";
183