diff mbox series

[v4,bpf-next,1/6] bpf: enable task local storage for tracing programs

Message ID 20210223012014.2087583-2-songliubraving@fb.com (mailing list archive)
State Superseded
Delegated to: BPF
Headers show
Series bpf: enable task local storage for tracing programs | expand

Checks

Context Check Description
netdev/cover_letter success Link
netdev/fixes_present success Link
netdev/patch_count success Link
netdev/tree_selection success Clearly marked for bpf-next
netdev/subject_prefix success Link
netdev/cc_maintainers warning 24 maintainers not CCed: bristot@redhat.com elver@google.com dietmar.eggemann@arm.com ebiederm@xmission.com shakeelb@google.com surenb@google.com keescook@chromium.org juri.lelli@redhat.com krisman@collabora.com mgorman@suse.de yhs@fb.com jackmanb@chromium.org vincent.guittot@linaro.org akpm@linux-foundation.org esyr@redhat.com revest@chromium.org john.fastabend@gmail.com bsegall@google.com andrii@kernel.org mingo@redhat.com christian.brauner@ubuntu.com rostedt@goodmis.org kafai@fb.com peterx@redhat.com
netdev/source_inline success Was 0 now: 0
netdev/verify_signedoff success Link
netdev/module_param success Was 0 now: 0
netdev/build_32bit success Errors and warnings before: 23506 this patch: 23506
netdev/kdoc success Errors and warnings before: 5 this patch: 5
netdev/verify_fixes success Link
netdev/checkpatch success total: 0 errors, 0 warnings, 0 checks, 328 lines checked
netdev/build_allmodconfig_warn success Errors and warnings before: 23283 this patch: 23283
netdev/header_inline success Link
netdev/stable success Stable not CCed

Commit Message

Song Liu Feb. 23, 2021, 1:20 a.m. UTC 
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).

Acked-by: KP Singh <kpsingh@kernel.org>
Signed-off-by: Song Liu <songliubraving@fb.com>
---
 include/linux/bpf.h            |  7 +++++++
 include/linux/bpf_lsm.h        | 22 ----------------------
 include/linux/bpf_types.h      |  2 +-
 include/linux/sched.h          |  5 +++++
 kernel/bpf/Makefile            |  3 +--
 kernel/bpf/bpf_local_storage.c | 28 +++++++++++++++++-----------
 kernel/bpf/bpf_lsm.c           |  4 ----
 kernel/bpf/bpf_task_storage.c  | 34 +++++++++-------------------------
 kernel/fork.c                  |  5 +++++
 kernel/trace/bpf_trace.c       |  4 ++++
 10 files changed, 49 insertions(+), 65 deletions(-)

Comments

kernel test robot Feb. 23, 2021, 3:08 a.m. UTC | #1 
Hi Song,

I love your patch! Yet something to improve:

[auto build test ERROR on bpf-next/master]

url:    https://github.com/0day-ci/linux/commits/Song-Liu/bpf-enable-task-local-storage-for-tracing-programs/20210223-092439
base:   https://git.kernel.org/pub/scm/linux/kernel/git/bpf/bpf-next.git master
config: m68k-randconfig-r031-20210222 (attached as .config)
compiler: m68k-linux-gcc (GCC) 9.3.0
reproduce (this is a W=1 build):
        wget https://raw.githubusercontent.com/intel/lkp-tests/master/sbin/make.cross -O ~/bin/make.cross
        chmod +x ~/bin/make.cross
        # https://github.com/0day-ci/linux/commit/a1afc383ea0cda5b0221b111f208140114752782
        git remote add linux-review https://github.com/0day-ci/linux
        git fetch --no-tags linux-review Song-Liu/bpf-enable-task-local-storage-for-tracing-programs/20210223-092439
        git checkout a1afc383ea0cda5b0221b111f208140114752782
        # save the attached .config to linux build tree
        COMPILER_INSTALL_PATH=$HOME/0day COMPILER=gcc-9.3.0 make.cross ARCH=m68k 

