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 
14 FIXTURE(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 
24 enum when_to_set {
25 	BEFORE_REUSEPORT,
26 	BEFORE_LISTEN,
27 	AFTER_LISTEN,
28 	AFTER_ALL_LISTEN,
29 };
30 
31 FIXTURE_VARIANT(so_incoming_cpu)
32 {
33 	int when_to_set;
34 };
35 
36 FIXTURE_VARIANT_ADD(so_incoming_cpu, before_reuseport)
37 {
38 	.when_to_set = BEFORE_REUSEPORT,
39 };
40 
41 FIXTURE_VARIANT_ADD(so_incoming_cpu, before_listen)
42 {
43 	.when_to_set = BEFORE_LISTEN,
44 };
45 
46 FIXTURE_VARIANT_ADD(so_incoming_cpu, after_listen)
47 {
48 	.when_to_set = AFTER_LISTEN,
49 };
50 
51 FIXTURE_VARIANT_ADD(so_incoming_cpu, after_all_listen)
52 {
53 	.when_to_set = AFTER_ALL_LISTEN,
54 };
55 
56 static 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 
69 static 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
80 static int nr_client_per_server, nr_server, nr_client;
81 
82 FIXTURE_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 
105 FIXTURE_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 
115 void 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 
123 int 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 
157 void 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 
178 void 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 
209 void 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 
235 TEST_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 
242 TEST_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 
256 TEST_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 
274 TEST_HARNESS_MAIN
275