1// SPDX-License-Identifier: GPL-2.0
2// Copyright (c) 2024 Meta
3
4#include <test_progs.h>
5#include "network_helpers.h"
6#include "sock_iter_batch.skel.h"
7
8#define TEST_NS "sock_iter_batch_netns"
9
10static const int nr_soreuse = 4;
11
12static void do_test(int sock_type, bool onebyone)
13{
14	int err, i, nread, to_read, total_read, iter_fd = -1;
15	int first_idx, second_idx, indices[nr_soreuse];
16	struct bpf_link *link = NULL;
17	struct sock_iter_batch *skel;
18	int *fds[2] = {};
19
20	skel = sock_iter_batch__open();
21	if (!ASSERT_OK_PTR(skel, "sock_iter_batch__open"))
22		return;
23
24	/* Prepare 2 buckets of sockets in the kernel hashtable */
25	for (i = 0; i < ARRAY_SIZE(fds); i++) {
26		int local_port;
27
28		fds[i] = start_reuseport_server(AF_INET6, sock_type, "::1", 0, 0,
29						nr_soreuse);
30		if (!ASSERT_OK_PTR(fds[i], "start_reuseport_server"))
31			goto done;
32		local_port = get_socket_local_port(*fds[i]);
33		if (!ASSERT_GE(local_port, 0, "get_socket_local_port"))
34			goto done;
35		skel->rodata->ports[i] = ntohs(local_port);
36	}
37
38	err = sock_iter_batch__load(skel);
39	if (!ASSERT_OK(err, "sock_iter_batch__load"))
40		goto done;
41
42	link = bpf_program__attach_iter(sock_type == SOCK_STREAM ?
43					skel->progs.iter_tcp_soreuse :
44					skel->progs.iter_udp_soreuse,
45					NULL);
46	if (!ASSERT_OK_PTR(link, "bpf_program__attach_iter"))
47		goto done;
48
49	iter_fd = bpf_iter_create(bpf_link__fd(link));
50	if (!ASSERT_GE(iter_fd, 0, "bpf_iter_create"))
51		goto done;
52
53	/* Test reading a bucket (either from fds[0] or fds[1]).
54	 * Only read "nr_soreuse - 1" number of sockets
55	 * from a bucket and leave one socket out from
56	 * that bucket on purpose.
57	 */
58	to_read = (nr_soreuse - 1) * sizeof(*indices);
59	total_read = 0;
60	first_idx = -1;
61	do {
62		nread = read(iter_fd, indices, onebyone ? sizeof(*indices) : to_read);
63		if (nread <= 0 || nread % sizeof(*indices))
64			break;
65		total_read += nread;
66
67		if (first_idx == -1)
68			first_idx = indices[0];
69		for (i = 0; i < nread / sizeof(*indices); i++)
70			ASSERT_EQ(indices[i], first_idx, "first_idx");
71	} while (total_read < to_read);
72	ASSERT_EQ(nread, onebyone ? sizeof(*indices) : to_read, "nread");
73	ASSERT_EQ(total_read, to_read, "total_read");
74
75	free_fds(fds[first_idx], nr_soreuse);
76	fds[first_idx] = NULL;
77
78	/* Read the "whole" second bucket */
79	to_read = nr_soreuse * sizeof(*indices);
80	total_read = 0;
81	second_idx = !first_idx;
82	do {
83		nread = read(iter_fd, indices, onebyone ? sizeof(*indices) : to_read);
84		if (nread <= 0 || nread % sizeof(*indices))
85			break;
86		total_read += nread;
87
88		for (i = 0; i < nread / sizeof(*indices); i++)
89			ASSERT_EQ(indices[i], second_idx, "second_idx");
90	} while (total_read <= to_read);
91	ASSERT_EQ(nread, 0, "nread");
92	/* Both so_reuseport ports should be in different buckets, so
93	 * total_read must equal to the expected to_read.
94	 *
95	 * For a very unlikely case, both ports collide at the same bucket,
96	 * the bucket offset (i.e. 3) will be skipped and it cannot
97	 * expect the to_read number of bytes.
98	 */
99	if (skel->bss->bucket[0] != skel->bss->bucket[1])
100		ASSERT_EQ(total_read, to_read, "total_read");
101
102done:
103	for (i = 0; i < ARRAY_SIZE(fds); i++)
104		free_fds(fds[i], nr_soreuse);
105	if (iter_fd < 0)
106		close(iter_fd);
107	bpf_link__destroy(link);
108	sock_iter_batch__destroy(skel);
109}
110
111void test_sock_iter_batch(void)
112{
113	struct nstoken *nstoken = NULL;
114
115	SYS_NOFAIL("ip netns del " TEST_NS);
116	SYS(done, "ip netns add %s", TEST_NS);
117	SYS(done, "ip -net %s link set dev lo up", TEST_NS);
118
119	nstoken = open_netns(TEST_NS);
120	if (!ASSERT_OK_PTR(nstoken, "open_netns"))
121		goto done;
122
123	if (test__start_subtest("tcp")) {
124		do_test(SOCK_STREAM, true);
125		do_test(SOCK_STREAM, false);
126	}
127	if (test__start_subtest("udp")) {
128		do_test(SOCK_DGRAM, true);
129		do_test(SOCK_DGRAM, false);
130	}
131	close_netns(nstoken);
132
133done:
134	SYS_NOFAIL("ip netns del " TEST_NS);
135}
136