1// SPDX-License-Identifier: GPL-2.0
  2/*
  3 * Copyright (c) 2022 Meta Platforms, Inc. and affiliates.
  4 */
  5
  6#include <linux/types.h>
  7#include <linux/bpf.h>
  8#include <linux/bpf_local_storage.h>
  9#include <uapi/linux/btf.h>
 10#include <linux/btf_ids.h>
 11
 12DEFINE_BPF_STORAGE_CACHE(cgroup_cache);
 13
 14static DEFINE_PER_CPU(int, bpf_cgrp_storage_busy);
 15
 16static void bpf_cgrp_storage_lock(void)
 17{
 18	migrate_disable();
 19	this_cpu_inc(bpf_cgrp_storage_busy);
 20}
 21
 22static void bpf_cgrp_storage_unlock(void)
 23{
 24	this_cpu_dec(bpf_cgrp_storage_busy);
 25	migrate_enable();
 26}
 27
 28static bool bpf_cgrp_storage_trylock(void)
 29{
 30	migrate_disable();
 31	if (unlikely(this_cpu_inc_return(bpf_cgrp_storage_busy) != 1)) {
 32		this_cpu_dec(bpf_cgrp_storage_busy);
 33		migrate_enable();
 34		return false;
 35	}
 36	return true;
 37}
 38
 39static struct bpf_local_storage __rcu **cgroup_storage_ptr(void *owner)
 40{
 41	struct cgroup *cg = owner;
 42
 43	return &cg->bpf_cgrp_storage;
 44}
 45
 46void bpf_cgrp_storage_free(struct cgroup *cgroup)
 47{
 48	struct bpf_local_storage *local_storage;
 49
 50	rcu_read_lock();
 51	local_storage = rcu_dereference(cgroup->bpf_cgrp_storage);
 52	if (!local_storage) {
 53		rcu_read_unlock();
 54		return;
 55	}
 56
 57	bpf_cgrp_storage_lock();
 58	bpf_local_storage_destroy(local_storage);
 59	bpf_cgrp_storage_unlock();
 60	rcu_read_unlock();
 61}
 62
 63static struct bpf_local_storage_data *
 64cgroup_storage_lookup(struct cgroup *cgroup, struct bpf_map *map, bool cacheit_lockit)
 65{
 66	struct bpf_local_storage *cgroup_storage;
 67	struct bpf_local_storage_map *smap;
 68
 69	cgroup_storage = rcu_dereference_check(cgroup->bpf_cgrp_storage,
 70					       bpf_rcu_lock_held());
 71	if (!cgroup_storage)
 72		return NULL;
 73
 74	smap = (struct bpf_local_storage_map *)map;
 75	return bpf_local_storage_lookup(cgroup_storage, smap, cacheit_lockit);
 76}
 77
 78static void *bpf_cgrp_storage_lookup_elem(struct bpf_map *map, void *key)
 79{
 80	struct bpf_local_storage_data *sdata;
 81	struct cgroup *cgroup;
 82	int fd;
 83
 84	fd = *(int *)key;
 85	cgroup = cgroup_v1v2_get_from_fd(fd);
 86	if (IS_ERR(cgroup))
 87		return ERR_CAST(cgroup);
 88
 89	bpf_cgrp_storage_lock();
 90	sdata = cgroup_storage_lookup(cgroup, map, true);
 91	bpf_cgrp_storage_unlock();
 92	cgroup_put(cgroup);
 93	return sdata ? sdata->data : NULL;
 94}
 95
 96static long bpf_cgrp_storage_update_elem(struct bpf_map *map, void *key,
 97					 void *value, u64 map_flags)
 98{
 99	struct bpf_local_storage_data *sdata;
100	struct cgroup *cgroup;
101	int fd;
102
103	fd = *(int *)key;
104	cgroup = cgroup_v1v2_get_from_fd(fd);
105	if (IS_ERR(cgroup))
106		return PTR_ERR(cgroup);
107
108	bpf_cgrp_storage_lock();
109	sdata = bpf_local_storage_update(cgroup, (struct bpf_local_storage_map *)map,
110					 value, map_flags, GFP_ATOMIC);
111	bpf_cgrp_storage_unlock();
112	cgroup_put(cgroup);
113	return PTR_ERR_OR_ZERO(sdata);
114}
115
116static int cgroup_storage_delete(struct cgroup *cgroup, struct bpf_map *map)
117{
118	struct bpf_local_storage_data *sdata;
119
120	sdata = cgroup_storage_lookup(cgroup, map, false);
121	if (!sdata)
122		return -ENOENT;
123
124	bpf_selem_unlink(SELEM(sdata), false);
125	return 0;
126}
127
128static long bpf_cgrp_storage_delete_elem(struct bpf_map *map, void *key)
129{
130	struct cgroup *cgroup;
131	int err, fd;
132
133	fd = *(int *)key;
134	cgroup = cgroup_v1v2_get_from_fd(fd);
135	if (IS_ERR(cgroup))
136		return PTR_ERR(cgroup);
137
138	bpf_cgrp_storage_lock();
139	err = cgroup_storage_delete(cgroup, map);
140	bpf_cgrp_storage_unlock();
141	cgroup_put(cgroup);
142	return err;
143}
144
145static int notsupp_get_next_key(struct bpf_map *map, void *key, void *next_key)
146{
147	return -ENOTSUPP;
148}
149
150static struct bpf_map *cgroup_storage_map_alloc(union bpf_attr *attr)
151{
152	return bpf_local_storage_map_alloc(attr, &cgroup_cache, true);
153}
154
155static void cgroup_storage_map_free(struct bpf_map *map)
156{
157	bpf_local_storage_map_free(map, &cgroup_cache, NULL);
158}
159
160/* *gfp_flags* is a hidden argument provided by the verifier */
161BPF_CALL_5(bpf_cgrp_storage_get, struct bpf_map *, map, struct cgroup *, cgroup,
162	   void *, value, u64, flags, gfp_t, gfp_flags)
163{
164	struct bpf_local_storage_data *sdata;
165
166	WARN_ON_ONCE(!bpf_rcu_lock_held());
167	if (flags & ~(BPF_LOCAL_STORAGE_GET_F_CREATE))
168		return (unsigned long)NULL;
169
170	if (!cgroup)
171		return (unsigned long)NULL;
172
173	if (!bpf_cgrp_storage_trylock())
174		return (unsigned long)NULL;
175
176	sdata = cgroup_storage_lookup(cgroup, map, true);
177	if (sdata)
178		goto unlock;
179
180	/* only allocate new storage, when the cgroup is refcounted */
181	if (!percpu_ref_is_dying(&cgroup->self.refcnt) &&
182	    (flags & BPF_LOCAL_STORAGE_GET_F_CREATE))
183		sdata = bpf_local_storage_update(cgroup, (struct bpf_local_storage_map *)map,
184						 value, BPF_NOEXIST, gfp_flags);
185
186unlock:
187	bpf_cgrp_storage_unlock();
188	return IS_ERR_OR_NULL(sdata) ? (unsigned long)NULL : (unsigned long)sdata->data;
189}
190
191BPF_CALL_2(bpf_cgrp_storage_delete, struct bpf_map *, map, struct cgroup *, cgroup)
192{
193	int ret;
194
195	WARN_ON_ONCE(!bpf_rcu_lock_held());
196	if (!cgroup)
197		return -EINVAL;
198
199	if (!bpf_cgrp_storage_trylock())
200		return -EBUSY;
201
202	ret = cgroup_storage_delete(cgroup, map);
203	bpf_cgrp_storage_unlock();
204	return ret;
205}
206
207const struct bpf_map_ops cgrp_storage_map_ops = {
208	.map_meta_equal = bpf_map_meta_equal,
209	.map_alloc_check = bpf_local_storage_map_alloc_check,
210	.map_alloc = cgroup_storage_map_alloc,
211	.map_free = cgroup_storage_map_free,
212	.map_get_next_key = notsupp_get_next_key,
213	.map_lookup_elem = bpf_cgrp_storage_lookup_elem,
214	.map_update_elem = bpf_cgrp_storage_update_elem,
215	.map_delete_elem = bpf_cgrp_storage_delete_elem,
216	.map_check_btf = bpf_local_storage_map_check_btf,
217	.map_mem_usage = bpf_local_storage_map_mem_usage,
218	.map_btf_id = &bpf_local_storage_map_btf_id[0],
219	.map_owner_storage_ptr = cgroup_storage_ptr,
220};
221
222const struct bpf_func_proto bpf_cgrp_storage_get_proto = {
223	.func		= bpf_cgrp_storage_get,
224	.gpl_only	= false,
225	.ret_type	= RET_PTR_TO_MAP_VALUE_OR_NULL,
226	.arg1_type	= ARG_CONST_MAP_PTR,
227	.arg2_type	= ARG_PTR_TO_BTF_ID_OR_NULL,
228	.arg2_btf_id	= &bpf_cgroup_btf_id[0],
229	.arg3_type	= ARG_PTR_TO_MAP_VALUE_OR_NULL,
230	.arg4_type	= ARG_ANYTHING,
231};
232
233const struct bpf_func_proto bpf_cgrp_storage_delete_proto = {
234	.func		= bpf_cgrp_storage_delete,
235	.gpl_only	= false,
236	.ret_type	= RET_INTEGER,
237	.arg1_type	= ARG_CONST_MAP_PTR,
238	.arg2_type	= ARG_PTR_TO_BTF_ID_OR_NULL,
239	.arg2_btf_id	= &bpf_cgroup_btf_id[0],
240};