xref: /linux/tools/testing/selftests/bpf/progs/local_kptr_stash.c (revision 48eb03dd26304c24f03bdbb9382e89c8564e71df)
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 
11 struct node_data {
12 	long key;
13 	long data;
14 	struct bpf_rb_node node;
15 };
16 
17 struct refcounted_node {
18 	long data;
19 	struct bpf_rb_node rb_node;
20 	struct bpf_refcount refcount;
21 };
22 
23 struct stash {
24 	struct bpf_spin_lock l;
25 	struct refcounted_node __kptr *stashed;
26 };
27 
28 struct {
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 
35 struct plain_local {
36 	long key;
37 	long data;
38 };
39 
40 struct map_value {
41 	struct prog_test_ref_kfunc *not_kptr;
42 	struct prog_test_ref_kfunc __kptr *val;
43 	struct node_data __kptr *node;
44 	struct plain_local __kptr *plain;
45 };
46 
47 /* This is necessary so that LLVM generates BTF for node_data struct
48  * If it's not included, a fwd reference for node_data will be generated but
49  * no struct. Example BTF of "node" field in map_value when not included:
50  *
51  * [10] PTR '(anon)' type_id=35
52  * [34] FWD 'node_data' fwd_kind=struct
53  * [35] TYPE_TAG 'kptr_ref' type_id=34
54  *
55  * (with no node_data struct defined)
56  * Had to do the same w/ bpf_kfunc_call_test_release below
57  */
58 struct node_data *just_here_because_btf_bug;
59 struct refcounted_node *just_here_because_btf_bug2;
60 
61 struct {
62 	__uint(type, BPF_MAP_TYPE_ARRAY);
63 	__type(key, int);
64 	__type(value, struct map_value);
65 	__uint(max_entries, 2);
66 } some_nodes SEC(".maps");
67 
68 static int create_and_stash(int idx, int val)
69 {
70 	struct map_value *mapval;
71 	struct node_data *res;
72 
73 	mapval = bpf_map_lookup_elem(&some_nodes, &idx);
74 	if (!mapval)
75 		return 1;
76 
77 	res = bpf_obj_new(typeof(*res));
78 	if (!res)
79 		return 1;
80 	res->key = val;
81 
82 	res = bpf_kptr_xchg(&mapval->node, res);
83 	if (res)
84 		bpf_obj_drop(res);
85 	return 0;
86 }
87 
88 SEC("tc")
89 long stash_rb_nodes(void *ctx)
90 {
91 	return create_and_stash(0, 41) ?: create_and_stash(1, 42);
92 }
93 
94 SEC("tc")
95 long stash_plain(void *ctx)
96 {
97 	struct map_value *mapval;
98 	struct plain_local *res;
99 	int idx = 0;
100 
101 	mapval = bpf_map_lookup_elem(&some_nodes, &idx);
102 	if (!mapval)
103 		return 1;
104 
105 	res = bpf_obj_new(typeof(*res));
106 	if (!res)
107 		return 1;
108 	res->key = 41;
109 
110 	res = bpf_kptr_xchg(&mapval->plain, res);
111 	if (res)
112 		bpf_obj_drop(res);
113 	return 0;
114 }
115 
116 SEC("tc")
117 long unstash_rb_node(void *ctx)
118 {
119 	struct map_value *mapval;
120 	struct node_data *res;
121 	long retval;
122 	int key = 1;
123 
124 	mapval = bpf_map_lookup_elem(&some_nodes, &key);
125 	if (!mapval)
126 		return 1;
127 
128 	res = bpf_kptr_xchg(&mapval->node, NULL);
129 	if (res) {
130 		retval = res->key;
131 		bpf_obj_drop(res);
132 		return retval;
133 	}
134 	return 1;
135 }
136 
137 SEC("tc")
138 long stash_test_ref_kfunc(void *ctx)
139 {
140 	struct prog_test_ref_kfunc *res;
141 	struct map_value *mapval;
142 	int key = 0;
143 
144 	mapval = bpf_map_lookup_elem(&some_nodes, &key);
145 	if (!mapval)
146 		return 1;
147 
148 	res = bpf_kptr_xchg(&mapval->val, NULL);
149 	if (res)
150 		bpf_kfunc_call_test_release(res);
151 	return 0;
152 }
153 
154 SEC("tc")
155 long refcount_acquire_without_unstash(void *ctx)
156 {
157 	struct refcounted_node *p;
158 	struct stash *s;
159 	int ret = 0;
160 
161 	s = bpf_map_lookup_elem(&refcounted_node_stash, &ret);
162 	if (!s)
163 		return 1;
164 
165 	if (!s->stashed)
166 		/* refcount_acquire failure is expected when no refcounted_node
167 		 * has been stashed before this program executes
168 		 */
169 		return 2;
170 
171 	p = bpf_refcount_acquire(s->stashed);
172 	if (!p)
173 		return 3;
174 
175 	ret = s->stashed ? s->stashed->data : -1;
176 	bpf_obj_drop(p);
177 	return ret;
178 }
179 
180 /* Helper for refcount_acquire_without_unstash test */
181 SEC("tc")
182 long stash_refcounted_node(void *ctx)
183 {
184 	struct refcounted_node *p;
185 	struct stash *s;
186 	int key = 0;
187 
188 	s = bpf_map_lookup_elem(&refcounted_node_stash, &key);
189 	if (!s)
190 		return 1;
191 
192 	p = bpf_obj_new(typeof(*p));
193 	if (!p)
194 		return 2;
195 	p->data = 42;
196 
197 	p = bpf_kptr_xchg(&s->stashed, p);
198 	if (p) {
199 		bpf_obj_drop(p);
200 		return 3;
201 	}
202 
203 	return 0;
204 }
205 
206 char _license[] SEC("license") = "GPL";
207