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
10 static const int nr_soreuse = 4;
11
do_test(int sock_type,bool onebyone)12 static 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
102 done:
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
test_sock_iter_batch(void)111 void 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
133 done:
134 SYS_NOFAIL("ip netns del " TEST_NS);
135 }
136