286 lines · c
1// SPDX-License-Identifier: GPL-2.02/* 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 plain_local;12 13struct 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 20struct refcounted_node {21 long data;22 struct bpf_rb_node rb_node;23 struct bpf_refcount refcount;24};25 26struct stash {27 struct bpf_spin_lock l;28 struct refcounted_node __kptr *stashed;29};30 31struct {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 38struct plain_local {39 long key;40 long data;41};42 43struct local_with_root {44 long key;45 struct bpf_spin_lock l;46 struct bpf_rb_root r __contains(node_data, node);47};48 49struct 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 struct58 * If it's not included, a fwd reference for node_data will be generated but59 * no struct. Example BTF of "node" field in map_value when not included:60 *61 * [10] PTR '(anon)' type_id=3562 * [34] FWD 'node_data' fwd_kind=struct63 * [35] TYPE_TAG 'kptr_ref' type_id=3464 *65 * (with no node_data struct defined)66 * Had to do the same w/ bpf_kfunc_call_test_release below67 */68struct node_data *just_here_because_btf_bug;69struct refcounted_node *just_here_because_btf_bug2;70 71struct {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 78static 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 89static 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 124SEC("tc")125long stash_rb_nodes(void *ctx)126{127 return create_and_stash(0, 41) ?: create_and_stash(1, 42);128}129 130SEC("tc")131long 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 152SEC("tc")153long 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 187SEC("tc")188long 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 216SEC("tc")217long 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 233SEC("tc")234long 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_node246 * has been stashed before this program executes247 */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 */260SEC("tc")261long 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 285char _license[] SEC("license") = "GPL";286