1 // SPDX-License-Identifier: GPL-2.0
2 // Copyright (c) 2022 Francis Laniel <[email protected]>
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
9 #include "tcptop.h"
10
11 /* Taken from kernel include/linux/socket.h. */
12 #define AF_INET 2 /* Internet IP Protocol */
13 #define AF_INET6 10 /* IP version 6 */
14
15 const volatile bool filter_cg = false;
16 const volatile pid_t target_pid = -1;
17 const volatile int target_family = -1;
18
19 struct {
20 __uint(type, BPF_MAP_TYPE_CGROUP_ARRAY);
21 __type(key, u32);
22 __type(value, u32);
23 __uint(max_entries, 1);
24 } cgroup_map SEC(".maps");
25
26 struct {
27 __uint(type, BPF_MAP_TYPE_HASH);
28 __uint(max_entries, 10240);
29 __type(key, struct ip_key_t);
30 __type(value, struct traffic_t);
31 } ip_map SEC(".maps");
32
probe_ip(bool receiving,struct sock * sk,size_t size)33 static int probe_ip(bool receiving, struct sock *sk, size_t size)
34 {
35 struct ip_key_t ip_key = {};
36 struct traffic_t *trafficp;
37 u16 family;
38 u32 pid;
39
40 if (filter_cg && !bpf_current_task_under_cgroup(&cgroup_map, 0))
41 return 0;
42
43 pid = bpf_get_current_pid_tgid() >> 32;
44 if (target_pid != -1 && target_pid != pid)
45 return 0;
46
47 family = BPF_CORE_READ(sk, __sk_common.skc_family);
48 if (target_family != -1 && target_family != family)
49 return 0;
50
51 /* drop */
52 if (family != AF_INET && family != AF_INET6)
53 return 0;
54
55 ip_key.pid = pid;
56 bpf_get_current_comm(&ip_key.name, sizeof(ip_key.name));
57 ip_key.lport = BPF_CORE_READ(sk, __sk_common.skc_num);
58 ip_key.dport = bpf_ntohs(BPF_CORE_READ(sk, __sk_common.skc_dport));
59 ip_key.family = family;
60
61 if (family == AF_INET) {
62 bpf_probe_read_kernel(&ip_key.saddr,
63 sizeof(sk->__sk_common.skc_rcv_saddr),
64 &sk->__sk_common.skc_rcv_saddr);
65 bpf_probe_read_kernel(&ip_key.daddr,
66 sizeof(sk->__sk_common.skc_daddr),
67 &sk->__sk_common.skc_daddr);
68 } else {
69 /*
70 * family == AF_INET6,
71 * we already checked above family is correct.
72 */
73 bpf_probe_read_kernel(&ip_key.saddr,
74 sizeof(sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32),
75 &sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
76 bpf_probe_read_kernel(&ip_key.daddr,
77 sizeof(sk->__sk_common.skc_v6_daddr.in6_u.u6_addr32),
78 &sk->__sk_common.skc_v6_daddr.in6_u.u6_addr32);
79 }
80
81 trafficp = bpf_map_lookup_elem(&ip_map, &ip_key);
82 if (!trafficp) {
83 struct traffic_t zero;
84
85 if (receiving) {
86 zero.sent = 0;
87 zero.received = size;
88 } else {
89 zero.sent = size;
90 zero.received = 0;
91 }
92
93 bpf_map_update_elem(&ip_map, &ip_key, &zero, BPF_NOEXIST);
94 } else {
95 if (receiving)
96 trafficp->received += size;
97 else
98 trafficp->sent += size;
99
100 bpf_map_update_elem(&ip_map, &ip_key, trafficp, BPF_EXIST);
101 }
102
103 return 0;
104 }
105
106 SEC("kprobe/tcp_sendmsg")
BPF_KPROBE(tcp_sendmsg,struct sock * sk,struct msghdr * msg,size_t size)107 int BPF_KPROBE(tcp_sendmsg, struct sock *sk, struct msghdr *msg, size_t size)
108 {
109 return probe_ip(false, sk, size);
110 }
111
112 /*
113 * tcp_recvmsg() would be obvious to trace, but is less suitable because:
114 * - we'd need to trace both entry and return, to have both sock and size
115 * - misses tcp_read_sock() traffic
116 * we'd much prefer tracepoints once they are available.
117 */
118 SEC("kprobe/tcp_cleanup_rbuf")
BPF_KPROBE(tcp_cleanup_rbuf,struct sock * sk,int copied)119 int BPF_KPROBE(tcp_cleanup_rbuf, struct sock *sk, int copied)
120 {
121 if (copied <= 0)
122 return 0;
123
124 return probe_ip(true, sk, copied);
125 }
126
127 char LICENSE[] SEC("license") = "GPL";
128