136 lines · c
1// SPDX-License-Identifier: GPL-2.02// Copyright (c) 2024 Meta3 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 sockets55 * from a bucket and leave one socket out from56 * 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, so93 * 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 cannot97 * 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