If you fix the issue, kindly add following tag as appropriate
Reported-by: kernel test robot <lkp@intel.com>

All errors (new ones prefixed by >>):

   m68k-linux-ld: kernel/bpf/bpf_task_storage.o: in function `task_storage_delete':
>> bpf_task_storage.c:(.text+0x28): undefined reference to `bpf_local_storage_lookup'
>> m68k-linux-ld: bpf_task_storage.c:(.text+0x3a): undefined reference to `bpf_selem_unlink'
   m68k-linux-ld: kernel/bpf/bpf_task_storage.o: in function `bpf_pid_task_storage_update_elem':
>> bpf_task_storage.c:(.text+0x120): undefined reference to `bpf_local_storage_update'
   m68k-linux-ld: kernel/bpf/bpf_task_storage.o: in function `bpf_pid_task_storage_lookup_elem':
   bpf_task_storage.c:(.text+0x1a8): undefined reference to `bpf_local_storage_lookup'
   m68k-linux-ld: kernel/bpf/bpf_task_storage.o: in function `task_storage_map_free':
>> bpf_task_storage.c:(.text+0x232): undefined reference to `bpf_local_storage_cache_idx_free'
>> m68k-linux-ld: bpf_task_storage.c:(.text+0x240): undefined reference to `bpf_local_storage_map_free'
   m68k-linux-ld: kernel/bpf/bpf_task_storage.o: in function `task_storage_map_alloc':
>> bpf_task_storage.c:(.text+0x24c): undefined reference to `bpf_local_storage_map_alloc'
>> m68k-linux-ld: bpf_task_storage.c:(.text+0x262): undefined reference to `bpf_local_storage_cache_idx_get'
   m68k-linux-ld: kernel/bpf/bpf_task_storage.o: in function `bpf_task_storage_get':
   bpf_task_storage.c:(.text+0x2b2): undefined reference to `bpf_local_storage_lookup'
>> m68k-linux-ld: bpf_task_storage.c:(.text+0x2ee): undefined reference to `bpf_local_storage_update'
   m68k-linux-ld: kernel/bpf/bpf_task_storage.o: in function `bpf_task_storage_free':
>> bpf_task_storage.c:(.text+0x328): undefined reference to `bpf_selem_unlink_map'
>> m68k-linux-ld: bpf_task_storage.c:(.text+0x32e): undefined reference to `bpf_selem_unlink_storage_nolock'
>> m68k-linux-ld: kernel/bpf/bpf_task_storage.o:(.rodata+0x74): undefined reference to `bpf_local_storage_map_alloc_check'
>> m68k-linux-ld: kernel/bpf/bpf_task_storage.o:(.rodata+0xcc): undefined reference to `bpf_local_storage_map_check_btf'
   m68k-linux-ld: drivers/macintosh/via-pmu.o: in function `pmu_set_rtc_time':
   (.text+0x187e): undefined reference to `rtc_tm_to_time64'

---
0-DAY CI Kernel Test Service, Intel Corporation
https://lists.01.org/hyperkitty/list/kbuild-all@lists.01.org
kernel test robot Feb. 23, 2021, 4:04 a.m. UTC | #2 
Hi Song,

I love your patch! Yet something to improve:

[auto build test ERROR on bpf-next/master]

url:    https://github.com/0day-ci/linux/commits/Song-Liu/bpf-enable-task-local-storage-for-tracing-programs/20210223-092439
base:   https://git.kernel.org/pub/scm/linux/kernel/git/bpf/bpf-next.git master
config: xtensa-randconfig-r034-20210222 (attached as .config)
compiler: xtensa-linux-gcc (GCC) 9.3.0
reproduce (this is a W=1 build):
        wget https://raw.githubusercontent.com/intel/lkp-tests/master/sbin/make.cross -O ~/bin/make.cross
        chmod +x ~/bin/make.cross
        # https://github.com/0day-ci/linux/commit/a1afc383ea0cda5b0221b111f208140114752782
        git remote add linux-review https://github.com/0day-ci/linux
        git fetch --no-tags linux-review Song-Liu/bpf-enable-task-local-storage-for-tracing-programs/20210223-092439
        git checkout a1afc383ea0cda5b0221b111f208140114752782
        # save the attached .config to linux build tree
        COMPILER_INSTALL_PATH=$HOME/0day COMPILER=gcc-9.3.0 make.cross ARCH=xtensa 

If you fix the issue, kindly add following tag as appropriate
Reported-by: kernel test robot <lkp@intel.com>

All errors (new ones prefixed by >>):

   xtensa-linux-ld: kernel/bpf/bpf_task_storage.o: in function `notsupp_get_next_key':
   bpf_task_storage.c:(.text+0x14): undefined reference to `bpf_local_storage_lookup'
>> xtensa-linux-ld: bpf_task_storage.c:(.text+0x18): undefined reference to `bpf_selem_unlink'
   xtensa-linux-ld: kernel/bpf/bpf_task_storage.o: in function `task_storage_delete':
   bpf_task_storage.c:(.text+0x32): undefined reference to `bpf_local_storage_lookup'
   xtensa-linux-ld: bpf_task_storage.c:(.text+0x3e): undefined reference to `bpf_selem_unlink'
   xtensa-linux-ld: kernel/bpf/bpf_task_storage.o: in function `bpf_pid_task_storage_delete_elem':
   bpf_task_storage.c:(.text+0x94): undefined reference to `bpf_local_storage_update'
   xtensa-linux-ld: kernel/bpf/bpf_task_storage.o: in function `bpf_pid_task_storage_update_elem':
   bpf_task_storage.c:(.text+0xce): undefined reference to `bpf_local_storage_update'
>> xtensa-linux-ld: bpf_task_storage.c:(.text+0x12c): undefined reference to `bpf_local_storage_lookup'
   xtensa-linux-ld: kernel/bpf/bpf_task_storage.o: in function `bpf_pid_task_storage_lookup_elem':
   bpf_task_storage.c:(.text+0x148): undefined reference to `bpf_local_storage_cache_idx_free'
