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(idx: 0, val: 41) ?: create_and_stash(idx: 1, val: 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(p: 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

source code of linux/tools/testing/selftests/bpf/progs/local_kptr_stash.c