1 | // SPDX-License-Identifier: GPL-2.0 |
2 | /* Copyright (c) 2023 Meta Platforms, Inc. and affiliates. */ |
3 | |
4 | #include "vmlinux.h" |
5 | #include "bpf_tracing_net.h" |
6 | #include <bpf/bpf_tracing.h> |
7 | #include <bpf/bpf_helpers.h> |
8 | |
9 | long create_errs = 0; |
10 | long create_cnts = 0; |
11 | long kmalloc_cnts = 0; |
12 | __u32 bench_pid = 0; |
13 | |
14 | struct storage { |
15 | __u8 data[64]; |
16 | }; |
17 | |
18 | struct { |
19 | __uint(type, BPF_MAP_TYPE_SK_STORAGE); |
20 | __uint(map_flags, BPF_F_NO_PREALLOC); |
21 | __type(key, int); |
22 | __type(value, struct storage); |
23 | } sk_storage_map SEC(".maps" ); |
24 | |
25 | struct { |
26 | __uint(type, BPF_MAP_TYPE_TASK_STORAGE); |
27 | __uint(map_flags, BPF_F_NO_PREALLOC); |
28 | __type(key, int); |
29 | __type(value, struct storage); |
30 | } task_storage_map SEC(".maps" ); |
31 | |
32 | SEC("raw_tp/kmalloc" ) |
33 | int BPF_PROG(kmalloc, unsigned long call_site, const void *ptr, |
34 | size_t bytes_req, size_t bytes_alloc, gfp_t gfp_flags, |
35 | int node) |
36 | { |
37 | __sync_fetch_and_add(&kmalloc_cnts, 1); |
38 | |
39 | return 0; |
40 | } |
41 | |
42 | SEC("tp_btf/sched_process_fork" ) |
43 | int BPF_PROG(sched_process_fork, struct task_struct *parent, struct task_struct *child) |
44 | { |
45 | struct storage *stg; |
46 | |
47 | if (parent->tgid != bench_pid) |
48 | return 0; |
49 | |
50 | stg = bpf_task_storage_get(&task_storage_map, child, NULL, |
51 | BPF_LOCAL_STORAGE_GET_F_CREATE); |
52 | if (stg) |
53 | __sync_fetch_and_add(&create_cnts, 1); |
54 | else |
55 | __sync_fetch_and_add(&create_errs, 1); |
56 | |
57 | return 0; |
58 | } |
59 | |
60 | SEC("lsm.s/socket_post_create" ) |
61 | int BPF_PROG(socket_post_create, struct socket *sock, int family, int type, |
62 | int protocol, int kern) |
63 | { |
64 | struct storage *stg; |
65 | __u32 pid; |
66 | |
67 | pid = bpf_get_current_pid_tgid() >> 32; |
68 | if (pid != bench_pid) |
69 | return 0; |
70 | |
71 | stg = bpf_sk_storage_get(&sk_storage_map, sock->sk, NULL, |
72 | BPF_LOCAL_STORAGE_GET_F_CREATE); |
73 | |
74 | if (stg) |
75 | __sync_fetch_and_add(&create_cnts, 1); |
76 | else |
77 | __sync_fetch_and_add(&create_errs, 1); |
78 | |
79 | return 0; |
80 | } |
81 | |
82 | char __license[] SEC("license" ) = "GPL" ; |
83 | |