>> xtensa-linux-ld: bpf_task_storage.c:(.text+0x14c): undefined reference to `bpf_local_storage_map_free'
>> xtensa-linux-ld: bpf_task_storage.c:(.text+0x159): undefined reference to `bpf_local_storage_cache_idx_free'
   xtensa-linux-ld: bpf_task_storage.c:(.text+0x162): undefined reference to `bpf_local_storage_map_free'
>> xtensa-linux-ld: bpf_task_storage.c:(.text+0x16c): undefined reference to `bpf_local_storage_map_alloc'
>> xtensa-linux-ld: bpf_task_storage.c:(.text+0x170): undefined reference to `bpf_local_storage_cache_idx_get'
   xtensa-linux-ld: bpf_task_storage.c:(.text+0x17a): undefined reference to `bpf_local_storage_map_alloc'
   xtensa-linux-ld: bpf_task_storage.c:(.text+0x18b): undefined reference to `bpf_local_storage_cache_idx_get'
   xtensa-linux-ld: kernel/bpf/bpf_task_storage.o: in function `task_storage_map_free':
   bpf_task_storage.c:(.text+0x1d3): undefined reference to `bpf_local_storage_lookup'
   xtensa-linux-ld: kernel/bpf/bpf_task_storage.o: in function `task_storage_map_alloc':
   bpf_task_storage.c:(.text+0x1f6): undefined reference to `bpf_local_storage_update'
   xtensa-linux-ld: bpf_task_storage.c:(.text+0x208): undefined reference to `bpf_selem_unlink_map'
