1 // SPDX-License-Identifier: GPL-2.0
2 // Copyright (c) 2021 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 <bpf/bpf_endian.h>
8 #include "tcprtt.h"
9 #include "bits.bpf.h"
10 #include "maps.bpf.h"
11
12 const volatile bool targ_laddr_hist = false;
13 const volatile bool targ_raddr_hist = false;
14 const volatile bool targ_show_ext = false;
15 const volatile __u16 targ_sport = 0;
16 const volatile __u16 targ_dport = 0;
17 const volatile __u32 targ_saddr = 0;
18 const volatile __u32 targ_daddr = 0;
19 const volatile bool targ_ms = false;
20
21 #define MAX_ENTRIES 10240
22
23 struct {
24 __uint(type, BPF_MAP_TYPE_HASH);
25 __uint(max_entries, MAX_ENTRIES);
26 __type(key, u64);
27 __type(value, struct hist);
28 } hists SEC(".maps");
29
30 static struct hist zero;
31
32 SEC("fentry/tcp_rcv_established")
BPF_PROG(tcp_rcv,struct sock * sk)33 int BPF_PROG(tcp_rcv, struct sock *sk)
34 {
35 const struct inet_sock *inet = (struct inet_sock *)(sk);
36 struct tcp_sock *ts;
37 struct hist *histp;
38 u64 key, slot;
39 u32 srtt;
40
41 if (targ_sport && targ_sport != inet->inet_sport)
42 return 0;
43 if (targ_dport && targ_dport != sk->__sk_common.skc_dport)
44 return 0;
45 if (targ_saddr && targ_saddr != inet->inet_saddr)
46 return 0;
47 if (targ_daddr && targ_daddr != sk->__sk_common.skc_daddr)
48 return 0;
49
50 if (targ_laddr_hist)
51 key = inet->inet_saddr;
52 else if (targ_raddr_hist)
53 key = inet->sk.__sk_common.skc_daddr;
54 else
55 key = 0;
56 histp = bpf_map_lookup_or_try_init(&hists, &key, &zero);
57 if (!histp)
58 return 0;
59 ts = (struct tcp_sock *)(sk);
60 srtt = BPF_CORE_READ(ts, srtt_us) >> 3;
61 if (targ_ms)
62 srtt /= 1000U;
63 slot = log2l(srtt);
64 if (slot >= MAX_SLOTS)
65 slot = MAX_SLOTS - 1;
66 __sync_fetch_and_add(&histp->slots[slot], 1);
67 if (targ_show_ext) {
68 __sync_fetch_and_add(&histp->latency, srtt);
69 __sync_fetch_and_add(&histp->cnt, 1);
70 }
71 return 0;
72 }
73
74 SEC("kprobe/tcp_rcv_established")
BPF_KPROBE(tcp_rcv_kprobe,struct sock * sk)75 int BPF_KPROBE(tcp_rcv_kprobe, struct sock *sk)
76 {
77 const struct inet_sock *inet = (struct inet_sock *)(sk);
78 u32 srtt, saddr, daddr;
79 struct tcp_sock *ts;
80 struct hist *histp;
81 u64 key, slot;
82
83 if (targ_sport) {
84 u16 sport;
85 bpf_probe_read_kernel(&sport, sizeof(sport), &inet->inet_sport);
86 if (targ_sport != sport)
87 return 0;
88 }
89 if (targ_dport) {
90 u16 dport;
91 bpf_probe_read_kernel(&dport, sizeof(dport), &sk->__sk_common.skc_dport);
92 if (targ_dport != dport)
93 return 0;
94 }
95 bpf_probe_read_kernel(&saddr, sizeof(saddr), &inet->inet_saddr);
96 if (targ_saddr && targ_saddr != saddr)
97 return 0;
98 bpf_probe_read_kernel(&daddr, sizeof(daddr), &sk->__sk_common.skc_daddr);
99 if (targ_daddr && targ_daddr != daddr)
100 return 0;
101
102 if (targ_laddr_hist)
103 key = saddr;
104 else if (targ_raddr_hist)
105 key = daddr;
106 else
107 key = 0;
108 histp = bpf_map_lookup_or_try_init(&hists, &key, &zero);
109 if (!histp)
110 return 0;
111 ts = (struct tcp_sock *)(sk);
112 bpf_probe_read_kernel(&srtt, sizeof(srtt), &ts->srtt_us);
113 srtt >>= 3;
114 if (targ_ms)
115 srtt /= 1000U;
116 slot = log2l(srtt);
117 if (slot >= MAX_SLOTS)
118 slot = MAX_SLOTS - 1;
119 __sync_fetch_and_add(&histp->slots[slot], 1);
120 if (targ_show_ext) {
121 __sync_fetch_and_add(&histp->latency, srtt);
122 __sync_fetch_and_add(&histp->cnt, 1);
123 }
124 return 0;
125 }
126
127 char LICENSE[] SEC("license") = "Dual BSD/GPL";
128