bpf: Enable task local storage for tracing programs

To access per-task data, BPF programs usually creates a hash table with pid as the key. This is not ideal because: 1. The user need to estimate the proper size of the hash table, which may be inaccurate; 2. Big hash tables are slow; 3. To clean up the data properly during task terminations, the user need to write extra logic. Task local storage overcomes these issues and offers a better option for these per-task data. Task local storage is only available to BPF_LSM. Now enable it for tracing programs. Unlike LSM programs, tracing programs can be called in IRQ contexts. Helpers that access task local storage are updated to use raw_spin_lock_irqsave() instead of raw_spin_lock_bh(). Tracing programs can attach to functions on the task free path, e.g. exit_creds(). To avoid allocating task local storage after bpf_task_storage_free(). bpf_task_storage_get() is updated to not allocate new storage when the task is not refcounted (task->usage == 0).

Song Liu committed 5 years ago

Commit a10787e1 parent 9c8f21e10 files · +51 −72Browse files

Showing the changes to kernel (6 of 10 files). Show all

Unified · Split

6 files

kernel/bpf/Makefile +1 −2
kernel/bpf/bpf_local_storage.c +17 −11
kernel/bpf/bpf_lsm.c +0 −4
kernel/bpf/bpf_task_storage.c +11 −32
kernel/fork.c +5 −0
kernel/trace/bpf_trace.c +4 −0

kernel/bpf/Makefile+1−2

 obj-$(CONFIG_BPF_SYSCALL) += syscall.o verifier.o inode.o helpers.o tnum.o bpf_iter.o map_iter.o task_iter.o prog_iter.o
 obj-$(CONFIG_BPF_SYSCALL) += hashtab.o arraymap.o percpu_freelist.o bpf_lru_list.o lpm_trie.o map_in_map.o
 obj-$(CONFIG_BPF_SYSCALL) += local_storage.o queue_stack_maps.o ringbuf.o
+obj-$(CONFIG_BPF_SYSCALL) += bpf_local_storage.o bpf_task_storage.o
 obj-${CONFIG_BPF_LSM}	  += bpf_inode_storage.o
-obj-${CONFIG_BPF_LSM}	  += bpf_task_storage.o
 obj-$(CONFIG_BPF_SYSCALL) += disasm.o
 obj-$(CONFIG_BPF_JIT) += trampoline.o
 obj-$(CONFIG_BPF_SYSCALL) += btf.o
 ifeq ($(CONFIG_NET),y)
 obj-$(CONFIG_BPF_SYSCALL) += devmap.o
 obj-$(CONFIG_BPF_SYSCALL) += cpumap.o
-obj-$(CONFIG_BPF_SYSCALL) += bpf_local_storage.o
 obj-$(CONFIG_BPF_SYSCALL) += offload.o
 obj-$(CONFIG_BPF_SYSCALL) += net_namespace.o
 endif