>> xtensa-linux-ld: bpf_task_storage.c:(.text+0x20c): undefined reference to `bpf_selem_unlink_storage_nolock'
   xtensa-linux-ld: kernel/bpf/bpf_task_storage.o: in function `bpf_task_storage_get':
   bpf_task_storage.c:(.text+0x252): undefined reference to `bpf_selem_unlink_map'
   xtensa-linux-ld: bpf_task_storage.c:(.text+0x25e): undefined reference to `bpf_selem_unlink_storage_nolock'
>> xtensa-linux-ld: kernel/bpf/bpf_task_storage.o:(.rodata+0x78): undefined reference to `bpf_local_storage_map_alloc_check'
>> xtensa-linux-ld: kernel/bpf/bpf_task_storage.o:(.rodata+0xd0): undefined reference to `bpf_local_storage_map_check_btf'

---
0-DAY CI Kernel Test Service, Intel Corporation
https://lists.01.org/hyperkitty/list/kbuild-all@lists.01.org
Martin KaFai Lau Feb. 23, 2021, 7:23 p.m. UTC | #3 
On Mon, Feb 22, 2021 at 05:20:09PM -0800, Song Liu wrote:
[ ... ]

> diff --git a/kernel/bpf/bpf_task_storage.c b/kernel/bpf/bpf_task_storage.c
> index e0da0258b732d..2034019966d44 100644
> --- a/kernel/bpf/bpf_task_storage.c
> +++ b/kernel/bpf/bpf_task_storage.c
> @@ -15,7 +15,6 @@
>  #include <linux/bpf_local_storage.h>
>  #include <linux/filter.h>
>  #include <uapi/linux/btf.h>
> -#include <linux/bpf_lsm.h>
>  #include <linux/btf_ids.h>
>  #include <linux/fdtable.h>
>  
> @@ -24,12 +23,8 @@ DEFINE_BPF_STORAGE_CACHE(task_cache);
>  static struct bpf_local_storage __rcu **task_storage_ptr(void *owner)
>  {
>  	struct task_struct *task = owner;
> -	struct bpf_storage_blob *bsb;
>  
> -	bsb = bpf_task(task);
> -	if (!bsb)
> -		return NULL;
task_storage_ptr() no longer returns NULL.  All "!task_storage_ptr(task)"
checks should be removed also.  e.g. In bpf_task_storage_get
and bpf_pid_task_storage_update_elem.

> -	return &bsb->storage;
> +	return &task->bpf_storage;
>  }
>
Song Liu Feb. 23, 2021, 8:51 p.m. UTC | #4 
> On Feb 23, 2021, at 11:23 AM, Martin Lau <kafai@fb.com> wrote:
> 
> On Mon, Feb 22, 2021 at 05:20:09PM -0800, Song Liu wrote:
> [ ... ]
> 
>> diff --git a/kernel/bpf/bpf_task_storage.c b/kernel/bpf/bpf_task_storage.c
>> index e0da0258b732d..2034019966d44 100644
>> --- a/kernel/bpf/bpf_task_storage.c
>> +++ b/kernel/bpf/bpf_task_storage.c
>> @@ -15,7 +15,6 @@
>> #include <linux/bpf_local_storage.h>
>> #include <linux/filter.h>
>> #include <uapi/linux/btf.h>
>> -#include <linux/bpf_lsm.h>
>> #include <linux/btf_ids.h>
>> #include <linux/fdtable.h>
>> 
>> @@ -24,12 +23,8 @@ DEFINE_BPF_STORAGE_CACHE(task_cache);
>> static struct bpf_local_storage __rcu **task_storage_ptr(void *owner)
>> {
>> 	struct task_struct *task = owner;
>> -	struct bpf_storage_blob *bsb;
>> 
>> -	bsb = bpf_task(task);
>> -	if (!bsb)
>> -		return NULL;
> task_storage_ptr() no longer returns NULL.  All "!task_storage_ptr(task)"
> checks should be removed also.  e.g. In bpf_task_storage_get
> and bpf_pid_task_storage_update_elem.

Good catch! Fixed it in v5. 

Thanks,
Song

> 
>> -	return &bsb->storage;
>> +	return &task->bpf_storage;
>> }
>>
diff mbox series

