1// SPDX-License-Identifier: GPL-2.0
  2/* Copyright Amazon.com Inc. or its affiliates. */
  3#define _GNU_SOURCE
  4#include <sched.h>
  5
  6#include <fcntl.h>
  7
  8#include <netinet/in.h>
  9#include <sys/socket.h>
 10#include <sys/sysinfo.h>
 11
 12#include "../kselftest_harness.h"
 13
 14FIXTURE(so_incoming_cpu)
 15{
 16	int *servers;
 17	union {
 18		struct sockaddr addr;
 19		struct sockaddr_in in_addr;
 20	};
 21	socklen_t addrlen;
 22};
 23
 24enum when_to_set {
 25	BEFORE_REUSEPORT,
 26	BEFORE_LISTEN,
 27	AFTER_LISTEN,
 28	AFTER_ALL_LISTEN,
 29};
 30
 31FIXTURE_VARIANT(so_incoming_cpu)
 32{
 33	int when_to_set;
 34};
 35
 36FIXTURE_VARIANT_ADD(so_incoming_cpu, before_reuseport)
 37{
 38	.when_to_set = BEFORE_REUSEPORT,
 39};
 40
 41FIXTURE_VARIANT_ADD(so_incoming_cpu, before_listen)
 42{
 43	.when_to_set = BEFORE_LISTEN,
 44};
 45
 46FIXTURE_VARIANT_ADD(so_incoming_cpu, after_listen)
 47{
 48	.when_to_set = AFTER_LISTEN,
 49};
 50
 51FIXTURE_VARIANT_ADD(so_incoming_cpu, after_all_listen)
 52{
 53	.when_to_set = AFTER_ALL_LISTEN,
 54};
 55
 56static void write_sysctl(struct __test_metadata *_metadata,
 57			 char *filename, char *string)
 58{
 59	int fd, len, ret;
 60
 61	fd = open(filename, O_WRONLY);
 62	ASSERT_NE(fd, -1);
 63
 64	len = strlen(string);
 65	ret = write(fd, string, len);
 66	ASSERT_EQ(ret, len);
 67}
 68
 69static void setup_netns(struct __test_metadata *_metadata)
 70{
 71	ASSERT_EQ(unshare(CLONE_NEWNET), 0);
 72	ASSERT_EQ(system("ip link set lo up"), 0);
 73
 74	write_sysctl(_metadata, "/proc/sys/net/ipv4/ip_local_port_range", "10000 60001");
 75	write_sysctl(_metadata, "/proc/sys/net/ipv4/tcp_tw_reuse", "0");
 76}
 77
 78#define NR_PORT				(60001 - 10000 - 1)
 79#define NR_CLIENT_PER_SERVER_DEFAULT	32
 80static int nr_client_per_server, nr_server, nr_client;
 81
 82FIXTURE_SETUP(so_incoming_cpu)
 83{
 84	setup_netns(_metadata);
 85
 86	nr_server = get_nprocs();
 87	ASSERT_LE(2, nr_server);
 88
 89	if (NR_CLIENT_PER_SERVER_DEFAULT * nr_server < NR_PORT)
 90		nr_client_per_server = NR_CLIENT_PER_SERVER_DEFAULT;
 91	else
 92		nr_client_per_server = NR_PORT / nr_server;
 93
 94	nr_client = nr_client_per_server * nr_server;
 95
 96	self->servers = malloc(sizeof(int) * nr_server);
 97	ASSERT_NE(self->servers, NULL);
 98
 99	self->in_addr.sin_family = AF_INET;
100	self->in_addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
101	self->in_addr.sin_port = htons(0);
102	self->addrlen = sizeof(struct sockaddr_in);
103}
104
105FIXTURE_TEARDOWN(so_incoming_cpu)
106{
107	int i;
108
109	for (i = 0; i < nr_server; i++)
110		close(self->servers[i]);
111
112	free(self->servers);
113}
114
115void set_so_incoming_cpu(struct __test_metadata *_metadata, int fd, int cpu)
116{
117	int ret;
118
119	ret = setsockopt(fd, SOL_SOCKET, SO_INCOMING_CPU, &cpu, sizeof(int));
120	ASSERT_EQ(ret, 0);
121}
122
123int create_server(struct __test_metadata *_metadata,
124		  FIXTURE_DATA(so_incoming_cpu) *self,
125		  const FIXTURE_VARIANT(so_incoming_cpu) *variant,
126		  int cpu)
127{
128	int fd, ret;
129
130	fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0);
131	ASSERT_NE(fd, -1);
132
133	if (variant->when_to_set == BEFORE_REUSEPORT)
134		set_so_incoming_cpu(_metadata, fd, cpu);
135
136	ret = setsockopt(fd, SOL_SOCKET, SO_REUSEPORT, &(int){1}, sizeof(int));
137	ASSERT_EQ(ret, 0);
138
139	ret = bind(fd, &self->addr, self->addrlen);
140	ASSERT_EQ(ret, 0);
141
142	if (variant->when_to_set == BEFORE_LISTEN)
143		set_so_incoming_cpu(_metadata, fd, cpu);
144
145	/* We don't use nr_client_per_server here not to block
146	 * this test at connect() if SO_INCOMING_CPU is broken.
147	 */
148	ret = listen(fd, nr_client);
149	ASSERT_EQ(ret, 0);
150
151	if (variant->when_to_set == AFTER_LISTEN)
152		set_so_incoming_cpu(_metadata, fd, cpu);
153
154	return fd;
155}
156
157void create_servers(struct __test_metadata *_metadata,
158		    FIXTURE_DATA(so_incoming_cpu) *self,
159		    const FIXTURE_VARIANT(so_incoming_cpu) *variant)
160{
161	int i, ret;
162
163	for (i = 0; i < nr_server; i++) {
164		self->servers[i] = create_server(_metadata, self, variant, i);
165
166		if (i == 0) {
167			ret = getsockname(self->servers[i], &self->addr, &self->addrlen);
168			ASSERT_EQ(ret, 0);
169		}
170	}
171
172	if (variant->when_to_set == AFTER_ALL_LISTEN) {
173		for (i = 0; i < nr_server; i++)
174			set_so_incoming_cpu(_metadata, self->servers[i], i);
175	}
176}
177
178void create_clients(struct __test_metadata *_metadata,
179		    FIXTURE_DATA(so_incoming_cpu) *self)
180{
181	cpu_set_t cpu_set;
182	int i, j, fd, ret;
183
184	for (i = 0; i < nr_server; i++) {
185		CPU_ZERO(&cpu_set);
186
187		CPU_SET(i, &cpu_set);
188		ASSERT_EQ(CPU_COUNT(&cpu_set), 1);
189		ASSERT_NE(CPU_ISSET(i, &cpu_set), 0);
190
191		/* Make sure SYN will be processed on the i-th CPU
192		 * and finally distributed to the i-th listener.
193		 */
194		ret = sched_setaffinity(0, sizeof(cpu_set), &cpu_set);
195		ASSERT_EQ(ret, 0);
196
197		for (j = 0; j < nr_client_per_server; j++) {
198			fd  = socket(AF_INET, SOCK_STREAM, 0);
199			ASSERT_NE(fd, -1);
200
201			ret = connect(fd, &self->addr, self->addrlen);
202			ASSERT_EQ(ret, 0);
203
204			close(fd);
205		}
206	}
207}
208
209void verify_incoming_cpu(struct __test_metadata *_metadata,
210			 FIXTURE_DATA(so_incoming_cpu) *self)
211{
212	int i, j, fd, cpu, ret, total = 0;
213	socklen_t len = sizeof(int);
214
215	for (i = 0; i < nr_server; i++) {
216		for (j = 0; j < nr_client_per_server; j++) {
217			/* If we see -EAGAIN here, SO_INCOMING_CPU is broken */
218			fd = accept(self->servers[i], &self->addr, &self->addrlen);
219			ASSERT_NE(fd, -1);
220
221			ret = getsockopt(fd, SOL_SOCKET, SO_INCOMING_CPU, &cpu, &len);
222			ASSERT_EQ(ret, 0);
223			ASSERT_EQ(cpu, i);
224
225			close(fd);
226			total++;
227		}
228	}
229
230	ASSERT_EQ(total, nr_client);
231	TH_LOG("SO_INCOMING_CPU is very likely to be "
232	       "working correctly with %d sockets.", total);
233}
234
235TEST_F(so_incoming_cpu, test1)
236{
237	create_servers(_metadata, self, variant);
238	create_clients(_metadata, self);
239	verify_incoming_cpu(_metadata, self);
240}
241
242TEST_F(so_incoming_cpu, test2)
243{
244	int server;
245
246	create_servers(_metadata, self, variant);
247
248	/* No CPU specified */
249	server = create_server(_metadata, self, variant, -1);
250	close(server);
251
252	create_clients(_metadata, self);
253	verify_incoming_cpu(_metadata, self);
254}
255
256TEST_F(so_incoming_cpu, test3)
257{
258	int server, client;
259
260	create_servers(_metadata, self, variant);
261
262	/* No CPU specified */
263	server = create_server(_metadata, self, variant, -1);
264
265	create_clients(_metadata, self);
266
267	/* Never receive any requests */
268	client = accept(server, &self->addr, &self->addrlen);
269	ASSERT_EQ(client, -1);
270
271	verify_incoming_cpu(_metadata, self);
272}
273
274TEST_HARNESS_MAIN