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