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