kernel/bpf/bpf_local_storage.c+17−11

 {
 	struct bpf_local_storage *local_storage;
 	bool free_local_storage = false;
+	unsigned long flags;
 	if (unlikely(!selem_linked_to_storage(selem)))
 		/* selem has already been unlinked from sk */
 		return;
 	local_storage = rcu_dereference(selem->local_storage);
-	raw_spin_lock_bh(&local_storage->lock);
+	raw_spin_lock_irqsave(&local_storage->lock, flags);
 	if (likely(selem_linked_to_storage(selem)))
 		free_local_storage = bpf_selem_unlink_storage_nolock(
 			local_storage, selem, true);
-	raw_spin_unlock_bh(&local_storage->lock);
+	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 	if (free_local_storage)
 		kfree_rcu(local_storage, rcu);
 {
 	struct bpf_local_storage_map *smap;
 	struct bpf_local_storage_map_bucket *b;
+	unsigned long flags;
 	if (unlikely(!selem_linked_to_map(selem)))
 		/* selem has already be unlinked from smap */
 	smap = rcu_dereference(SDATA(selem)->smap);
 	b = select_bucket(smap, selem);
-	raw_spin_lock_bh(&b->lock);
+	raw_spin_lock_irqsave(&b->lock, flags);
 	if (likely(selem_linked_to_map(selem)))
 		hlist_del_init_rcu(&selem->map_node);
-	raw_spin_unlock_bh(&b->lock);
+	raw_spin_unlock_irqrestore(&b->lock, flags);
 }
 void bpf_selem_link_map(struct bpf_local_storage_map *smap,
 			struct bpf_local_storage_elem *selem)
 {
 	struct bpf_local_storage_map_bucket *b = select_bucket(smap, selem);
+	unsigned long flags;
-	raw_spin_lock_bh(&b->lock);
+	raw_spin_lock_irqsave(&b->lock, flags);
 	RCU_INIT_POINTER(SDATA(selem)->smap, smap);
 	hlist_add_head_rcu(&selem->map_node, &b->list);
-	raw_spin_unlock_bh(&b->lock);
+	raw_spin_unlock_irqrestore(&b->lock, flags);
 }
 void bpf_selem_unlink(struct bpf_local_storage_elem *selem)
 	sdata = SDATA(selem);
 	if (cacheit_lockit) {
+		unsigned long flags;
 		/* spinlock is needed to avoid racing with the
 		 * parallel delete.  Otherwise, publishing an already
 		 * deleted sdata to the cache will become a use-after-free
 		 * problem in the next bpf_local_storage_lookup().
 		 */
-		raw_spin_lock_bh(&local_storage->lock);
+		raw_spin_lock_irqsave(&local_storage->lock, flags);
 		if (selem_linked_to_storage(selem))
 			rcu_assign_pointer(local_storage->cache[smap->cache_idx],
 					   sdata);
-		raw_spin_unlock_bh(&local_storage->lock);
+		raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 	}
 	return sdata;
 	struct bpf_local_storage_data *old_sdata = NULL;
 	struct bpf_local_storage_elem *selem;
 	struct bpf_local_storage *local_storage;
+	unsigned long flags;
 	int err;
 	/* BPF_EXIST and BPF_NOEXIST cannot be both set */
 		}
 	}
-	raw_spin_lock_bh(&local_storage->lock);
+	raw_spin_lock_irqsave(&local_storage->lock, flags);
 	/* Recheck local_storage->list under local_storage->lock */
 	if (unlikely(hlist_empty(&local_storage->list))) {
 	}
 unlock:
-	raw_spin_unlock_bh(&local_storage->lock);
+	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 	return SDATA(selem);
 unlock_err:
-	raw_spin_unlock_bh(&local_storage->lock);
+	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 	return ERR_PTR(err);
 }

kernel/bpf/bpf_lsm.c+0−4

 		return &bpf_spin_lock_proto;
 	case BPF_FUNC_spin_unlock:
 		return &bpf_spin_unlock_proto;
-	case BPF_FUNC_task_storage_get:
-		return &bpf_task_storage_get_proto;
-	case BPF_FUNC_task_storage_delete:
-		return &bpf_task_storage_delete_proto;
 	case BPF_FUNC_bprm_opts_set:
 		return &bpf_bprm_opts_set_proto;
 	case BPF_FUNC_ima_inode_hash:

kernel/bpf/bpf_task_storage.c+11−32

