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 
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")
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")
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")
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")
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