1 // SPDX-License-Identifier: (LGPL-2.1 OR BSD-2-Clause)
2 /* Copyright (c) 2021 Hengqi Chen */
3 #include <vmlinux.h>
4 #include <bpf/bpf_helpers.h>
5 #include <bpf/bpf_tracing.h>
6 #include <bpf/bpf_core_read.h>
7 #include "tcpstates.h"
8
9 #define MAX_ENTRIES 10240
10 #define AF_INET 2
11 #define AF_INET6 10
12
13 const volatile bool filter_by_sport = false;
14 const volatile bool filter_by_dport = false;
15 const volatile short target_family = 0;
16
17 struct {
18 __uint(type, BPF_MAP_TYPE_HASH);
19 __uint(max_entries, MAX_ENTRIES);
20 __type(key, __u16);
21 __type(value, __u16);
22 } sports SEC(".maps");
23
24 struct {
25 __uint(type, BPF_MAP_TYPE_HASH);
26 __uint(max_entries, MAX_ENTRIES);
27 __type(key, __u16);
28 __type(value, __u16);
29 } dports SEC(".maps");
30
31 struct {
32 __uint(type, BPF_MAP_TYPE_HASH);
33 __uint(max_entries, MAX_ENTRIES);
34 __type(key, struct sock *);
35 __type(value, __u64);
36 } timestamps SEC(".maps");
37
38 struct {
39 __uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
40 __uint(key_size, sizeof(__u32));
41 __uint(value_size, sizeof(__u32));
42 } events SEC(".maps");
43
44 SEC("tracepoint/sock/inet_sock_set_state")
handle_set_state(struct trace_event_raw_inet_sock_set_state * ctx)45 int handle_set_state(struct trace_event_raw_inet_sock_set_state *ctx)
46 {
47 struct sock *sk = (struct sock *)ctx->skaddr;
48 __u16 family = ctx->family;
49 __u16 sport = ctx->sport;
50 __u16 dport = ctx->dport;
51 __u64 *tsp, delta_us, ts;
52 struct event event = {};
53
54 if (ctx->protocol != IPPROTO_TCP)
55 return 0;
56
57 if (target_family && target_family != family)
58 return 0;
59
60 if (filter_by_sport && !bpf_map_lookup_elem(&sports, &sport))
61 return 0;
62
63 if (filter_by_dport && !bpf_map_lookup_elem(&dports, &dport))
64 return 0;
65
66 tsp = bpf_map_lookup_elem(×tamps, &sk);
67 ts = bpf_ktime_get_ns();
68 if (!tsp)
69 delta_us = 0;
70 else
71 delta_us = (ts - *tsp) / 1000;
72
73 event.skaddr = (__u64)sk;
74 event.ts_us = ts / 1000;
75 event.delta_us = delta_us;
76 event.pid = bpf_get_current_pid_tgid() >> 32;
77 event.oldstate = ctx->oldstate;
78 event.newstate = ctx->newstate;
79 event.family = family;
80 event.sport = sport;
81 event.dport = dport;
82 bpf_get_current_comm(&event.task, sizeof(event.task));
83
84 if (family == AF_INET) {
85 bpf_probe_read_kernel(&event.saddr, sizeof(event.saddr), &sk->__sk_common.skc_rcv_saddr);
86 bpf_probe_read_kernel(&event.daddr, sizeof(event.daddr), &sk->__sk_common.skc_daddr);
87 } else { /* family == AF_INET6 */
88 bpf_probe_read_kernel(&event.saddr, sizeof(event.saddr), &sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
89 bpf_probe_read_kernel(&event.daddr, sizeof(event.daddr), &sk->__sk_common.skc_v6_daddr.in6_u.u6_addr32);
90 }
91
92 bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU, &event, sizeof(event));
93
94 if (ctx->newstate == TCP_CLOSE)
95 bpf_map_delete_elem(×tamps, &sk);
96 else
97 bpf_map_update_elem(×tamps, &sk, &ts, BPF_ANY);
98
99 return 0;
100 }
101
102 char LICENSE[] SEC("license") = "Dual BSD/GPL";
103