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 local_with_root { 41 long key; 42 struct bpf_spin_lock l; 43 struct bpf_rb_root r __contains(node_data, node); 44 }; 45 46 struct 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 */ 65 struct node_data *just_here_because_btf_bug; 66 struct refcounted_node *just_here_because_btf_bug2; 67 68 struct { 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 75 static 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 86 static 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 106 SEC("tc") 107 long stash_rb_nodes(void *ctx) 108 { 109 return create_and_stash(0, 41) ?: create_and_stash(1, 42); 110 } 111 112 SEC("tc") 113 long 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 134 SEC("tc") 135 long 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 169 SEC("tc") 170 long 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 190 SEC("tc") 191 long 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 207 SEC("tc") 208 long 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 */ 234 SEC("tc") 235 long 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 259 char _license[] SEC("license") = "GPL"; 260