task_local_data.bpf.h

  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
 73typedef struct {
 74	__s16 off;
 75} tld_key_t;
 76
 77struct tld_metadata {
 78	char name[TLD_NAME_LEN];
 79	__u16 size;
 80};
 81
 82struct tld_meta_u {
 83	__u8 cnt;
 84	__u16 size;
 85	struct tld_metadata metadata[TLD_MAX_DATA_CNT];
 86};
 87
 88struct tld_data_u {
 89	__u64 start; /* offset of tld_data_u->data in a page */
 90	char data[__PAGE_SIZE - sizeof(__u64)];
 91};
 92
 93struct tld_map_value {
 94	struct tld_data_u __uptr *data;
 95	struct tld_meta_u __uptr *meta;
 96};
 97
 98typedef struct tld_uptr_dummy {
 99	struct tld_data_u data[0];
100	struct tld_meta_u meta[0];
101} *tld_uptr_dummy_t;
102
103struct 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 */
118struct tld_keys;
119
120struct {
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
127struct {
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))
145static 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))
171static 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