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}