1 // SPDX-License-Identifier: GPL-2.0
2 // Copyright (c) 2023 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
8 #include "compat.bpf.h"
9 #include "core_fixes.bpf.h"
10 #include "tcppktlat.h"
11
12 #define MAX_ENTRIES 10240
13 #define AF_INET 2
14
15 const volatile pid_t targ_pid = 0;
16 const volatile pid_t targ_tid = 0;
17 const volatile __u16 targ_sport = 0;
18 const volatile __u16 targ_dport = 0;
19 const volatile __u64 targ_min_us = 0;
20
21 struct {
22 __uint(type, BPF_MAP_TYPE_HASH);
23 __uint(max_entries, MAX_ENTRIES);
24 __type(key, u64);
25 __type(value, u64);
26 } start SEC(".maps");
27
handle_tcp_probe(struct sock * sk,struct sk_buff * skb)28 static int handle_tcp_probe(struct sock *sk, struct sk_buff *skb)
29 {
30 const struct inet_sock *inet = (struct inet_sock *)(sk);
31 u64 sock_ident, ts, len, doff;
32 const struct tcphdr *th;
33
34 if (targ_sport && targ_sport != BPF_CORE_READ(inet, inet_sport))
35 return 0;
36 if (targ_dport && targ_dport != BPF_CORE_READ(sk, __sk_common.skc_dport))
37 return 0;
38 th = (const struct tcphdr*)BPF_CORE_READ(skb, data);
39 doff = BPF_CORE_READ_BITFIELD_PROBED(th, doff);
40 len = BPF_CORE_READ(skb, len);
41 /* `doff * 4` means `__tcp_hdrlen` */
42 if (len <= doff * 4)
43 return 0;
44 sock_ident = get_sock_ident(sk);
45 ts = bpf_ktime_get_ns();
46 bpf_map_update_elem(&start, &sock_ident, &ts, 0);
47 return 0;
48 }
49
handle_tcp_rcv_space_adjust(void * ctx,struct sock * sk)50 static int handle_tcp_rcv_space_adjust(void *ctx, struct sock *sk)
51 {
52 const struct inet_sock *inet = (struct inet_sock *)(sk);
53 u64 sock_ident = get_sock_ident(sk);
54 u64 id = bpf_get_current_pid_tgid(), *tsp;
55 u32 pid = id >> 32, tid = id;
56 struct event *eventp;
57 s64 delta_us;
58 u16 family;
59
60 tsp = bpf_map_lookup_elem(&start, &sock_ident);
61 if (!tsp)
62 return 0;
63
64 if (targ_pid && targ_pid != pid)
65 goto cleanup;
66 if (targ_tid && targ_tid != tid)
67 goto cleanup;
68
69 delta_us = (bpf_ktime_get_ns() - *tsp) / 1000;
70 if (delta_us < 0 || delta_us <= targ_min_us)
71 goto cleanup;
72
73 eventp = reserve_buf(sizeof(*eventp));
74 if (!eventp)
75 goto cleanup;
76
77 eventp->pid = pid;
78 eventp->tid = tid;
79 eventp->delta_us = delta_us;
80 eventp->sport = BPF_CORE_READ(inet, inet_sport);
81 eventp->dport = BPF_CORE_READ(sk, __sk_common.skc_dport);
82 bpf_get_current_comm(&eventp->comm, TASK_COMM_LEN);
83 family = BPF_CORE_READ(sk, __sk_common.skc_family);
84 if (family == AF_INET) {
85 eventp->saddr[0] = BPF_CORE_READ(sk, __sk_common.skc_rcv_saddr);
86 eventp->daddr[0] = BPF_CORE_READ(sk, __sk_common.skc_daddr);
87 } else { /* family == AF_INET6 */
88 BPF_CORE_READ_INTO(eventp->saddr, sk, __sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
89 BPF_CORE_READ_INTO(eventp->daddr, sk, __sk_common.skc_v6_daddr.in6_u.u6_addr32);
90 }
91 eventp->family = family;
92 submit_buf(ctx, eventp, sizeof(*eventp));
93
94 cleanup:
95 bpf_map_delete_elem(&start, &sock_ident);
96 return 0;
97 }
98
handle_tcp_destroy_sock(void * ctx,struct sock * sk)99 static int handle_tcp_destroy_sock(void *ctx, struct sock *sk)
100 {
101 u64 sock_ident = get_sock_ident(sk);
102
103 bpf_map_delete_elem(&start, &sock_ident);
104 return 0;
105 }
106
107 SEC("tp_btf/tcp_probe")
BPF_PROG(tcp_probe_btf,struct sock * sk,struct sk_buff * skb)108 int BPF_PROG(tcp_probe_btf, struct sock *sk, struct sk_buff *skb)
109 {
110 return handle_tcp_probe(sk, skb);
111 }
112
113 SEC("tp_btf/tcp_rcv_space_adjust")
BPF_PROG(tcp_rcv_space_adjust_btf,struct sock * sk)114 int BPF_PROG(tcp_rcv_space_adjust_btf, struct sock *sk)
115 {
116 return handle_tcp_rcv_space_adjust(ctx, sk);
117 }
118
119 SEC("tp_btf/tcp_destroy_sock")
BPF_PROG(tcp_destroy_sock_btf,struct sock * sk)120 int BPF_PROG(tcp_destroy_sock_btf, struct sock *sk)
121 {
122 return handle_tcp_destroy_sock(ctx, sk);
123 }
124
125 SEC("raw_tp/tcp_probe")
BPF_PROG(tcp_probe,struct sock * sk,struct sk_buff * skb)126 int BPF_PROG(tcp_probe, struct sock *sk, struct sk_buff *skb) {
127 return handle_tcp_probe(sk, skb);
128 }
129
130 SEC("raw_tp/tcp_rcv_space_adjust")
BPF_PROG(tcp_rcv_space_adjust,struct sock * sk)131 int BPF_PROG(tcp_rcv_space_adjust, struct sock *sk)
132 {
133 return handle_tcp_rcv_space_adjust(ctx, sk);
134 }
135
136 SEC("raw_tp/tcp_destroy_sock")
BPF_PROG(tcp_destroy_sock,struct sock * sk)137 int BPF_PROG(tcp_destroy_sock, struct sock *sk)
138 {
139 return handle_tcp_destroy_sock(ctx, sk);
140 }
141
142 char LICENSE[] SEC("license") = "GPL";
143