1 // SPDX-License-Identifier: GPL-2.0
2
3 #include "vmlinux.h"
4 #include <bpf/bpf_helpers.h>
5 #include <bpf/bpf_endian.h>
6
7 #include "bpf_tracing_net.h"
8
9 __be16 serv_port = 0;
10 int tw_found = 0;
11 int tw_destroy_err = 0;
12
13 int bpf_sock_destroy(struct sock_common *sk) __ksym;
14
15 struct {
16 __uint(type, BPF_MAP_TYPE_ARRAY);
17 __uint(max_entries, 1);
18 __type(key, __u32);
19 __type(value, __u64);
20 } tcp_conn_sockets SEC(".maps");
21
22 struct {
23 __uint(type, BPF_MAP_TYPE_ARRAY);
24 __uint(max_entries, 1);
25 __type(key, __u32);
26 __type(value, __u64);
27 } udp_conn_sockets SEC(".maps");
28
29 SEC("cgroup/connect6")
sock_connect(struct bpf_sock_addr * ctx)30 int sock_connect(struct bpf_sock_addr *ctx)
31 {
32 __u64 sock_cookie = 0;
33 int key = 0;
34 __u32 keyc = 0;
35
36 if (ctx->family != AF_INET6 || ctx->user_family != AF_INET6)
37 return 1;
38
39 sock_cookie = bpf_get_socket_cookie(ctx);
40 if (ctx->protocol == IPPROTO_TCP)
41 bpf_map_update_elem(&tcp_conn_sockets, &key, &sock_cookie, 0);
42 else if (ctx->protocol == IPPROTO_UDP)
43 bpf_map_update_elem(&udp_conn_sockets, &keyc, &sock_cookie, 0);
44 else
45 return 1;
46
47 return 1;
48 }
49
50 SEC("iter/tcp")
iter_tcp6_client(struct bpf_iter__tcp * ctx)51 int iter_tcp6_client(struct bpf_iter__tcp *ctx)
52 {
53 struct sock_common *sk_common = ctx->sk_common;
54 __u64 sock_cookie = 0;
55 __u64 *val;
56 int key = 0;
57
58 if (!sk_common)
59 return 0;
60
61 if (sk_common->skc_family != AF_INET6)
62 return 0;
63
64 sock_cookie = bpf_get_socket_cookie(sk_common);
65 val = bpf_map_lookup_elem(&tcp_conn_sockets, &key);
66 if (!val)
67 return 0;
68 /* Destroy connected client sockets. */
69 if (sock_cookie == *val)
70 bpf_sock_destroy(sk_common);
71
72 return 0;
73 }
74
75 SEC("iter/tcp")
iter_tcp6_server(struct bpf_iter__tcp * ctx)76 int iter_tcp6_server(struct bpf_iter__tcp *ctx)
77 {
78 struct sock_common *sk_common = ctx->sk_common;
79 const struct inet_connection_sock *icsk;
80 const struct inet_sock *inet;
81 struct tcp6_sock *tcp_sk;
82 __be16 srcp;
83
84 if (!sk_common)
85 return 0;
86
87 if (sk_common->skc_family != AF_INET6)
88 return 0;
89
90 tcp_sk = bpf_skc_to_tcp6_sock(sk_common);
91 if (!tcp_sk)
92 return 0;
93
94 icsk = &tcp_sk->tcp.inet_conn;
95 inet = &icsk->icsk_inet;
96 srcp = inet->inet_sport;
97
98 /* Destroy server sockets. */
99 if (srcp == serv_port)
100 bpf_sock_destroy(sk_common);
101
102 return 0;
103 }
104
105 SEC("iter/tcp")
iter_tcp6_timewait(struct bpf_iter__tcp * ctx)106 int iter_tcp6_timewait(struct bpf_iter__tcp *ctx)
107 {
108 struct sock_common *sk_common = ctx->sk_common;
109 __u64 *val;
110 int key = 0;
111
112 if (!sk_common)
113 return 0;
114
115 if (sk_common->skc_family != AF_INET6)
116 return 0;
117
118 if (!bpf_skc_to_tcp_timewait_sock(sk_common))
119 return 0;
120
121 val = bpf_map_lookup_elem(&tcp_conn_sockets, &key);
122 if (!val)
123 return 0;
124 /* The timewait sock inherits the cookie of the closed client sock. */
125 if (bpf_get_socket_cookie(sk_common) != *val)
126 return 0;
127
128 tw_found++;
129 tw_destroy_err = bpf_sock_destroy(sk_common);
130
131 return 0;
132 }
133
134 SEC("iter/udp")
iter_udp6_client(struct bpf_iter__udp * ctx)135 int iter_udp6_client(struct bpf_iter__udp *ctx)
136 {
137 struct udp_sock *udp_sk = ctx->udp_sk;
138 struct sock *sk = (struct sock *) udp_sk;
139 __u64 sock_cookie = 0, *val;
140 int key = 0;
141
142 if (!sk)
143 return 0;
144
145 sock_cookie = bpf_get_socket_cookie(sk);
146 val = bpf_map_lookup_elem(&udp_conn_sockets, &key);
147 if (!val)
148 return 0;
149 /* Destroy connected client sockets. */
150 if (sock_cookie == *val)
151 bpf_sock_destroy((struct sock_common *)sk);
152
153 return 0;
154 }
155
156 SEC("iter/udp")
iter_udp6_server(struct bpf_iter__udp * ctx)157 int iter_udp6_server(struct bpf_iter__udp *ctx)
158 {
159 struct udp_sock *udp_sk = ctx->udp_sk;
160 struct sock *sk = (struct sock *) udp_sk;
161 struct inet_sock *inet;
162 __be16 srcp;
163
164 if (!sk)
165 return 0;
166
167 inet = &udp_sk->inet;
168 srcp = inet->inet_sport;
169 if (srcp == serv_port)
170 bpf_sock_destroy((struct sock_common *)sk);
171
172 return 0;
173 }
174
175 char _license[] SEC("license") = "GPL";
176