Patch

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index cccaef1088eae..e2cfc4809219c 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1499,6 +1499,7 @@  struct bpf_prog *bpf_prog_by_id(u32 id);
 struct bpf_link *bpf_link_by_id(u32 id);
 
 const struct bpf_func_proto *bpf_base_func_proto(enum bpf_func_id func_id);
+void bpf_task_storage_free(struct task_struct *task);
 #else /* !CONFIG_BPF_SYSCALL */
 static inline struct bpf_prog *bpf_prog_get(u32 ufd)
 {
@@ -1684,6 +1685,10 @@  bpf_base_func_proto(enum bpf_func_id func_id)
 {
 	return NULL;
 }
+
+static inline void bpf_task_storage_free(struct task_struct *task)
+{
+}
 #endif /* CONFIG_BPF_SYSCALL */
 
 void __bpf_free_used_btfs(struct bpf_prog_aux *aux,
@@ -1886,6 +1891,8 @@  extern const struct bpf_func_proto bpf_this_cpu_ptr_proto;
 extern const struct bpf_func_proto bpf_ktime_get_coarse_ns_proto;
 extern const struct bpf_func_proto bpf_sock_from_file_proto;
 extern const struct bpf_func_proto bpf_get_socket_ptr_cookie_proto;
+extern const struct bpf_func_proto bpf_task_storage_get_proto;
+extern const struct bpf_func_proto bpf_task_storage_delete_proto;
 
 const struct bpf_func_proto *bpf_tracing_func_proto(
 	enum bpf_func_id func_id, const struct bpf_prog *prog);
diff --git a/include/linux/bpf_lsm.h b/include/linux/bpf_lsm.h
index 0d1c33ace3987..479c101546ad1 100644
--- a/include/linux/bpf_lsm.h
+++ b/include/linux/bpf_lsm.h
@@ -38,21 +38,9 @@  static inline struct bpf_storage_blob *bpf_inode(
 	return inode->i_security + bpf_lsm_blob_sizes.lbs_inode;
 }
 
-static inline struct bpf_storage_blob *bpf_task(
-	const struct task_struct *task)
-{
-	if (unlikely(!task->security))
-		return NULL;
-
-	return task->security + bpf_lsm_blob_sizes.lbs_task;
-}
-
 extern const struct bpf_func_proto bpf_inode_storage_get_proto;
 extern const struct bpf_func_proto bpf_inode_storage_delete_proto;
-extern const struct bpf_func_proto bpf_task_storage_get_proto;
-extern const struct bpf_func_proto bpf_task_storage_delete_proto;
 void bpf_inode_storage_free(struct inode *inode);
-void bpf_task_storage_free(struct task_struct *task);
 
 #else /* !CONFIG_BPF_LSM */
 
@@ -73,20 +61,10 @@  static inline struct bpf_storage_blob *bpf_inode(
 	return NULL;
 }
 
-static inline struct bpf_storage_blob *bpf_task(
-	const struct task_struct *task)
-{
-	return NULL;
-}
-
 static inline void bpf_inode_storage_free(struct inode *inode)
 {
 }
 
-static inline void bpf_task_storage_free(struct task_struct *task)
-{
-}
-
 #endif /* CONFIG_BPF_LSM */
 
 #endif /* _LINUX_BPF_LSM_H */
diff --git a/include/linux/bpf_types.h b/include/linux/bpf_types.h
index 99f7fd657d87a..b9edee336d804 100644
--- a/include/linux/bpf_types.h
+++ b/include/linux/bpf_types.h
@@ -109,8 +109,8 @@  BPF_MAP_TYPE(BPF_MAP_TYPE_SOCKHASH, sock_hash_ops)
 #endif
 #ifdef CONFIG_BPF_LSM
 BPF_MAP_TYPE(BPF_MAP_TYPE_INODE_STORAGE, inode_storage_map_ops)
-BPF_MAP_TYPE(BPF_MAP_TYPE_TASK_STORAGE, task_storage_map_ops)
 #endif
+BPF_MAP_TYPE(BPF_MAP_TYPE_TASK_STORAGE, task_storage_map_ops)
 BPF_MAP_TYPE(BPF_MAP_TYPE_CPUMAP, cpu_map_ops)
 #if defined(CONFIG_XDP_SOCKETS)
 BPF_MAP_TYPE(BPF_MAP_TYPE_XSKMAP, xsk_map_ops)
diff --git a/include/linux/sched.h b/include/linux/sched.h
index 4d568288abf9f..e5fbf8e6952ab 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -42,6 +42,7 @@  struct audit_context;
 struct backing_dev_info;
 struct bio_list;
 struct blk_plug;
+struct bpf_local_storage;
 struct capture_control;
 struct cfs_rq;
 struct fs_struct;
@@ -1348,6 +1349,10 @@  struct task_struct {
 	/* Used by LSM modules for access restriction: */
 	void				*security;
 #endif
+#ifdef CONFIG_BPF_SYSCALL
+	/* Used by BPF task local storage */
+	struct bpf_local_storage __rcu	*bpf_storage;
+#endif
 
 #ifdef CONFIG_GCC_PLUGIN_STACKLEAK
 	unsigned long			lowest_stack;
diff --git a/kernel/bpf/Makefile b/kernel/bpf/Makefile
index d1249340fd6ba..ca995fdfa45e7 100644
--- a/kernel/bpf/Makefile
+++ b/kernel/bpf/Makefile
@@ -8,9 +8,8 @@  CFLAGS_core.o += $(call cc-disable-warning, override-init) $(cflags-nogcse-yy)
 
 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) += local_storage.o queue_stack_maps.o ringbuf.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
diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index dd5aedee99e73..9bd47ad2b26f1 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -140,17 +140,18 @@  static void __bpf_selem_unlink_storage(struct bpf_local_storage_elem *selem)
 {
 	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);
@@ -167,6 +168,7 @@  void bpf_selem_unlink_map(struct bpf_local_storage_elem *selem)
 {
 	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 */
@@ -174,21 +176,22 @@  void bpf_selem_unlink_map(struct bpf_local_storage_elem *selem)
 
 	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)
@@ -224,16 +227,18 @@  bpf_local_storage_lookup(struct bpf_local_storage *local_storage,
 
 	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;
@@ -327,6 +332,7 @@  bpf_local_storage_update(void *owner, struct bpf_local_storage_map *smap,
 	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 */
@@ -374,7 +380,7 @@  bpf_local_storage_update(void *owner, struct bpf_local_storage_map *smap,
 		}
 	}
 
-	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))) {
@@ -428,11 +434,11 @@  bpf_local_storage_update(void *owner, struct bpf_local_storage_map *smap,
 	}
 
 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);
 }
 
