165ffd797SStanislav Fomichev // SPDX-License-Identifier: GPL-2.0-only
265ffd797SStanislav Fomichev
365ffd797SStanislav Fomichev #include <sys/socket.h>
465ffd797SStanislav Fomichev #include <linux/bpf.h>
565ffd797SStanislav Fomichev #include <bpf/bpf_helpers.h>
665ffd797SStanislav Fomichev
765ffd797SStanislav Fomichev int invocations = 0, in_use = 0;
865ffd797SStanislav Fomichev
94fb5f949SStanislav Fomichev struct {
104fb5f949SStanislav Fomichev __uint(type, BPF_MAP_TYPE_SK_STORAGE);
114fb5f949SStanislav Fomichev __uint(map_flags, BPF_F_NO_PREALLOC);
124fb5f949SStanislav Fomichev __type(key, int);
134fb5f949SStanislav Fomichev __type(value, int);
144fb5f949SStanislav Fomichev } sk_map SEC(".maps");
154fb5f949SStanislav Fomichev
1665ffd797SStanislav Fomichev SEC("cgroup/sock_create")
sock(struct bpf_sock * ctx)1765ffd797SStanislav Fomichev int sock(struct bpf_sock *ctx)
1865ffd797SStanislav Fomichev {
194fb5f949SStanislav Fomichev int *sk_storage;
2065ffd797SStanislav Fomichev
2165ffd797SStanislav Fomichev if (ctx->type != SOCK_DGRAM)
2265ffd797SStanislav Fomichev return 1;
2365ffd797SStanislav Fomichev
244fb5f949SStanislav Fomichev sk_storage = bpf_sk_storage_get(&sk_map, ctx, 0,
254fb5f949SStanislav Fomichev BPF_SK_STORAGE_GET_F_CREATE);
264fb5f949SStanislav Fomichev if (!sk_storage)
274fb5f949SStanislav Fomichev return 0;
284fb5f949SStanislav Fomichev *sk_storage = 0xdeadbeef;
294fb5f949SStanislav Fomichev
3065ffd797SStanislav Fomichev __sync_fetch_and_add(&invocations, 1);
3165ffd797SStanislav Fomichev
3265ffd797SStanislav Fomichev if (in_use > 0) {
3365ffd797SStanislav Fomichev /* BPF_CGROUP_INET_SOCK_RELEASE is _not_ called
3465ffd797SStanislav Fomichev * when we return an error from the BPF
3565ffd797SStanislav Fomichev * program!
3665ffd797SStanislav Fomichev */
3765ffd797SStanislav Fomichev return 0;
3865ffd797SStanislav Fomichev }
3965ffd797SStanislav Fomichev
4065ffd797SStanislav Fomichev __sync_fetch_and_add(&in_use, 1);
4165ffd797SStanislav Fomichev return 1;
4265ffd797SStanislav Fomichev }
4365ffd797SStanislav Fomichev
4465ffd797SStanislav Fomichev SEC("cgroup/sock_release")
sock_release(struct bpf_sock * ctx)4565ffd797SStanislav Fomichev int sock_release(struct bpf_sock *ctx)
4665ffd797SStanislav Fomichev {
474fb5f949SStanislav Fomichev int *sk_storage;
4865ffd797SStanislav Fomichev
4965ffd797SStanislav Fomichev if (ctx->type != SOCK_DGRAM)
5065ffd797SStanislav Fomichev return 1;
5165ffd797SStanislav Fomichev
524fb5f949SStanislav Fomichev sk_storage = bpf_sk_storage_get(&sk_map, ctx, 0, 0);
534fb5f949SStanislav Fomichev if (!sk_storage || *sk_storage != 0xdeadbeef)
544fb5f949SStanislav Fomichev return 0;
554fb5f949SStanislav Fomichev
5665ffd797SStanislav Fomichev __sync_fetch_and_add(&invocations, 1);
5765ffd797SStanislav Fomichev __sync_fetch_and_add(&in_use, -1);
5865ffd797SStanislav Fomichev return 1;
5965ffd797SStanislav Fomichev }
60