1// SPDX-License-Identifier: GPL-2.0
2/* Copyright (c) 2023 Meta Platforms, Inc. and affiliates. */
3
4#include <vmlinux.h>
5#include <bpf/bpf_tracing.h>
6#include <bpf/bpf_helpers.h>
7#include <bpf/bpf_core_read.h>
8#include "../bpf_experimental.h"
9#include "../bpf_testmod/bpf_testmod_kfunc.h"
10
11struct node_data {
12	long key;
13	long data;
14	struct bpf_rb_node node;
15};
16
17struct refcounted_node {
18	long data;
19	struct bpf_rb_node rb_node;
20	struct bpf_refcount refcount;
21};
22
23struct stash {
24	struct bpf_spin_lock l;
25	struct refcounted_node __kptr *stashed;
26};
27
28struct {
29	__uint(type, BPF_MAP_TYPE_ARRAY);
30	__type(key, int);
31	__type(value, struct stash);
32	__uint(max_entries, 10);
33} refcounted_node_stash SEC(".maps");
34
35struct plain_local {
36	long key;
37	long data;
38};
39
40struct local_with_root {
41	long key;
42	struct bpf_spin_lock l;
43	struct bpf_rb_root r __contains(node_data, node);
44};
45
46struct map_value {
47	struct prog_test_ref_kfunc *not_kptr;
48	struct prog_test_ref_kfunc __kptr *val;
49	struct node_data __kptr *node;
50	struct plain_local __kptr *plain;
51	struct local_with_root __kptr *local_root;
52};
53
54/* This is necessary so that LLVM generates BTF for node_data struct
55 * If it's not included, a fwd reference for node_data will be generated but
56 * no struct. Example BTF of "node" field in map_value when not included:
57 *
58 * [10] PTR '(anon)' type_id=35
59 * [34] FWD 'node_data' fwd_kind=struct
60 * [35] TYPE_TAG 'kptr_ref' type_id=34
61 *
62 * (with no node_data struct defined)
63 * Had to do the same w/ bpf_kfunc_call_test_release below
64 */
65struct node_data *just_here_because_btf_bug;
66struct refcounted_node *just_here_because_btf_bug2;
67
68struct {
69	__uint(type, BPF_MAP_TYPE_ARRAY);
70	__type(key, int);
71	__type(value, struct map_value);
72	__uint(max_entries, 2);
73} some_nodes SEC(".maps");
74
75static bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
76{
77	struct node_data *node_a;
78	struct node_data *node_b;
79
80	node_a = container_of(a, struct node_data, node);
81	node_b = container_of(b, struct node_data, node);
82
83	return node_a->key < node_b->key;
84}
85
86static int create_and_stash(int idx, int val)
87{
88	struct map_value *mapval;
89	struct node_data *res;
90
91	mapval = bpf_map_lookup_elem(&some_nodes, &idx);
92	if (!mapval)
93		return 1;
94
95	res = bpf_obj_new(typeof(*res));
96	if (!res)
97		return 1;
98	res->key = val;
99
100	res = bpf_kptr_xchg(&mapval->node, res);
101	if (res)
102		bpf_obj_drop(res);
103	return 0;
104}
105
106SEC("tc")
107long stash_rb_nodes(void *ctx)
108{
109	return create_and_stash(0, 41) ?: create_and_stash(1, 42);
110}
111
112SEC("tc")
113long stash_plain(void *ctx)
114{
115	struct map_value *mapval;
116	struct plain_local *res;
117	int idx = 0;
118
119	mapval = bpf_map_lookup_elem(&some_nodes, &idx);
120	if (!mapval)
121		return 1;
122
123	res = bpf_obj_new(typeof(*res));
124	if (!res)
125		return 1;
126	res->key = 41;
127
128	res = bpf_kptr_xchg(&mapval->plain, res);
129	if (res)
130		bpf_obj_drop(res);
131	return 0;
132}
133
134SEC("tc")
135long stash_local_with_root(void *ctx)
136{
137	struct local_with_root *res;
138	struct map_value *mapval;
139	struct node_data *n;
140	int idx = 0;
141
142	mapval = bpf_map_lookup_elem(&some_nodes, &idx);
143	if (!mapval)
144		return 1;
145
146	res = bpf_obj_new(typeof(*res));
147	if (!res)
148		return 2;
149	res->key = 41;
150
151	n = bpf_obj_new(typeof(*n));
152	if (!n) {
153		bpf_obj_drop(res);
154		return 3;
155	}
156
157	bpf_spin_lock(&res->l);
158	bpf_rbtree_add(&res->r, &n->node, less);
159	bpf_spin_unlock(&res->l);
160
161	res = bpf_kptr_xchg(&mapval->local_root, res);
162	if (res) {
163		bpf_obj_drop(res);
164		return 4;
165	}
166	return 0;
167}
168
169SEC("tc")
170long unstash_rb_node(void *ctx)
171{
172	struct map_value *mapval;
173	struct node_data *res;
174	long retval;
175	int key = 1;
176
177	mapval = bpf_map_lookup_elem(&some_nodes, &key);
178	if (!mapval)
179		return 1;
180
181	res = bpf_kptr_xchg(&mapval->node, NULL);
182	if (res) {
183		retval = res->key;
184		bpf_obj_drop(res);
185		return retval;
186	}
187	return 1;
188}
189
190SEC("tc")
191long stash_test_ref_kfunc(void *ctx)
192{
193	struct prog_test_ref_kfunc *res;
194	struct map_value *mapval;
195	int key = 0;
196
197	mapval = bpf_map_lookup_elem(&some_nodes, &key);
198	if (!mapval)
199		return 1;
200
201	res = bpf_kptr_xchg(&mapval->val, NULL);
202	if (res)
203		bpf_kfunc_call_test_release(res);
204	return 0;
205}
206
207SEC("tc")
208long refcount_acquire_without_unstash(void *ctx)
209{
210	struct refcounted_node *p;
211	struct stash *s;
212	int ret = 0;
213
214	s = bpf_map_lookup_elem(&refcounted_node_stash, &ret);
215	if (!s)
216		return 1;
217
218	if (!s->stashed)
219		/* refcount_acquire failure is expected when no refcounted_node
220		 * has been stashed before this program executes
221		 */
222		return 2;
223
224	p = bpf_refcount_acquire(s->stashed);
225	if (!p)
226		return 3;
227
228	ret = s->stashed ? s->stashed->data : -1;
229	bpf_obj_drop(p);
230	return ret;
231}
232
233/* Helper for refcount_acquire_without_unstash test */
234SEC("tc")
235long stash_refcounted_node(void *ctx)
236{
237	struct refcounted_node *p;
238	struct stash *s;
239	int key = 0;
240
241	s = bpf_map_lookup_elem(&refcounted_node_stash, &key);
242	if (!s)
243		return 1;
244
245	p = bpf_obj_new(typeof(*p));
246	if (!p)
247		return 2;
248	p->data = 42;
249
250	p = bpf_kptr_xchg(&s->stashed, p);
251	if (p) {
252		bpf_obj_drop(p);
253		return 3;
254	}
255
256	return 0;
257}
258
259char _license[] SEC("license") = "GPL";
260