1 // SPDX-License-Identifier: GPL-2.0
2 // Copyright (c) 2020 Wenbo Zhang
3 #include <vmlinux.h>
4 #include <bpf/bpf_helpers.h>
5 #include <bpf/bpf_core_read.h>
6 #include <bpf/bpf_tracing.h>
7 #include "tcpconnlat.h"
8
9 #define AF_INET 2
10 #define AF_INET6 10
11
12 const volatile __u64 targ_min_us = 0;
13 const volatile pid_t targ_tgid = 0;
14
15 struct piddata {
16 char comm[TASK_COMM_LEN];
17 u64 ts;
18 u32 tgid;
19 };
20
21 struct {
22 __uint(type, BPF_MAP_TYPE_HASH);
23 __uint(max_entries, 4096);
24 __type(key, struct sock *);
25 __type(value, struct piddata);
26 } start SEC(".maps");
27
28 struct {
29 __uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
30 __uint(key_size, sizeof(u32));
31 __uint(value_size, sizeof(u32));
32 } events SEC(".maps");
33
trace_connect(struct sock * sk)34 static __always_inline int trace_connect(struct sock *sk)
35 {
36 u32 tgid = bpf_get_current_pid_tgid() >> 32;
37 struct piddata piddata = {};
38
39 if (targ_tgid && targ_tgid != tgid)
40 return 0;
41
42 bpf_get_current_comm(&piddata.comm, sizeof(piddata.comm));
43 piddata.ts = bpf_ktime_get_ns();
44 piddata.tgid = tgid;
45 bpf_map_update_elem(&start, &sk, &piddata, 0);
46 return 0;
47 }
48
49 SEC("fentry/tcp_v4_connect")
BPF_PROG(tcp_v4_connect,struct sock * sk)50 int BPF_PROG(tcp_v4_connect, struct sock *sk)
51 {
52 return trace_connect(sk);
53 }
54
55 SEC("kprobe/tcp_v6_connect")
BPF_KPROBE(tcp_v6_connect,struct sock * sk)56 int BPF_KPROBE(tcp_v6_connect, struct sock *sk)
57 {
58 return trace_connect(sk);
59 }
60
61 SEC("fentry/tcp_rcv_state_process")
BPF_PROG(tcp_rcv_state_process,struct sock * sk)62 int BPF_PROG(tcp_rcv_state_process, struct sock *sk)
63 {
64 struct piddata *piddatap;
65 struct event event = {};
66 s64 delta;
67 u64 ts;
68
69 if (sk->__sk_common.skc_state != TCP_SYN_SENT)
70 return 0;
71
72 piddatap = bpf_map_lookup_elem(&start, &sk);
73 if (!piddatap)
74 return 0;
75
76 ts = bpf_ktime_get_ns();
77 delta = (s64)(ts - piddatap->ts);
78 if (delta < 0)
79 goto cleanup;
80
81 event.delta_us = delta / 1000U;
82 if (targ_min_us && event.delta_us < targ_min_us)
83 goto cleanup;
84 __builtin_memcpy(&event.comm, piddatap->comm,
85 sizeof(event.comm));
86 event.ts_us = ts / 1000;
87 event.tgid = piddatap->tgid;
88 event.lport = sk->__sk_common.skc_num;
89 event.dport = sk->__sk_common.skc_dport;
90 event.af = sk->__sk_common.skc_family;
91 if (event.af == AF_INET) {
92 event.saddr_v4 = sk->__sk_common.skc_rcv_saddr;
93 event.daddr_v4 = sk->__sk_common.skc_daddr;
94 } else {
95 BPF_CORE_READ_INTO(&event.saddr_v6, sk,
96 __sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
97 BPF_CORE_READ_INTO(&event.daddr_v6, sk,
98 __sk_common.skc_v6_daddr.in6_u.u6_addr32);
99 }
100 bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU,
101 &event, sizeof(event));
102
103 cleanup:
104 bpf_map_delete_elem(&start, &sk);
105 return 0;
106 }
107
108 char LICENSE[] SEC("license") = "GPL";
109