1// SPDX-License-Identifier: GPL-2.0
2/* Copyright (c) 2023 Meta Platforms, Inc. and affiliates. */
3#include <linux/bpf.h>
4#include <bpf/bpf_endian.h>
5#include <bpf/bpf_helpers.h>
6
7#include <linux/if_ether.h>
8#include <linux/in.h>
9#include <linux/in6.h>
10#include <linux/ipv6.h>
11#include <linux/tcp.h>
12
13#include <sys/types.h>
14#include <sys/socket.h>
15
16#include "cgroup_tcp_skb.h"
17
18char _license[] SEC("license") = "GPL";
19
20__u16 g_sock_port = 0;
21__u32 g_sock_state = 0;
22int g_unexpected = 0;
23__u32 g_packet_count = 0;
24
25int needed_tcp_pkt(struct __sk_buff *skb, struct tcphdr *tcph)
26{
27	struct ipv6hdr ip6h;
28
29	if (skb->protocol != bpf_htons(ETH_P_IPV6))
30		return 0;
31	if (bpf_skb_load_bytes(skb, 0, &ip6h, sizeof(ip6h)))
32		return 0;
33
34	if (ip6h.nexthdr != IPPROTO_TCP)
35		return 0;
36
37	if (bpf_skb_load_bytes(skb, sizeof(ip6h), tcph, sizeof(*tcph)))
38		return 0;
39
40	if (tcph->source != bpf_htons(g_sock_port) &&
41	    tcph->dest != bpf_htons(g_sock_port))
42		return 0;
43
44	return 1;
45}
46
47/* Run accept() on a socket in the cgroup to receive a new connection. */
48static int egress_accept(struct tcphdr *tcph)
49{
50	if (g_sock_state ==  SYN_RECV_SENDING_SYN_ACK) {
51		if (tcph->fin || !tcph->syn || !tcph->ack)
52			g_unexpected++;
53		else
54			g_sock_state = SYN_RECV;
55		return 1;
56	}
57
58	return 0;
59}
60
61static int ingress_accept(struct tcphdr *tcph)
62{
63	switch (g_sock_state) {
64	case INIT:
65		if (!tcph->syn || tcph->fin || tcph->ack)
66			g_unexpected++;
67		else
68			g_sock_state = SYN_RECV_SENDING_SYN_ACK;
69		break;
70	case SYN_RECV:
71		if (tcph->fin || tcph->syn || !tcph->ack)
72			g_unexpected++;
73		else
74			g_sock_state = ESTABLISHED;
75		break;
76	default:
77		return 0;
78	}
79
80	return 1;
81}
82
83/* Run connect() on a socket in the cgroup to start a new connection. */
84static int egress_connect(struct tcphdr *tcph)
85{
86	if (g_sock_state == INIT) {
87		if (!tcph->syn || tcph->fin || tcph->ack)
88			g_unexpected++;
89		else
90			g_sock_state = SYN_SENT;
91		return 1;
92	}
93
94	return 0;
95}
96
97static int ingress_connect(struct tcphdr *tcph)
98{
99	if (g_sock_state == SYN_SENT) {
100		if (tcph->fin || !tcph->syn || !tcph->ack)
101			g_unexpected++;
102		else
103			g_sock_state = ESTABLISHED;
104		return 1;
105	}
106
107	return 0;
108}
109
110/* The connection is closed by the peer outside the cgroup. */
111static int egress_close_remote(struct tcphdr *tcph)
112{
113	switch (g_sock_state) {
114	case ESTABLISHED:
115		break;
116	case CLOSE_WAIT_SENDING_ACK:
117		if (tcph->fin || tcph->syn || !tcph->ack)
118			g_unexpected++;
119		else
120			g_sock_state = CLOSE_WAIT;
121		break;
122	case CLOSE_WAIT:
123		if (!tcph->fin)
124			g_unexpected++;
125		else
126			g_sock_state = LAST_ACK;
127		break;
128	default:
129		return 0;
130	}
131
132	return 1;
133}
134
135static int ingress_close_remote(struct tcphdr *tcph)
136{
137	switch (g_sock_state) {
138	case ESTABLISHED:
139		if (tcph->fin)
140			g_sock_state = CLOSE_WAIT_SENDING_ACK;
141		break;
142	case LAST_ACK:
143		if (tcph->fin || tcph->syn || !tcph->ack)
144			g_unexpected++;
145		else
146			g_sock_state = CLOSED;
147		break;
148	default:
149		return 0;
150	}
151
152	return 1;
153}
154
155/* The connection is closed by the endpoint inside the cgroup. */
156static int egress_close_local(struct tcphdr *tcph)
157{
158	switch (g_sock_state) {
159	case ESTABLISHED:
160		if (tcph->fin)
161			g_sock_state = FIN_WAIT1;
162		break;
163	case TIME_WAIT_SENDING_ACK:
164		if (tcph->fin || tcph->syn || !tcph->ack)
165			g_unexpected++;
166		else
167			g_sock_state = TIME_WAIT;
168		break;
169	default:
170		return 0;
171	}
172
173	return 1;
174}
175
176static int ingress_close_local(struct tcphdr *tcph)
177{
178	switch (g_sock_state) {
179	case ESTABLISHED:
180		break;
181	case FIN_WAIT1:
182		if (tcph->fin || tcph->syn || !tcph->ack)
183			g_unexpected++;
184		else
185			g_sock_state = FIN_WAIT2;
186		break;
187	case FIN_WAIT2:
188		if (!tcph->fin || tcph->syn || !tcph->ack)
189			g_unexpected++;
190		else
191			g_sock_state = TIME_WAIT_SENDING_ACK;
192		break;
193	default:
194		return 0;
195	}
196
197	return 1;
198}
199
200/* Check the types of outgoing packets of a server socket to make sure they
201 * are consistent with the state of the server socket.
202 *
203 * The connection is closed by the client side.
204 */
205SEC("cgroup_skb/egress")
206int server_egress(struct __sk_buff *skb)
207{
208	struct tcphdr tcph;
209
210	if (!needed_tcp_pkt(skb, &tcph))
211		return 1;
212
213	g_packet_count++;
214
215	/* Egress of the server socket. */
216	if (egress_accept(&tcph) || egress_close_remote(&tcph))
217		return 1;
218
219	g_unexpected++;
220	return 1;
221}
222
223/* Check the types of incoming packets of a server socket to make sure they
224 * are consistent with the state of the server socket.
225 *
226 * The connection is closed by the client side.
227 */
228SEC("cgroup_skb/ingress")
229int server_ingress(struct __sk_buff *skb)
230{
231	struct tcphdr tcph;
232
233	if (!needed_tcp_pkt(skb, &tcph))
234		return 1;
235
236	g_packet_count++;
237
238	/* Ingress of the server socket. */
239	if (ingress_accept(&tcph) || ingress_close_remote(&tcph))
240		return 1;
241
242	g_unexpected++;
243	return 1;
244}
245
246/* Check the types of outgoing packets of a server socket to make sure they
247 * are consistent with the state of the server socket.
248 *
249 * The connection is closed by the server side.
250 */
251SEC("cgroup_skb/egress")
252int server_egress_srv(struct __sk_buff *skb)
253{
254	struct tcphdr tcph;
255
256	if (!needed_tcp_pkt(skb, &tcph))
257		return 1;
258
259	g_packet_count++;
260
261	/* Egress of the server socket. */
262	if (egress_accept(&tcph) || egress_close_local(&tcph))
263		return 1;
264
265	g_unexpected++;
266	return 1;
267}
268
269/* Check the types of incoming packets of a server socket to make sure they
270 * are consistent with the state of the server socket.
271 *
272 * The connection is closed by the server side.
273 */
274SEC("cgroup_skb/ingress")
275int server_ingress_srv(struct __sk_buff *skb)
276{
277	struct tcphdr tcph;
278
279	if (!needed_tcp_pkt(skb, &tcph))
280		return 1;
281
282	g_packet_count++;
283
284	/* Ingress of the server socket. */
285	if (ingress_accept(&tcph) || ingress_close_local(&tcph))
286		return 1;
287
288	g_unexpected++;
289	return 1;
290}
291
292/* Check the types of outgoing packets of a client socket to make sure they
293 * are consistent with the state of the client socket.
294 *
295 * The connection is closed by the server side.
296 */
297SEC("cgroup_skb/egress")
298int client_egress_srv(struct __sk_buff *skb)
299{
300	struct tcphdr tcph;
301
302	if (!needed_tcp_pkt(skb, &tcph))
303		return 1;
304
305	g_packet_count++;
306
307	/* Egress of the server socket. */
308	if (egress_connect(&tcph) || egress_close_remote(&tcph))
309		return 1;
310
311	g_unexpected++;
312	return 1;
313}
314
315/* Check the types of incoming packets of a client socket to make sure they
316 * are consistent with the state of the client socket.
317 *
318 * The connection is closed by the server side.
319 */
320SEC("cgroup_skb/ingress")
321int client_ingress_srv(struct __sk_buff *skb)
322{
323	struct tcphdr tcph;
324
325	if (!needed_tcp_pkt(skb, &tcph))
326		return 1;
327
328	g_packet_count++;
329
330	/* Ingress of the server socket. */
331	if (ingress_connect(&tcph) || ingress_close_remote(&tcph))
332		return 1;
333
334	g_unexpected++;
335	return 1;
336}
337
338/* Check the types of outgoing packets of a client socket to make sure they
339 * are consistent with the state of the client socket.
340 *
341 * The connection is closed by the client side.
342 */
343SEC("cgroup_skb/egress")
344int client_egress(struct __sk_buff *skb)
345{
346	struct tcphdr tcph;
347
348	if (!needed_tcp_pkt(skb, &tcph))
349		return 1;
350
351	g_packet_count++;
352
353	/* Egress of the server socket. */
354	if (egress_connect(&tcph) || egress_close_local(&tcph))
355		return 1;
356
357	g_unexpected++;
358	return 1;
359}
360
361/* Check the types of incoming packets of a client socket to make sure they
362 * are consistent with the state of the client socket.
363 *
364 * The connection is closed by the client side.
365 */
366SEC("cgroup_skb/ingress")
367int client_ingress(struct __sk_buff *skb)
368{
369	struct tcphdr tcph;
370
371	if (!needed_tcp_pkt(skb, &tcph))
372		return 1;
373
374	g_packet_count++;
375
376	/* Ingress of the server socket. */
377	if (ingress_connect(&tcph) || ingress_close_local(&tcph))
378		return 1;
379
380	g_unexpected++;
381	return 1;
382}
383