kernel/bpf/bpf_task_storage.c

0001 // SPDX-License-Identifier: GPL-2.0
0002 /*
0003  * Copyright (c) 2020 Facebook
0004  * Copyright 2020 Google LLC.
0005  */
0006
0007 #include <linux/pid.h>
0008 #include <linux/sched.h>
0009 #include <linux/rculist.h>
0010 #include <linux/list.h>
0011 #include <linux/hash.h>
0012 #include <linux/types.h>
0013 #include <linux/spinlock.h>
0014 #include <linux/bpf.h>
0015 #include <linux/bpf_local_storage.h>
0016 #include <linux/filter.h>
0017 #include <uapi/linux/btf.h>
0018 #include <linux/btf_ids.h>
0019 #include <linux/fdtable.h>
0020 #include <linux/rcupdate_trace.h>
0021
0022 DEFINE_BPF_STORAGE_CACHE(task_cache);
0023
0024 static DEFINE_PER_CPU(int, bpf_task_storage_busy);
0025
0026 static void bpf_task_storage_lock(void)
0027 {
0028     migrate_disable();
0029     __this_cpu_inc(bpf_task_storage_busy);
0030 }
0031
0032 static void bpf_task_storage_unlock(void)
0033 {
0034     __this_cpu_dec(bpf_task_storage_busy);
0035     migrate_enable();
0036 }
0037
0038 static bool bpf_task_storage_trylock(void)
0039 {
0040     migrate_disable();
0041     if (unlikely(__this_cpu_inc_return(bpf_task_storage_busy) != 1)) {
0042         __this_cpu_dec(bpf_task_storage_busy);
0043         migrate_enable();
0044         return false;
0045     }
0046     return true;
0047 }
0048
0049 static struct bpf_local_storage __rcu **task_storage_ptr(void *owner)
0050 {
0051     struct task_struct *task = owner;
0052
0053     return &task->bpf_storage;
0054 }
0055
0056 static struct bpf_local_storage_data *
0057 task_storage_lookup(struct task_struct *task, struct bpf_map *map,
0058             bool cacheit_lockit)
0059 {
0060     struct bpf_local_storage *task_storage;
0061     struct bpf_local_storage_map *smap;
0062
0063     task_storage =
0064         rcu_dereference_check(task->bpf_storage, bpf_rcu_lock_held());
0065     if (!task_storage)
0066         return NULL;
0067
0068     smap = (struct bpf_local_storage_map *)map;
0069     return bpf_local_storage_lookup(task_storage, smap, cacheit_lockit);
0070 }
0071
0072 void bpf_task_storage_free(struct task_struct *task)
0073 {
0074     struct bpf_local_storage_elem *selem;
0075     struct bpf_local_storage *local_storage;
0076     bool free_task_storage = false;
0077     struct hlist_node *n;
0078     unsigned long flags;
0079
0080     rcu_read_lock();
0081
0082     local_storage = rcu_dereference(task->bpf_storage);
0083     if (!local_storage) {
0084         rcu_read_unlock();
0085         return;
0086     }
0087
0088     /* Neither the bpf_prog nor the bpf-map's syscall
0089      * could be modifying the local_storage->list now.
0090      * Thus, no elem can be added-to or deleted-from the
0091      * local_storage->list by the bpf_prog or by the bpf-map's syscall.
0092      *
0093      * It is racing with bpf_local_storage_map_free() alone
0094      * when unlinking elem from the local_storage->list and
0095      * the map's bucket->list.
0096      */
0097     bpf_task_storage_lock();
0098     raw_spin_lock_irqsave(&local_storage->lock, flags);
0099     hlist_for_each_entry_safe(selem, n, &local_storage->list, snode) {
0100         /* Always unlink from map before unlinking from
0101          * local_storage.
0102          */
0103         bpf_selem_unlink_map(selem);
0104         free_task_storage = bpf_selem_unlink_storage_nolock(
0105             local_storage, selem, false, false);
0106     }
0107     raw_spin_unlock_irqrestore(&local_storage->lock, flags);
0108     bpf_task_storage_unlock();
0109     rcu_read_unlock();
0110
0111     /* free_task_storage should always be true as long as
0112      * local_storage->list was non-empty.
0113      */
0114     if (free_task_storage)
0115         kfree_rcu(local_storage, rcu);
0116 }
0117
0118 static void *bpf_pid_task_storage_lookup_elem(struct bpf_map *map, void *key)
0119 {
0120     struct bpf_local_storage_data *sdata;
0121     struct task_struct *task;
0122     unsigned int f_flags;
0123     struct pid *pid;
0124     int fd, err;
0125
0126     fd = *(int *)key;
0127     pid = pidfd_get_pid(fd, &f_flags);
0128     if (IS_ERR(pid))
0129         return ERR_CAST(pid);
0130
0131     /* We should be in an RCU read side critical section, it should be safe
0132      * to call pid_task.
0133      */
0134     WARN_ON_ONCE(!rcu_read_lock_held());
0135     task = pid_task(pid, PIDTYPE_PID);
0136     if (!task) {
0137         err = -ENOENT;
0138         goto out;
0139     }
0140
0141     bpf_task_storage_lock();
0142     sdata = task_storage_lookup(task, map, true);
0143     bpf_task_storage_unlock();
0144     put_pid(pid);
0145     return sdata ? sdata->data : NULL;
0146 out:
0147     put_pid(pid);
0148     return ERR_PTR(err);
0149 }
0150
0151 static int bpf_pid_task_storage_update_elem(struct bpf_map *map, void *key,
0152                         void *value, u64 map_flags)
0153 {
0154     struct bpf_local_storage_data *sdata;
0155     struct task_struct *task;
0156     unsigned int f_flags;
0157     struct pid *pid;
0158     int fd, err;
0159
0160     fd = *(int *)key;
0161     pid = pidfd_get_pid(fd, &f_flags);
0162     if (IS_ERR(pid))
0163         return PTR_ERR(pid);
0164
0165     /* We should be in an RCU read side critical section, it should be safe
0166      * to call pid_task.
0167      */
0168     WARN_ON_ONCE(!rcu_read_lock_held());
0169     task = pid_task(pid, PIDTYPE_PID);
0170     if (!task) {
0171         err = -ENOENT;
0172         goto out;
0173     }
0174
0175     bpf_task_storage_lock();
0176     sdata = bpf_local_storage_update(
0177         task, (struct bpf_local_storage_map *)map, value, map_flags,
0178         GFP_ATOMIC);
0179     bpf_task_storage_unlock();
0180
0181     err = PTR_ERR_OR_ZERO(sdata);
0182 out:
0183     put_pid(pid);
0184     return err;
0185 }
0186
0187 static int task_storage_delete(struct task_struct *task, struct bpf_map *map)
0188 {
0189     struct bpf_local_storage_data *sdata;
0190
0191     sdata = task_storage_lookup(task, map, false);
0192     if (!sdata)
0193         return -ENOENT;
0194
0195     bpf_selem_unlink(SELEM(sdata), true);
0196
0197     return 0;
0198 }
0199
0200 static int bpf_pid_task_storage_delete_elem(struct bpf_map *map, void *key)
0201 {
0202     struct task_struct *task;
0203     unsigned int f_flags;
0204     struct pid *pid;
0205     int fd, err;
0206
0207     fd = *(int *)key;
0208     pid = pidfd_get_pid(fd, &f_flags);
0209     if (IS_ERR(pid))
0210         return PTR_ERR(pid);
0211
0212     /* We should be in an RCU read side critical section, it should be safe
0213      * to call pid_task.
0214      */
0215     WARN_ON_ONCE(!rcu_read_lock_held());
0216     task = pid_task(pid, PIDTYPE_PID);
0217     if (!task) {
0218         err = -ENOENT;
0219         goto out;
0220     }
0221
0222     bpf_task_storage_lock();
0223     err = task_storage_delete(task, map);
0224     bpf_task_storage_unlock();
0225 out:
0226     put_pid(pid);
0227     return err;
0228 }
0229
0230 /* *gfp_flags* is a hidden argument provided by the verifier */
0231 BPF_CALL_5(bpf_task_storage_get, struct bpf_map *, map, struct task_struct *,
0232        task, void *, value, u64, flags, gfp_t, gfp_flags)
0233 {
0234     struct bpf_local_storage_data *sdata;
0235
0236     WARN_ON_ONCE(!bpf_rcu_lock_held());
0237     if (flags & ~(BPF_LOCAL_STORAGE_GET_F_CREATE))
0238         return (unsigned long)NULL;
0239
0240     if (!task)
0241         return (unsigned long)NULL;
0242
0243     if (!bpf_task_storage_trylock())
0244         return (unsigned long)NULL;
0245
0246     sdata = task_storage_lookup(task, map, true);
0247     if (sdata)
0248         goto unlock;
0249
0250     /* only allocate new storage, when the task is refcounted */
0251     if (refcount_read(&task->usage) &&
0252         (flags & BPF_LOCAL_STORAGE_GET_F_CREATE))
0253         sdata = bpf_local_storage_update(
0254             task, (struct bpf_local_storage_map *)map, value,
0255             BPF_NOEXIST, gfp_flags);
0256
0257 unlock:
0258     bpf_task_storage_unlock();
0259     return IS_ERR_OR_NULL(sdata) ? (unsigned long)NULL :
0260         (unsigned long)sdata->data;
0261 }
0262
0263 BPF_CALL_2(bpf_task_storage_delete, struct bpf_map *, map, struct task_struct *,
0264        task)
0265 {
0266     int ret;
0267
0268     WARN_ON_ONCE(!bpf_rcu_lock_held());
0269     if (!task)
0270         return -EINVAL;
0271
0272     if (!bpf_task_storage_trylock())
0273         return -EBUSY;
0274
0275     /* This helper must only be called from places where the lifetime of the task
0276      * is guaranteed. Either by being refcounted or by being protected
0277      * by an RCU read-side critical section.
0278      */
0279     ret = task_storage_delete(task, map);
0280     bpf_task_storage_unlock();
0281     return ret;
0282 }
0283
0284 static int notsupp_get_next_key(struct bpf_map *map, void *key, void *next_key)
0285 {
0286     return -ENOTSUPP;
0287 }
0288
0289 static struct bpf_map *task_storage_map_alloc(union bpf_attr *attr)
0290 {
0291     struct bpf_local_storage_map *smap;
0292
0293     smap = bpf_local_storage_map_alloc(attr);
0294     if (IS_ERR(smap))
0295         return ERR_CAST(smap);
0296
0297     smap->cache_idx = bpf_local_storage_cache_idx_get(&task_cache);
0298     return &smap->map;
0299 }
0300
0301 static void task_storage_map_free(struct bpf_map *map)
0302 {
0303     struct bpf_local_storage_map *smap;
0304
0305     smap = (struct bpf_local_storage_map *)map;
0306     bpf_local_storage_cache_idx_free(&task_cache, smap->cache_idx);
0307     bpf_local_storage_map_free(smap, &bpf_task_storage_busy);
0308 }
0309
0310 BTF_ID_LIST_SINGLE(task_storage_map_btf_ids, struct, bpf_local_storage_map)
0311 const struct bpf_map_ops task_storage_map_ops = {
0312     .map_meta_equal = bpf_map_meta_equal,
0313     .map_alloc_check = bpf_local_storage_map_alloc_check,
0314     .map_alloc = task_storage_map_alloc,
0315     .map_free = task_storage_map_free,
0316     .map_get_next_key = notsupp_get_next_key,
0317     .map_lookup_elem = bpf_pid_task_storage_lookup_elem,
0318     .map_update_elem = bpf_pid_task_storage_update_elem,
0319     .map_delete_elem = bpf_pid_task_storage_delete_elem,
0320     .map_check_btf = bpf_local_storage_map_check_btf,
0321     .map_btf_id = &task_storage_map_btf_ids[0],
0322     .map_owner_storage_ptr = task_storage_ptr,
0323 };
0324
0325 const struct bpf_func_proto bpf_task_storage_get_proto = {
0326     .func = bpf_task_storage_get,
0327     .gpl_only = false,
0328     .ret_type = RET_PTR_TO_MAP_VALUE_OR_NULL,
0329     .arg1_type = ARG_CONST_MAP_PTR,
0330     .arg2_type = ARG_PTR_TO_BTF_ID,
0331     .arg2_btf_id = &btf_tracing_ids[BTF_TRACING_TYPE_TASK],
0332     .arg3_type = ARG_PTR_TO_MAP_VALUE_OR_NULL,
0333     .arg4_type = ARG_ANYTHING,
0334 };
0335
0336 const struct bpf_func_proto bpf_task_storage_delete_proto = {
0337     .func = bpf_task_storage_delete,
0338     .gpl_only = false,
0339     .ret_type = RET_INTEGER,
0340     .arg1_type = ARG_CONST_MAP_PTR,
0341     .arg2_type = ARG_PTR_TO_BTF_ID,
0342     .arg2_btf_id = &btf_tracing_ids[BTF_TRACING_TYPE_TASK],
0343 };