xref: /linux/tools/testing/selftests/bpf/progs/task_local_data.bpf.h (revision 7a5f1cd22d47f8ca4b760b6334378ae42c1bd24b)
1 /* SPDX-License-Identifier: (LGPL-2.1 OR BSD-2-Clause) */
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 	__u16 cnt;
84 	__u16 size;
85 	struct tld_metadata metadata[TLD_MAX_DATA_CNT];
86 };
87 
88 struct tld_data_u {
89 	__u64 unused;
90 	char data[__PAGE_SIZE - sizeof(__u64)] __attribute__((aligned(8)));
91 };
92 
93 struct tld_map_value {
94 	struct tld_data_u __uptr *data;
95 	struct tld_meta_u __uptr *meta;
96 	__u16 start; /* offset of tld_data_u->data in a page */
97 };
98 
99 typedef struct tld_uptr_dummy {
100 	struct tld_data_u data[0];
101 	struct tld_meta_u meta[0];
102 } *tld_uptr_dummy_t;
103 
104 struct tld_object {
105 	struct tld_map_value *data_map;
106 	struct tld_keys *key_map;
107 	/*
108 	 * Force the compiler to generate the actual definition of tld_meta_u
109 	 * and tld_data_u in BTF. Without it, tld_meta_u and u_tld_data will
110 	 * be BTF_KIND_FWD.
111 	 */
112 	tld_uptr_dummy_t dummy[0];
113 };
114 
115 /*
116  * Map value of tld_key_map for caching keys. Must be defined by the developer.
117  * Members should be tld_key_t and passed to the 3rd argument of tld_fetch_key().
118  */
119 struct tld_keys;
120 
121 struct {
122 	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
123 	__uint(map_flags, BPF_F_NO_PREALLOC);
124 	__type(key, int);
125 	__type(value, struct tld_map_value);
126 } tld_data_map SEC(".maps");
127 
128 struct {
129 	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
130 	__uint(map_flags, BPF_F_NO_PREALLOC);
131 	__type(key, int);
132 	__type(value, struct tld_keys);
133 } tld_key_map SEC(".maps");
134 
135 /**
136  * tld_object_init() - Initialize a tld_object.
137  *
138  * @task: The task_struct of the target task
139  * @tld_obj: A pointer to a tld_object to be initialized
140  *
141  * Return 0 on success; -ENODATA if the user space did not initialize task local data
142  * for the current task through tld_get_data(); -ENOMEM if the creation of tld_key_map
143  * fails
144  */
145 __attribute__((unused))
146 static int tld_object_init(struct task_struct *task, struct tld_object *tld_obj)
147 {
148 	int i;
149 
150 	tld_obj->data_map = bpf_task_storage_get(&tld_data_map, task, 0, 0);
151 	if (!tld_obj->data_map)
152 		return -ENODATA;
153 
154 	bpf_for(i, 0, TLD_KEY_MAP_CREATE_RETRY) {
155 		tld_obj->key_map = bpf_task_storage_get(&tld_key_map, task, 0,
156 							BPF_LOCAL_STORAGE_GET_F_CREATE);
157 		if (likely(tld_obj->key_map))
158 			break;
159 	}
160 	if (!tld_obj->key_map)
161 		return -ENOMEM;
162 
163 	return 0;
164 }
165 
166 /*
167  * Return the offset of TLD if @name is found. Otherwise, return the current TLD count
168  * using the nonpositive range so that the next tld_get_data() can skip fetching key if
169  * no new TLD is added or start comparing name from the first newly added TLD.
170  */
171 __attribute__((unused))
172 static int __tld_fetch_key(struct tld_object *tld_obj, const char *name, int i_start)
173 {
174 	struct tld_metadata *metadata;
175 	int i, cnt, start, off = 0;
176 
177 	if (!tld_obj->data_map || !tld_obj->data_map->data || !tld_obj->data_map->meta)
178 		return 0;
179 
180 	start = tld_obj->data_map->start;
181 	cnt = tld_obj->data_map->meta->cnt;
182 	metadata = tld_obj->data_map->meta->metadata;
183 
184 	bpf_for(i, 0, cnt) {
185 		if (i >= TLD_MAX_DATA_CNT)
186 			break;
187 
188 		if (i >= i_start && !bpf_strncmp(metadata[i].name, TLD_NAME_LEN, name))
189 			return start + off;
190 
191 		off += TLD_ROUND_UP(metadata[i].size, 8);
192 	}
193 
194 	return -cnt;
195 }
196 
197 /**
198  * tld_get_data() - Retrieve a pointer to the TLD associated with the name.
199  *
200  * @tld_obj: A pointer to a valid tld_object initialized by tld_object_init()
201  * @key: The cached key of the TLD in tld_key_map
202  * @name: The name of the key associated with a TLD
203  * @size: The size of the TLD. Must be a known constant value
204  *
205  * Return a pointer to the TLD associated with @name; NULL if not found or @size is too
206  * big. @key is used to cache the key if the TLD is found to speed up subsequent calls.
207  * It should be defined as an member of tld_keys of tld_key_t type by the developer.
208  */
209 #define tld_get_data(tld_obj, key, name, size)						\
210 	({										\
211 		void *data = NULL, *_data = (tld_obj)->data_map->data;			\
212 		long off = (tld_obj)->key_map->key.off;					\
213 		int cnt;								\
214 											\
215 		if (likely(_data)) {							\
216 			if (likely(off > 0)) {						\
217 				barrier_var(off);					\
218 				if (likely(off < __PAGE_SIZE - size))			\
219 					data = _data + off;				\
220 			} else {							\
221 				cnt = -(off);						\
222 				if (likely((tld_obj)->data_map->meta) &&		\
223 				    cnt < (tld_obj)->data_map->meta->cnt) {		\
224 					off = __tld_fetch_key(tld_obj, name, cnt);	\
225 					(tld_obj)->key_map->key.off = off;		\
226 											\
227 					if (likely(off < __PAGE_SIZE - size)) {		\
228 						barrier_var(off);			\
229 						if (off > 0)				\
230 							data = _data + off;		\
231 					}						\
232 				}							\
233 			}								\
234 		}									\
235 		data;									\
236 	})
237 
238 #endif
239