1 // SPDX-License-Identifier: GPL-2.0-only
2 
3 #include <sys/socket.h>
4 #include <linux/bpf.h>
5 #include <bpf/bpf_helpers.h>
6 
7 int invocations = 0, in_use = 0;
8 
9 struct {
10 	__uint(type, BPF_MAP_TYPE_SK_STORAGE);
11 	__uint(map_flags, BPF_F_NO_PREALLOC);
12 	__type(key, int);
13 	__type(value, int);
14 } sk_map SEC(".maps");
15 
16 SEC("cgroup/sock_create")
17 int sock(struct bpf_sock *ctx)
18 {
19 	int *sk_storage;
20 	__u32 key;
21 
22 	if (ctx->type != SOCK_DGRAM)
23 		return 1;
24 
25 	sk_storage = bpf_sk_storage_get(&sk_map, ctx, 0,
26 					BPF_SK_STORAGE_GET_F_CREATE);
27 	if (!sk_storage)
28 		return 0;
29 	*sk_storage = 0xdeadbeef;
30 
31 	__sync_fetch_and_add(&invocations, 1);
32 
33 	if (in_use > 0) {
34 		/* BPF_CGROUP_INET_SOCK_RELEASE is _not_ called
35 		 * when we return an error from the BPF
36 		 * program!
37 		 */
38 		return 0;
39 	}
40 
41 	__sync_fetch_and_add(&in_use, 1);
42 	return 1;
43 }
44 
45 SEC("cgroup/sock_release")
46 int sock_release(struct bpf_sock *ctx)
47 {
48 	int *sk_storage;
49 	__u32 key;
50 
51 	if (ctx->type != SOCK_DGRAM)
52 		return 1;
53 
54 	sk_storage = bpf_sk_storage_get(&sk_map, ctx, 0, 0);
55 	if (!sk_storage || *sk_storage != 0xdeadbeef)
56 		return 0;
57 
58 	__sync_fetch_and_add(&invocations, 1);
59 	__sync_fetch_and_add(&in_use, -1);
60 	return 1;
61 }
62