udp_limit.c (1163B)
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 7int invocations = 0, in_use = 0; 8 9struct { 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 16SEC("cgroup/sock_create") 17int 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 45SEC("cgroup/sock_release") 46int 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}