1 // SPDX-License-Identifier: GPL-2.0
2 /* Copyright (c) 2022 Hengqi Chen */
3 #include <vmlinux.h>
4 #include <bpf/bpf_core_read.h>
5 #include <bpf/bpf_helpers.h>
6 #include <bpf/bpf_tracing.h>
7 #include "tcplife.h"
8
9 #define MAX_ENTRIES 10240
10 #define AF_INET 2
11 #define AF_INET6 10
12
13 const volatile bool filter_sport = false;
14 const volatile bool filter_dport = false;
15 const volatile __u16 target_sports[MAX_PORTS] = {};
16 const volatile __u16 target_dports[MAX_PORTS] = {};
17 const volatile pid_t target_pid = 0;
18 const volatile __u16 target_family = 0;
19
20 struct {
21 __uint(type, BPF_MAP_TYPE_HASH);
22 __uint(max_entries, MAX_ENTRIES);
23 __type(key, struct sock *);
24 __type(value, __u64);
25 } birth SEC(".maps");
26
27 struct {
28 __uint(type, BPF_MAP_TYPE_HASH);
29 __uint(max_entries, MAX_ENTRIES);
30 __type(key, struct sock *);
31 __type(value, struct ident);
32 } idents SEC(".maps");
33
34 struct {
35 __uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
36 __uint(key_size, sizeof(__u32));
37 __uint(value_size, sizeof(__u32));
38 } events SEC(".maps");
39
40 SEC("tracepoint/sock/inet_sock_set_state")
inet_sock_set_state(struct trace_event_raw_inet_sock_set_state * args)41 int inet_sock_set_state(struct trace_event_raw_inet_sock_set_state *args)
42 {
43 __u64 ts, *start, delta_us, rx_b, tx_b;
44 struct ident ident = {}, *identp;
45 __u16 sport, dport, family;
46 struct event event = {};
47 struct tcp_sock *tp;
48 struct sock *sk;
49 bool found;
50 __u32 pid;
51 int i;
52
53 if (BPF_CORE_READ(args, protocol) != IPPROTO_TCP)
54 return 0;
55
56 family = BPF_CORE_READ(args, family);
57 if (target_family && family != target_family)
58 return 0;
59
60 sport = BPF_CORE_READ(args, sport);
61 if (filter_sport) {
62 found = false;
63 for (i = 0; i < MAX_PORTS; i++) {
64 if (!target_sports[i])
65 return 0;
66 if (sport != target_sports[i])
67 continue;
68 found = true;
69 break;
70 }
71 if (!found)
72 return 0;
73 }
74
75 dport = BPF_CORE_READ(args, dport);
76 if (filter_dport) {
77 found = false;
78 for (i = 0; i < MAX_PORTS; i++) {
79 if (!target_dports[i])
80 return 0;
81 if (dport != target_dports[i])
82 continue;
83 found = true;
84 break;
85 }
86 if (!found)
87 return 0;
88 }
89
90 sk = (struct sock *)BPF_CORE_READ(args, skaddr);
91 if (BPF_CORE_READ(args, newstate) < TCP_FIN_WAIT1) {
92 ts = bpf_ktime_get_ns();
93 bpf_map_update_elem(&birth, &sk, &ts, BPF_ANY);
94 }
95
96 if (BPF_CORE_READ(args, newstate) == TCP_SYN_SENT || BPF_CORE_READ(args, newstate) == TCP_LAST_ACK) {
97 pid = bpf_get_current_pid_tgid() >> 32;
98 if (target_pid && pid != target_pid)
99 return 0;
100 ident.pid = pid;
101 bpf_get_current_comm(ident.comm, sizeof(ident.comm));
102 bpf_map_update_elem(&idents, &sk, &ident, BPF_ANY);
103 }
104
105 if (BPF_CORE_READ(args, newstate) != TCP_CLOSE)
106 return 0;
107
108 start = bpf_map_lookup_elem(&birth, &sk);
109 if (!start) {
110 bpf_map_delete_elem(&idents, &sk);
111 return 0;
112 }
113 ts = bpf_ktime_get_ns();
114 delta_us = (ts - *start) / 1000;
115
116 identp = bpf_map_lookup_elem(&idents, &sk);
117 pid = identp ? identp->pid : bpf_get_current_pid_tgid() >> 32;
118 if (target_pid && pid != target_pid)
119 goto cleanup;
120
121 tp = (struct tcp_sock *)sk;
122 rx_b = BPF_CORE_READ(tp, bytes_received);
123 tx_b = BPF_CORE_READ(tp, bytes_acked);
124
125 event.ts_us = ts / 1000;
126 event.span_us = delta_us;
127 event.rx_b = rx_b;
128 event.tx_b = tx_b;
129 event.pid = pid;
130 event.sport = sport;
131 event.dport = dport;
132 event.family = family;
133 if (!identp)
134 bpf_get_current_comm(event.comm, sizeof(event.comm));
135 else
136 bpf_probe_read_kernel(event.comm, sizeof(event.comm), (void *)identp->comm);
137 if (family == AF_INET) {
138 bpf_probe_read_kernel(&event.saddr, sizeof(args->saddr), BPF_CORE_READ(args, saddr));
139 bpf_probe_read_kernel(&event.daddr, sizeof(args->daddr), BPF_CORE_READ(args, daddr));
140 } else { /* AF_INET6 */
141 bpf_probe_read_kernel(&event.saddr, sizeof(args->saddr_v6), BPF_CORE_READ(args, saddr_v6));
142 bpf_probe_read_kernel(&event.daddr, sizeof(args->daddr_v6), BPF_CORE_READ(args, daddr_v6));
143 }
144 bpf_perf_event_output(args, &events, BPF_F_CURRENT_CPU, &event, sizeof(event));
145
146 cleanup:
147 bpf_map_delete_elem(&birth, &sk);
148 bpf_map_delete_elem(&idents, &sk);
149 return 0;
150 }
151
152 char LICENSE[] SEC("license") = "GPL";
153