| // SPDX-License-Identifier: GPL-2.0 |
| /* Copyright (c) 2023 Meta Platforms, Inc. and affiliates. */ |
| |
| #include <vmlinux.h> |
| #include <bpf/bpf_tracing.h> |
| #include <bpf/bpf_helpers.h> |
| #include <bpf/bpf_core_read.h> |
| #include "../bpf_experimental.h" |
| #include "../bpf_testmod/bpf_testmod_kfunc.h" |
| |
| struct node_data { |
| long key; |
| long data; |
| struct bpf_rb_node node; |
| }; |
| |
| struct refcounted_node { |
| long data; |
| struct bpf_rb_node rb_node; |
| struct bpf_refcount refcount; |
| }; |
| |
| struct stash { |
| struct bpf_spin_lock l; |
| struct refcounted_node __kptr *stashed; |
| }; |
| |
| struct { |
| __uint(type, BPF_MAP_TYPE_ARRAY); |
| __type(key, int); |
| __type(value, struct stash); |
| __uint(max_entries, 10); |
| } refcounted_node_stash SEC(".maps"); |
| |
| struct plain_local { |
| long key; |
| long data; |
| }; |
| |
| struct local_with_root { |
| long key; |
| struct bpf_spin_lock l; |
| struct bpf_rb_root r __contains(node_data, node); |
| }; |
| |
| struct map_value { |
| struct prog_test_ref_kfunc *not_kptr; |
| struct prog_test_ref_kfunc __kptr *val; |
| struct node_data __kptr *node; |
| struct plain_local __kptr *plain; |
| struct local_with_root __kptr *local_root; |
| }; |
| |
| /* This is necessary so that LLVM generates BTF for node_data struct |
| * If it's not included, a fwd reference for node_data will be generated but |
| * no struct. Example BTF of "node" field in map_value when not included: |
| * |
| * [10] PTR '(anon)' type_id=35 |
| * [34] FWD 'node_data' fwd_kind=struct |
| * [35] TYPE_TAG 'kptr_ref' type_id=34 |
| * |
| * (with no node_data struct defined) |
| * Had to do the same w/ bpf_kfunc_call_test_release below |
| */ |
| struct node_data *just_here_because_btf_bug; |
| struct refcounted_node *just_here_because_btf_bug2; |
| |
| struct { |
| __uint(type, BPF_MAP_TYPE_ARRAY); |
| __type(key, int); |
| __type(value, struct map_value); |
| __uint(max_entries, 2); |
| } some_nodes SEC(".maps"); |
| |
| static bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b) |
| { |
| struct node_data *node_a; |
| struct node_data *node_b; |
| |
| node_a = container_of(a, struct node_data, node); |
| node_b = container_of(b, struct node_data, node); |
| |
| return node_a->key < node_b->key; |
| } |
| |
| static int create_and_stash(int idx, int val) |
| { |
| struct map_value *mapval; |
| struct node_data *res; |
| |
| mapval = bpf_map_lookup_elem(&some_nodes, &idx); |
| if (!mapval) |
| return 1; |
| |
| res = bpf_obj_new(typeof(*res)); |
| if (!res) |
| return 1; |
| res->key = val; |
| |
| res = bpf_kptr_xchg(&mapval->node, res); |
| if (res) |
| bpf_obj_drop(res); |
| return 0; |
| } |
| |
| SEC("tc") |
| long stash_rb_nodes(void *ctx) |
| { |
| return create_and_stash(0, 41) ?: create_and_stash(1, 42); |
| } |
| |
| SEC("tc") |
| long stash_plain(void *ctx) |
| { |
| struct map_value *mapval; |
| struct plain_local *res; |
| int idx = 0; |
| |
| mapval = bpf_map_lookup_elem(&some_nodes, &idx); |
| if (!mapval) |
| return 1; |
| |
| res = bpf_obj_new(typeof(*res)); |
| if (!res) |
| return 1; |
| res->key = 41; |
| |
| res = bpf_kptr_xchg(&mapval->plain, res); |
| if (res) |
| bpf_obj_drop(res); |
| return 0; |
| } |
| |
| SEC("tc") |
| long stash_local_with_root(void *ctx) |
| { |
| struct local_with_root *res; |
| struct map_value *mapval; |
| struct node_data *n; |
| int idx = 0; |
| |
| mapval = bpf_map_lookup_elem(&some_nodes, &idx); |
| if (!mapval) |
| return 1; |
| |
| res = bpf_obj_new(typeof(*res)); |
| if (!res) |
| return 2; |
| res->key = 41; |
| |
| n = bpf_obj_new(typeof(*n)); |
| if (!n) { |
| bpf_obj_drop(res); |
| return 3; |
| } |
| |
| bpf_spin_lock(&res->l); |
| bpf_rbtree_add(&res->r, &n->node, less); |
| bpf_spin_unlock(&res->l); |
| |
| res = bpf_kptr_xchg(&mapval->local_root, res); |
| if (res) { |
| bpf_obj_drop(res); |
| return 4; |
| } |
| return 0; |
| } |
| |
| SEC("tc") |
| long unstash_rb_node(void *ctx) |
| { |
| struct map_value *mapval; |
| struct node_data *res; |
| long retval; |
| int key = 1; |
| |
| mapval = bpf_map_lookup_elem(&some_nodes, &key); |
| if (!mapval) |
| return 1; |
| |
| res = bpf_kptr_xchg(&mapval->node, NULL); |
| if (res) { |
| retval = res->key; |
| bpf_obj_drop(res); |
| return retval; |
| } |
| return 1; |
| } |
| |
| SEC("tc") |
| long stash_test_ref_kfunc(void *ctx) |
| { |
| struct prog_test_ref_kfunc *res; |
| struct map_value *mapval; |
| int key = 0; |
| |
| mapval = bpf_map_lookup_elem(&some_nodes, &key); |
| if (!mapval) |
| return 1; |
| |
| res = bpf_kptr_xchg(&mapval->val, NULL); |
| if (res) |
| bpf_kfunc_call_test_release(res); |
| return 0; |
| } |
| |
| SEC("tc") |
| long refcount_acquire_without_unstash(void *ctx) |
| { |
| struct refcounted_node *p; |
| struct stash *s; |
| int ret = 0; |
| |
| s = bpf_map_lookup_elem(&refcounted_node_stash, &ret); |
| if (!s) |
| return 1; |
| |
| if (!s->stashed) |
| /* refcount_acquire failure is expected when no refcounted_node |
| * has been stashed before this program executes |
| */ |
| return 2; |
| |
| p = bpf_refcount_acquire(s->stashed); |
| if (!p) |
| return 3; |
| |
| ret = s->stashed ? s->stashed->data : -1; |
| bpf_obj_drop(p); |
| return ret; |
| } |
| |
| /* Helper for refcount_acquire_without_unstash test */ |
| SEC("tc") |
| long stash_refcounted_node(void *ctx) |
| { |
| struct refcounted_node *p; |
| struct stash *s; |
| int key = 0; |
| |
| s = bpf_map_lookup_elem(&refcounted_node_stash, &key); |
| if (!s) |
| return 1; |
| |
| p = bpf_obj_new(typeof(*p)); |
| if (!p) |
| return 2; |
| p->data = 42; |
| |
| p = bpf_kptr_xchg(&s->stashed, p); |
| if (p) { |
| bpf_obj_drop(p); |
| return 3; |
| } |
| |
| return 0; |
| } |
| |
| char _license[] SEC("license") = "GPL"; |