diff --git a/kernel/bpf/bpf_lsm.c b/kernel/bpf/bpf_lsm.c
index 1622a44d1617e..9829f381b51c5 100644
--- a/kernel/bpf/bpf_lsm.c
+++ b/kernel/bpf/bpf_lsm.c
@@ -115,10 +115,6 @@  bpf_lsm_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 		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:
diff --git a/kernel/bpf/bpf_task_storage.c b/kernel/bpf/bpf_task_storage.c
index e0da0258b732d..2034019966d44 100644
--- a/kernel/bpf/bpf_task_storage.c
+++ b/kernel/bpf/bpf_task_storage.c
@@ -15,7 +15,6 @@ 
 #include <linux/bpf_local_storage.h>
 #include <linux/filter.h>
 #include <uapi/linux/btf.h>
-#include <linux/bpf_lsm.h>
 #include <linux/btf_ids.h>
 #include <linux/fdtable.h>
 
@@ -24,12 +23,8 @@  DEFINE_BPF_STORAGE_CACHE(task_cache);
 static struct bpf_local_storage __rcu **task_storage_ptr(void *owner)
 {
 	struct task_struct *task = owner;
-	struct bpf_storage_blob *bsb;
 
-	bsb = bpf_task(task);
-	if (!bsb)
-		return NULL;
-	return &bsb->storage;
+	return &task->bpf_storage;
 }
 
 static struct bpf_local_storage_data *
