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 
10 struct node_data {
11 	long key;
12 	long data;
13 	struct bpf_rb_node node;
14 };
15 
16 struct map_value {
17 	struct prog_test_ref_kfunc *not_kptr;
18 	struct prog_test_ref_kfunc __kptr *val;
19 	struct node_data __kptr *node;
20 };
21 
22 /* This is necessary so that LLVM generates BTF for node_data struct
23  * If it's not included, a fwd reference for node_data will be generated but
24  * no struct. Example BTF of "node" field in map_value when not included:
25  *
26  * [10] PTR '(anon)' type_id=35
27  * [34] FWD 'node_data' fwd_kind=struct
28  * [35] TYPE_TAG 'kptr_ref' type_id=34
29  *
30  * (with no node_data struct defined)
31  * Had to do the same w/ bpf_kfunc_call_test_release below
32  */
33 struct node_data *just_here_because_btf_bug;
34 
35 extern void bpf_kfunc_call_test_release(struct prog_test_ref_kfunc *p) __ksym;
36 
37 struct {
38 	__uint(type, BPF_MAP_TYPE_ARRAY);
39 	__type(key, int);
40 	__type(value, struct map_value);
41 	__uint(max_entries, 2);
42 } some_nodes SEC(".maps");
43 
44 static int create_and_stash(int idx, int val)
45 {
46 	struct map_value *mapval;
47 	struct node_data *res;
48 
49 	mapval = bpf_map_lookup_elem(&some_nodes, &idx);
50 	if (!mapval)
51 		return 1;
52 
53 	res = bpf_obj_new(typeof(*res));
54 	if (!res)
55 		return 1;
56 	res->key = val;
57 
58 	res = bpf_kptr_xchg(&mapval->node, res);
59 	if (res)
60 		bpf_obj_drop(res);
61 	return 0;
62 }
63 
64 SEC("tc")
65 long stash_rb_nodes(void *ctx)
66 {
67 	return create_and_stash(0, 41) ?: create_and_stash(1, 42);
68 }
69 
70 SEC("tc")
71 long unstash_rb_node(void *ctx)
72 {
73 	struct map_value *mapval;
74 	struct node_data *res;
75 	long retval;
76 	int key = 1;
77 
78 	mapval = bpf_map_lookup_elem(&some_nodes, &key);
79 	if (!mapval)
80 		return 1;
81 
82 	res = bpf_kptr_xchg(&mapval->node, NULL);
83 	if (res) {
84 		retval = res->key;
85 		bpf_obj_drop(res);
86 		return retval;
87 	}
88 	return 1;
89 }
90 
91 SEC("tc")
92 long stash_test_ref_kfunc(void *ctx)
93 {
94 	struct prog_test_ref_kfunc *res;
95 	struct map_value *mapval;
96 	int key = 0;
97 
98 	mapval = bpf_map_lookup_elem(&some_nodes, &key);
99 	if (!mapval)
100 		return 1;
101 
102 	res = bpf_kptr_xchg(&mapval->val, NULL);
103 	if (res)
104 		bpf_kfunc_call_test_release(res);
105 	return 0;
106 }
107 
108 char _license[] SEC("license") = "GPL";
109