bpf_task_storage.c

  1// SPDX-License-Identifier: GPL-2.0
  2/*
  3 * Copyright (c) 2020 Facebook
  4 * Copyright 2020 Google LLC.
  5 */
  6
  7#include <linux/pid.h>
  8#include <linux/sched.h>
  9#include <linux/rculist.h>
 10#include <linux/list.h>
 11#include <linux/hash.h>
 12#include <linux/types.h>
 13#include <linux/spinlock.h>
 14#include <linux/bpf.h>
 15#include <linux/bpf_local_storage.h>
 16#include <linux/filter.h>
 17#include <uapi/linux/btf.h>
 18#include <linux/btf_ids.h>
 19#include <linux/rcupdate_trace.h>
 20
 21DEFINE_BPF_STORAGE_CACHE(task_cache);
 22
 23static DEFINE_PER_CPU(int, bpf_task_storage_busy);
 24
 25static void bpf_task_storage_lock(void)
 26{
 27	cant_migrate();
 28	this_cpu_inc(bpf_task_storage_busy);
 29}
 30
 31static void bpf_task_storage_unlock(void)
 32{
 33	this_cpu_dec(bpf_task_storage_busy);
 34}
 35
 36static bool bpf_task_storage_trylock(void)
 37{
 38	cant_migrate();
 39	if (unlikely(this_cpu_inc_return(bpf_task_storage_busy) != 1)) {
 40		this_cpu_dec(bpf_task_storage_busy);
 41		return false;
 42	}
 43	return true;
 44}
 45
 46static struct bpf_local_storage __rcu **task_storage_ptr(void *owner)
 47{
 48	struct task_struct *task = owner;
 49
 50	return &task->bpf_storage;
 51}
 52
 53static struct bpf_local_storage_data *
 54task_storage_lookup(struct task_struct *task, struct bpf_map *map,
 55		    bool cacheit_lockit)
 56{
 57	struct bpf_local_storage *task_storage;
 58	struct bpf_local_storage_map *smap;
 59
 60	task_storage =
 61		rcu_dereference_check(task->bpf_storage, bpf_rcu_lock_held());
 62	if (!task_storage)
 63		return NULL;
 64
 65	smap = (struct bpf_local_storage_map *)map;
 66	return bpf_local_storage_lookup(task_storage, smap, cacheit_lockit);
 67}
 68
 69void bpf_task_storage_free(struct task_struct *task)
 70{
 71	struct bpf_local_storage *local_storage;
 72
 73	rcu_read_lock_dont_migrate();
 74
 75	local_storage = rcu_dereference(task->bpf_storage);
 76	if (!local_storage)
 77		goto out;
 78
 79	bpf_task_storage_lock();
 80	bpf_local_storage_destroy(local_storage);
 81	bpf_task_storage_unlock();
 82out:
 83	rcu_read_unlock_migrate();
 84}
 85
 86static void *bpf_pid_task_storage_lookup_elem(struct bpf_map *map, void *key)
 87{
 88	struct bpf_local_storage_data *sdata;
 89	struct task_struct *task;
 90	unsigned int f_flags;
 91	struct pid *pid;
 92	int fd, err;
 93
 94	fd = *(int *)key;
 95	pid = pidfd_get_pid(fd, &f_flags);
 96	if (IS_ERR(pid))
 97		return ERR_CAST(pid);
 98
 99	/* We should be in an RCU read side critical section, it should be safe
100	 * to call pid_task.
101	 */
102	WARN_ON_ONCE(!rcu_read_lock_held());
103	task = pid_task(pid, PIDTYPE_PID);
104	if (!task) {
105		err = -ENOENT;
106		goto out;
107	}
108
109	bpf_task_storage_lock();
110	sdata = task_storage_lookup(task, map, true);
111	bpf_task_storage_unlock();
112	put_pid(pid);
113	return sdata ? sdata->data : NULL;
114out:
115	put_pid(pid);
116	return ERR_PTR(err);
117}
118
119static long bpf_pid_task_storage_update_elem(struct bpf_map *map, void *key,
120					     void *value, u64 map_flags)
121{
122	struct bpf_local_storage_data *sdata;
123	struct task_struct *task;
124	unsigned int f_flags;
125	struct pid *pid;
126	int fd, err;
127
128	if ((map_flags & BPF_F_LOCK) && btf_record_has_field(map->record, BPF_UPTR))
129		return -EOPNOTSUPP;
130
131	fd = *(int *)key;
132	pid = pidfd_get_pid(fd, &f_flags);
133	if (IS_ERR(pid))
134		return PTR_ERR(pid);
135
136	/* We should be in an RCU read side critical section, it should be safe
137	 * to call pid_task.
138	 */
139	WARN_ON_ONCE(!rcu_read_lock_held());
140	task = pid_task(pid, PIDTYPE_PID);
141	if (!task) {
142		err = -ENOENT;
143		goto out;
144	}
145
146	bpf_task_storage_lock();
147	sdata = bpf_local_storage_update(
148		task, (struct bpf_local_storage_map *)map, value, map_flags,
149		true, GFP_ATOMIC);
150	bpf_task_storage_unlock();
151
152	err = PTR_ERR_OR_ZERO(sdata);
153out:
154	put_pid(pid);
155	return err;
156}
157
158static int task_storage_delete(struct task_struct *task, struct bpf_map *map,
159			       bool nobusy)
160{
161	struct bpf_local_storage_data *sdata;
162
163	sdata = task_storage_lookup(task, map, false);
164	if (!sdata)
165		return -ENOENT;
166
167	if (!nobusy)
168		return -EBUSY;
169
170	bpf_selem_unlink(SELEM(sdata), false);
171
172	return 0;
173}
174
175static long bpf_pid_task_storage_delete_elem(struct bpf_map *map, void *key)
176{
177	struct task_struct *task;
178	unsigned int f_flags;
179	struct pid *pid;
180	int fd, err;
181
182	fd = *(int *)key;
183	pid = pidfd_get_pid(fd, &f_flags);
184	if (IS_ERR(pid))
185		return PTR_ERR(pid);
186
187	/* We should be in an RCU read side critical section, it should be safe
188	 * to call pid_task.
189	 */
190	WARN_ON_ONCE(!rcu_read_lock_held());
191	task = pid_task(pid, PIDTYPE_PID);
192	if (!task) {
193		err = -ENOENT;
194		goto out;
195	}
196
197	bpf_task_storage_lock();
198	err = task_storage_delete(task, map, true);
199	bpf_task_storage_unlock();
200out:
201	put_pid(pid);
202	return err;
203}
204
205/* Called by bpf_task_storage_get*() helpers */
206static void *__bpf_task_storage_get(struct bpf_map *map,
207				    struct task_struct *task, void *value,
208				    u64 flags, gfp_t gfp_flags, bool nobusy)
209{
210	struct bpf_local_storage_data *sdata;
211
212	sdata = task_storage_lookup(task, map, nobusy);
213	if (sdata)
214		return sdata->data;
215
216	/* only allocate new storage, when the task is refcounted */
217	if (refcount_read(&task->usage) &&
218	    (flags & BPF_LOCAL_STORAGE_GET_F_CREATE) && nobusy) {
219		sdata = bpf_local_storage_update(
220			task, (struct bpf_local_storage_map *)map, value,
221			BPF_NOEXIST, false, gfp_flags);
222		return IS_ERR(sdata) ? NULL : sdata->data;
223	}
224
225	return NULL;
226}
227
228/* *gfp_flags* is a hidden argument provided by the verifier */
229BPF_CALL_5(bpf_task_storage_get_recur, struct bpf_map *, map, struct task_struct *,
230	   task, void *, value, u64, flags, gfp_t, gfp_flags)
231{
232	bool nobusy;
233	void *data;
234
235	WARN_ON_ONCE(!bpf_rcu_lock_held());
236	if (flags & ~BPF_LOCAL_STORAGE_GET_F_CREATE || !task)
237		return (unsigned long)NULL;
238
239	nobusy = bpf_task_storage_trylock();
240	data = __bpf_task_storage_get(map, task, value, flags,
241				      gfp_flags, nobusy);
242	if (nobusy)
243		bpf_task_storage_unlock();
244	return (unsigned long)data;
245}
246
247/* *gfp_flags* is a hidden argument provided by the verifier */
248BPF_CALL_5(bpf_task_storage_get, struct bpf_map *, map, struct task_struct *,
249	   task, void *, value, u64, flags, gfp_t, gfp_flags)
250{
251	void *data;
252
253	WARN_ON_ONCE(!bpf_rcu_lock_held());
254	if (flags & ~BPF_LOCAL_STORAGE_GET_F_CREATE || !task)
255		return (unsigned long)NULL;
256
257	bpf_task_storage_lock();
258	data = __bpf_task_storage_get(map, task, value, flags,
259				      gfp_flags, true);
260	bpf_task_storage_unlock();
261	return (unsigned long)data;
262}
263
264BPF_CALL_2(bpf_task_storage_delete_recur, struct bpf_map *, map, struct task_struct *,
265	   task)
266{
267	bool nobusy;
268	int ret;
269
270	WARN_ON_ONCE(!bpf_rcu_lock_held());
271	if (!task)
272		return -EINVAL;
273
274	nobusy = bpf_task_storage_trylock();
275	/* This helper must only be called from places where the lifetime of the task
276	 * is guaranteed. Either by being refcounted or by being protected
277	 * by an RCU read-side critical section.
278	 */
279	ret = task_storage_delete(task, map, nobusy);
280	if (nobusy)
281		bpf_task_storage_unlock();
282	return ret;
283}
284
285BPF_CALL_2(bpf_task_storage_delete, struct bpf_map *, map, struct task_struct *,
286	   task)
287{
288	int ret;
289
290	WARN_ON_ONCE(!bpf_rcu_lock_held());
291	if (!task)
292		return -EINVAL;
293
294	bpf_task_storage_lock();
295	/* This helper must only be called from places where the lifetime of the task
296	 * is guaranteed. Either by being refcounted or by being protected
297	 * by an RCU read-side critical section.
298	 */
299	ret = task_storage_delete(task, map, true);
300	bpf_task_storage_unlock();
301	return ret;
302}
303
304static int notsupp_get_next_key(struct bpf_map *map, void *key, void *next_key)
305{
306	return -ENOTSUPP;
307}
308
309static struct bpf_map *task_storage_map_alloc(union bpf_attr *attr)
310{
311	return bpf_local_storage_map_alloc(attr, &task_cache, true);
312}
313
314static void task_storage_map_free(struct bpf_map *map)
315{
316	bpf_local_storage_map_free(map, &task_cache, &bpf_task_storage_busy);
317}
318
319BTF_ID_LIST_GLOBAL_SINGLE(bpf_local_storage_map_btf_id, struct, bpf_local_storage_map)
320const struct bpf_map_ops task_storage_map_ops = {
321	.map_meta_equal = bpf_map_meta_equal,
322	.map_alloc_check = bpf_local_storage_map_alloc_check,
323	.map_alloc = task_storage_map_alloc,
324	.map_free = task_storage_map_free,
325	.map_get_next_key = notsupp_get_next_key,
326	.map_lookup_elem = bpf_pid_task_storage_lookup_elem,
327	.map_update_elem = bpf_pid_task_storage_update_elem,
328	.map_delete_elem = bpf_pid_task_storage_delete_elem,
329	.map_check_btf = bpf_local_storage_map_check_btf,
330	.map_mem_usage = bpf_local_storage_map_mem_usage,
331	.map_btf_id = &bpf_local_storage_map_btf_id[0],
332	.map_owner_storage_ptr = task_storage_ptr,
333};
334
335const struct bpf_func_proto bpf_task_storage_get_recur_proto = {
336	.func = bpf_task_storage_get_recur,
337	.gpl_only = false,
338	.ret_type = RET_PTR_TO_MAP_VALUE_OR_NULL,
339	.arg1_type = ARG_CONST_MAP_PTR,
340	.arg2_type = ARG_PTR_TO_BTF_ID_OR_NULL,
341	.arg2_btf_id = &btf_tracing_ids[BTF_TRACING_TYPE_TASK],
342	.arg3_type = ARG_PTR_TO_MAP_VALUE_OR_NULL,
343	.arg4_type = ARG_ANYTHING,
344};
345
346const struct bpf_func_proto bpf_task_storage_get_proto = {
347	.func = bpf_task_storage_get,
348	.gpl_only = false,
349	.ret_type = RET_PTR_TO_MAP_VALUE_OR_NULL,
350	.arg1_type = ARG_CONST_MAP_PTR,
351	.arg2_type = ARG_PTR_TO_BTF_ID_OR_NULL,
352	.arg2_btf_id = &btf_tracing_ids[BTF_TRACING_TYPE_TASK],
353	.arg3_type = ARG_PTR_TO_MAP_VALUE_OR_NULL,
354	.arg4_type = ARG_ANYTHING,
355};
356
357const struct bpf_func_proto bpf_task_storage_delete_recur_proto = {
358	.func = bpf_task_storage_delete_recur,
359	.gpl_only = false,
360	.ret_type = RET_INTEGER,
361	.arg1_type = ARG_CONST_MAP_PTR,
362	.arg2_type = ARG_PTR_TO_BTF_ID_OR_NULL,
363	.arg2_btf_id = &btf_tracing_ids[BTF_TRACING_TYPE_TASK],
364};
365
366const struct bpf_func_proto bpf_task_storage_delete_proto = {
367	.func = bpf_task_storage_delete,
368	.gpl_only = false,
369	.ret_type = RET_INTEGER,
370	.arg1_type = ARG_CONST_MAP_PTR,
371	.arg2_type = ARG_PTR_TO_BTF_ID_OR_NULL,
372	.arg2_btf_id = &btf_tracing_ids[BTF_TRACING_TYPE_TASK],
373};