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 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 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