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