@@ -38,13 +33,8 @@  task_storage_lookup(struct task_struct *task, struct bpf_map *map,
 {
 	struct bpf_local_storage *task_storage;
 	struct bpf_local_storage_map *smap;
-	struct bpf_storage_blob *bsb;
-
-	bsb = bpf_task(task);
-	if (!bsb)
-		return NULL;
 
-	task_storage = rcu_dereference(bsb->storage);
+	task_storage = rcu_dereference(task->bpf_storage);
 	if (!task_storage)
 		return NULL;
 
@@ -57,16 +47,12 @@  void bpf_task_storage_free(struct task_struct *task)
 	struct bpf_local_storage_elem *selem;
 	struct bpf_local_storage *local_storage;
 	bool free_task_storage = false;
-	struct bpf_storage_blob *bsb;
 	struct hlist_node *n;
-
-	bsb = bpf_task(task);
-	if (!bsb)
-		return;
+	unsigned long flags;
 
 	rcu_read_lock();
 
-	local_storage = rcu_dereference(bsb->storage);
+	local_storage = rcu_dereference(task->bpf_storage);
 	if (!local_storage) {
 		rcu_read_unlock();
 		return;
@@ -81,7 +67,7 @@  void bpf_task_storage_free(struct task_struct *task)
 	 * when unlinking elem from the local_storage->list and
 	 * the map's bucket->list.
 	 */
-	raw_spin_lock_bh(&local_storage->lock);
+	raw_spin_lock_irqsave(&local_storage->lock, flags);
 	hlist_for_each_entry_safe(selem, n, &local_storage->list, snode) {
 		/* Always unlink from map before unlinking from
 		 * local_storage.
@@ -90,7 +76,7 @@  void bpf_task_storage_free(struct task_struct *task)
 		free_task_storage = bpf_selem_unlink_storage_nolock(
 			local_storage, selem, false);
 	}
-	raw_spin_unlock_bh(&local_storage->lock);
+	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 	rcu_read_unlock();
 
 	/* free_task_storage should always be true as long as
@@ -225,11 +211,9 @@  BPF_CALL_4(bpf_task_storage_get, struct bpf_map *, map, struct task_struct *,
 	if (sdata)
 		return (unsigned long)sdata->data;
 
-	/* This helper must only be called from places where the lifetime of the task
-	 * is guaranteed. Either by being refcounted or by being protected
-	 * by an RCU read-side critical section.
-	 */
-	if (flags & BPF_LOCAL_STORAGE_GET_F_CREATE) {
+	/* only allocate new storage, when the task is refcounted */
+	if (refcount_read(&task->usage) &&
+	    (flags & BPF_LOCAL_STORAGE_GET_F_CREATE)) {
 		sdata = bpf_local_storage_update(
 			task, (struct bpf_local_storage_map *)map, value,
 			BPF_NOEXIST);
diff --git a/kernel/fork.c b/kernel/fork.c
index d66cd1014211b..181604db2d65e 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -96,6 +96,7 @@ 
 #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>
@@ -734,6 +735,7 @@  void __put_task_struct(struct task_struct *tsk)
 	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);
@@ -2062,6 +2064,9 @@  static __latent_entropy struct task_struct *copy_process(
 	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);
diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
index b0c45d923f0f9..e9701744d8e4c 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -1367,6 +1367,10 @@  bpf_tracing_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 		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;
 	}