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