@@ -15,7 +15,6 @@ · 14 unmodified lines
15	#include <linux/bpf_local_storage.h>	15	#include <linux/bpf_local_storage.h>
16	#include <linux/filter.h>	16	#include <linux/filter.h>
17	#include <uapi/linux/btf.h>	17	#include <uapi/linux/btf.h>
18	#include <linux/bpf_lsm.h>
19	#include <linux/btf_ids.h>	18	#include <linux/btf_ids.h>
20	#include <linux/fdtable.h>	19	#include <linux/fdtable.h>
21		20
@@ -24,12 +23,8 @@ DEFINE_BPF_STORAGE_CACHE(task_cache); · 2 unmodified lines
24	static struct bpf_local_storage __rcu *task_storage_ptr(void owner)	23	static struct bpf_local_storage __rcu *task_storage_ptr(void owner)
25	{	24	{
26	struct task_struct *task = owner;	25	struct task_struct *task = owner;
27	struct bpf_storage_blob *bsb;
28		26
29	bsb = bpf_task(task);	27	return &task->bpf_storage;
30	if (!bsb)
31	return NULL;
32	return &bsb->storage;
33	}	28	}
34		29
35	static struct bpf_local_storage_data *	30	static struct bpf_local_storage_data *
@@ -38,13 +33,8 @@ task_storage_lookup(struct task_struct task, struct bpf_map map, · 2 unmodified lines
38	{	33	{
39	struct bpf_local_storage *task_storage;	34	struct bpf_local_storage *task_storage;
40	struct bpf_local_storage_map *smap;	35	struct bpf_local_storage_map *smap;
41	struct bpf_storage_blob *bsb;
42
43	bsb = bpf_task(task);
44	if (!bsb)
45	return NULL;
46		36
47	task_storage = rcu_dereference(bsb->storage);	37	task_storage = rcu_dereference(task->bpf_storage);
48	if (!task_storage)	38	if (!task_storage)
49	return NULL;	39	return NULL;
50		40
@@ -57,16 +47,12 @@ void bpf_task_storage_free(struct task_struct *task) · 6 unmodified lines
57	struct bpf_local_storage_elem *selem;	47	struct bpf_local_storage_elem *selem;
58	struct bpf_local_storage *local_storage;	48	struct bpf_local_storage *local_storage;
59	bool free_task_storage = false;	49	bool free_task_storage = false;
60	struct bpf_storage_blob *bsb;
61	struct hlist_node *n;	50	struct hlist_node *n;
62		51	unsigned long flags;
63	bsb = bpf_task(task);
64	if (!bsb)
65	return;
66		52
67	rcu_read_lock();	53	rcu_read_lock();
68		54
69	local_storage = rcu_dereference(bsb->storage);	55	local_storage = rcu_dereference(task->bpf_storage);
70	if (!local_storage) {	56	if (!local_storage) {
71	rcu_read_unlock();	57	rcu_read_unlock();
72	return;	58	return;
@@ -81,7 +67,7 @@ void bpf_task_storage_free(struct task_struct *task) · 8 unmodified lines
81	* when unlinking elem from the local_storage->list and	67	* when unlinking elem from the local_storage->list and
82	* the map's bucket->list.	68	* the map's bucket->list.
83	*/	69	*/
84	raw_spin_lock_bh(&local_storage->lock);	70	raw_spin_lock_irqsave(&local_storage->lock, flags);
85	hlist_for_each_entry_safe(selem, n, &local_storage->list, snode) {	71	hlist_for_each_entry_safe(selem, n, &local_storage->list, snode) {
86	/* Always unlink from map before unlinking from	72	/* Always unlink from map before unlinking from
87	* local_storage.	73	* local_storage.
@@ -90,7 +76,7 @@ void bpf_task_storage_free(struct task_struct *task) · 2 unmodified lines
90	free_task_storage = bpf_selem_unlink_storage_nolock(	76	free_task_storage = bpf_selem_unlink_storage_nolock(
91	local_storage, selem, false);	77	local_storage, selem, false);
92	}	78	}
93	raw_spin_unlock_bh(&local_storage->lock);	79	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
94	rcu_read_unlock();	80	rcu_read_unlock();
95		81
96	/* free_task_storage should always be true as long as	82	/* free_task_storage should always be true as long as
@@ -150,7 +136,7 @@ static int bpf_pid_task_storage_update_elem(struct bpf_map map, void key, · 53 unmodified lines
150	*/	136	*/
151	WARN_ON_ONCE(!rcu_read_lock_held());	137	WARN_ON_ONCE(!rcu_read_lock_held());
152	task = pid_task(pid, PIDTYPE_PID);	138	task = pid_task(pid, PIDTYPE_PID);
153	if (!task \|\| !task_storage_ptr(task)) {	139	if (!task) {
154	err = -ENOENT;	140	err = -ENOENT;
155	goto out;	141	goto out;
156	}	142	}
@@ -213,23 +199,16 @@ BPF_CALL_4(bpf_task_storage_get, struct bpf_map , map, struct task_struct , · 56 unmodified lines
213	if (flags & ~(BPF_LOCAL_STORAGE_GET_F_CREATE))	199	if (flags & ~(BPF_LOCAL_STORAGE_GET_F_CREATE))
214	return (unsigned long)NULL;	200	return (unsigned long)NULL;
215		201
216	/* explicitly check that the task_storage_ptr is not	202	if (!task)
217	* NULL as task_storage_lookup returns NULL in this case and
218	* bpf_local_storage_update expects the owner to have a
219	* valid storage pointer.
220	*/
221	if (!task \|\| !task_storage_ptr(task))
222	return (unsigned long)NULL;	203	return (unsigned long)NULL;
223		204
224	sdata = task_storage_lookup(task, map, true);	205	sdata = task_storage_lookup(task, map, true);
225	if (sdata)	206	if (sdata)
226	return (unsigned long)sdata->data;	207	return (unsigned long)sdata->data;
227		208
228	/* This helper must only be called from places where the lifetime of the task	209	/* only allocate new storage, when the task is refcounted */
229	* is guaranteed. Either by being refcounted or by being protected	210	if (refcount_read(&task->usage) &&
230	* by an RCU read-side critical section.	211	(flags & BPF_LOCAL_STORAGE_GET_F_CREATE)) {
231	*/
232	if (flags & BPF_LOCAL_STORAGE_GET_F_CREATE) {
233	sdata = bpf_local_storage_update(	212	sdata = bpf_local_storage_update(
234	task, (struct bpf_local_storage_map *)map, value,	213	task, (struct bpf_local_storage_map *)map, value,
235	BPF_NOEXIST);	214	BPF_NOEXIST);

kernel/fork.c+5−0

 #include <linux/kasan.h>
 #include <linux/scs.h>
 #include <linux/io_uring.h>
+#include <linux/bpf.h>
 #include <asm/pgalloc.h>
 #include <linux/uaccess.h>
 	cgroup_free(tsk);
 	task_numa_free(tsk, true);
 	security_task_free(tsk);
+	bpf_task_storage_free(tsk);
 	exit_creds(tsk);
 	delayacct_tsk_free(tsk);
 	put_signal_struct(tsk->signal);
 	p->sequential_io	= 0;
 	p->sequential_io_avg	= 0;
 #endif
+#ifdef CONFIG_BPF_SYSCALL
+	RCU_INIT_POINTER(p->bpf_storage, NULL);
+#endif
 	/* Perform scheduler related setup. Assign this task to a CPU. */
 	retval = sched_fork(clone_flags, p);

kernel/trace/bpf_trace.c+4−0

 		return &bpf_per_cpu_ptr_proto;
 	case BPF_FUNC_this_cpu_ptr:
 		return &bpf_this_cpu_ptr_proto;
+	case BPF_FUNC_task_storage_get:
+		return &bpf_task_storage_get_proto;
+	case BPF_FUNC_task_storage_delete:
+		return &bpf_task_storage_delete_proto;
 	default:
 		return NULL;
 	}

bpf: Enable task local storage for tracing programs

Keyboard shortcuts

General

Issues and pull requests

Issues, pull requests and Inbox

Repository code