xref: /linux/tools/testing/selftests/bpf/progs/task_local_data.bpf.h (revision 07fdad3a93756b872da7b53647715c48d0f4a2d0)
1 /* SPDX-License-Identifier: GPL-2.0 */
2 #ifndef __TASK_LOCAL_DATA_BPF_H
3 #define __TASK_LOCAL_DATA_BPF_H
4 
5 /*
6  * Task local data is a library that facilitates sharing per-task data
7  * between user space and bpf programs.
8  *
9  *
10  * USAGE
11  *
12  * A TLD, an entry of data in task local data, first needs to be created by the
13  * user space. This is done by calling user space API, TLD_DEFINE_KEY() or
14  * tld_create_key(), with the name of the TLD and the size.
15  *
16  * TLD_DEFINE_KEY(prio, "priority", sizeof(int));
17  *
18  * or
19  *
20  * void func_call(...) {
21  *     tld_key_t prio, in_cs;
22  *
23  *     prio = tld_create_key("priority", sizeof(int));
24  *     in_cs = tld_create_key("in_critical_section", sizeof(bool));
25  *     ...
26  *
27  * A key associated with the TLD, which has an opaque type tld_key_t, will be
28  * initialized or returned. It can be used to get a pointer to the TLD in the
29  * user space by calling tld_get_data().
30  *
31  * In a bpf program, tld_object_init() first needs to be called to initialized a
32  * tld_object on the stack. Then, TLDs can be accessed by calling tld_get_data().
33  * The API will try to fetch the key by the name and use it to locate the data.
34  * A pointer to the TLD will be returned. It also caches the key in a task local
35  * storage map, tld_key_map, whose value type, struct tld_keys, must be defined
36  * by the developer.
37  *
38  * struct tld_keys {
39  *     tld_key_t prio;
40  *     tld_key_t in_cs;
41  * };
42  *
43  * SEC("struct_ops")
44  * void prog(struct task_struct task, ...)
45  * {
46  *     struct tld_object tld_obj;
47  *     int err, *p;
48  *
49  *     err = tld_object_init(task, &tld_obj);
50  *     if (err)
51  *         return;
52  *
53  *     p = tld_get_data(&tld_obj, prio, "priority", sizeof(int));
54  *     if (p)
55  *         // do something depending on *p
56  */
57 #include <errno.h>
58 #include <bpf/bpf_helpers.h>
59 
60 #define TLD_ROUND_MASK(x, y) ((__typeof__(x))((y) - 1))
61 #define TLD_ROUND_UP(x, y) ((((x) - 1) | TLD_ROUND_MASK(x, y)) + 1)
62 
63 #define TLD_MAX_DATA_CNT (__PAGE_SIZE / sizeof(struct tld_metadata) - 1)
64 
65 #ifndef TLD_NAME_LEN
66 #define TLD_NAME_LEN 62
67 #endif
68 
69 #ifndef TLD_KEY_MAP_CREATE_RETRY
70 #define TLD_KEY_MAP_CREATE_RETRY 10
71 #endif
72 
73 typedef struct {
74 	__s16 off;
75 } tld_key_t;
76 
77 struct tld_metadata {
78 	char name[TLD_NAME_LEN];
79 	__u16 size;
80 };
81 
82 struct tld_meta_u {
83 	__u8 cnt;
84 	__u16 size;
85 	struct tld_metadata metadata[TLD_MAX_DATA_CNT];
86 };
87 
88 struct tld_data_u {
89 	__u64 start; /* offset of tld_data_u->data in a page */
90 	char data[__PAGE_SIZE - sizeof(__u64)];
91 };
92 
93 struct tld_map_value {
94 	struct tld_data_u __uptr *data;
95 	struct tld_meta_u __uptr *meta;
96 };
97 
98 typedef struct tld_uptr_dummy {
99 	struct tld_data_u data[0];
100 	struct tld_meta_u meta[0];
101 } *tld_uptr_dummy_t;
102 
103 struct tld_object {
104 	struct tld_map_value *data_map;
105 	struct tld_keys *key_map;
106 	/*
107 	 * Force the compiler to generate the actual definition of tld_meta_u
108 	 * and tld_data_u in BTF. Without it, tld_meta_u and u_tld_data will
109 	 * be BTF_KIND_FWD.
110 	 */
111 	tld_uptr_dummy_t dummy[0];
112 };
113 
114 /*
115  * Map value of tld_key_map for caching keys. Must be defined by the developer.
116  * Members should be tld_key_t and passed to the 3rd argument of tld_fetch_key().
117  */
118 struct tld_keys;
119 
120 struct {
121 	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
122 	__uint(map_flags, BPF_F_NO_PREALLOC);
123 	__type(key, int);
124 	__type(value, struct tld_map_value);
125 } tld_data_map SEC(".maps");
126 
127 struct {
128 	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
129 	__uint(map_flags, BPF_F_NO_PREALLOC);
130 	__type(key, int);
131 	__type(value, struct tld_keys);
132 } tld_key_map SEC(".maps");
133 
134 /**
135  * tld_object_init() - Initialize a tld_object.
136  *
137  * @task: The task_struct of the target task
138  * @tld_obj: A pointer to a tld_object to be initialized
139  *
140  * Return 0 on success; -ENODATA if the user space did not initialize task local data
141  * for the current task through tld_get_data(); -ENOMEM if the creation of tld_key_map
142  * fails
143  */
144 __attribute__((unused))
145 static int tld_object_init(struct task_struct *task, struct tld_object *tld_obj)
146 {
147 	int i;
148 
149 	tld_obj->data_map = bpf_task_storage_get(&tld_data_map, task, 0, 0);
150 	if (!tld_obj->data_map)
151 		return -ENODATA;
152 
153 	bpf_for(i, 0, TLD_KEY_MAP_CREATE_RETRY) {
154 		tld_obj->key_map = bpf_task_storage_get(&tld_key_map, task, 0,
155 							BPF_LOCAL_STORAGE_GET_F_CREATE);
156 		if (likely(tld_obj->key_map))
157 			break;
158 	}
159 	if (!tld_obj->key_map)
160 		return -ENOMEM;
161 
162 	return 0;
163 }
164 
165 /*
166  * Return the offset of TLD if @name is found. Otherwise, return the current TLD count
167  * using the nonpositive range so that the next tld_get_data() can skip fetching key if
168  * no new TLD is added or start comparing name from the first newly added TLD.
169  */
170 __attribute__((unused))
171 static int __tld_fetch_key(struct tld_object *tld_obj, const char *name, int i_start)
172 {
173 	struct tld_metadata *metadata;
174 	int i, cnt, start, off = 0;
175 
176 	if (!tld_obj->data_map || !tld_obj->data_map->data || !tld_obj->data_map->meta)
177 		return 0;
178 
179 	start = tld_obj->data_map->data->start;
180 	cnt = tld_obj->data_map->meta->cnt;
181 	metadata = tld_obj->data_map->meta->metadata;
182 
183 	bpf_for(i, 0, cnt) {
184 		if (i >= TLD_MAX_DATA_CNT)
185 			break;
186 
187 		if (i >= i_start && !bpf_strncmp(metadata[i].name, TLD_NAME_LEN, name))
188 			return start + off;
189 
190 		off += TLD_ROUND_UP(metadata[i].size, 8);
191 	}
192 
193 	return -cnt;
194 }
195 
196 /**
197  * tld_get_data() - Retrieve a pointer to the TLD associated with the name.
198  *
199  * @tld_obj: A pointer to a valid tld_object initialized by tld_object_init()
200  * @key: The cached key of the TLD in tld_key_map
201  * @name: The name of the key associated with a TLD
202  * @size: The size of the TLD. Must be a known constant value
203  *
204  * Return a pointer to the TLD associated with @name; NULL if not found or @size is too
205  * big. @key is used to cache the key if the TLD is found to speed up subsequent calls.
206  * It should be defined as an member of tld_keys of tld_key_t type by the developer.
207  */
208 #define tld_get_data(tld_obj, key, name, size)						\
209 	({										\
210 		void *data = NULL, *_data = (tld_obj)->data_map->data;			\
211 		long off = (tld_obj)->key_map->key.off;					\
212 		int cnt;								\
213 											\
214 		if (likely(_data)) {							\
215 			if (likely(off > 0)) {						\
216 				barrier_var(off);					\
217 				if (likely(off < __PAGE_SIZE - size))			\
218 					data = _data + off;				\
219 			} else {							\
220 				cnt = -(off);						\
221 				if (likely((tld_obj)->data_map->meta) &&		\
222 				    cnt < (tld_obj)->data_map->meta->cnt) {		\
223 					off = __tld_fetch_key(tld_obj, name, cnt);	\
224 					(tld_obj)->key_map->key.off = off;		\
225 											\
226 					if (likely(off < __PAGE_SIZE - size)) {		\
227 						barrier_var(off);			\
228 						if (off > 0)				\
229 							data = _data + off;		\
230 					}						\
231 				}							\
232 			}								\
233 		}									\
234 		data;									\
235 	})
236 
237 #endif
238