1 // SPDX-License-Identifier: GPL-2.0
2 /*
3  * Test functionality of BPF filters with SO_REUSEPORT.  This program creates
4  * an SO_REUSEPORT receiver group containing one socket per CPU core. It then
5  * creates a BPF program that will select a socket from this group based
6  * on the core id that receives the packet.  The sending code artificially
7  * moves itself to run on different core ids and sends one message from
8  * each core.  Since these packets are delivered over loopback, they should
9  * arrive on the same core that sent them.  The receiving code then ensures
10  * that the packet was received on the socket for the corresponding core id.
11  * This entire process is done for several different core id permutations
12  * and for each IPv4/IPv6 and TCP/UDP combination.
13  */
14 
15 #define _GNU_SOURCE
16 
17 #include <arpa/inet.h>
18 #include <errno.h>
19 #include <error.h>
20 #include <linux/filter.h>
21 #include <linux/in.h>
22 #include <linux/unistd.h>
23 #include <sched.h>
24 #include <stdio.h>
25 #include <stdlib.h>
26 #include <string.h>
27 #include <sys/epoll.h>
28 #include <sys/types.h>
29 #include <sys/socket.h>
30 #include <unistd.h>
31 
32 static const int PORT = 8888;
33 
build_rcv_group(int * rcv_fd,size_t len,int family,int proto)34 static void build_rcv_group(int *rcv_fd, size_t len, int family, int proto)
35 {
36 	struct sockaddr_storage addr;
37 	struct sockaddr_in  *addr4;
38 	struct sockaddr_in6 *addr6;
39 	size_t i;
40 	int opt;
41 
42 	switch (family) {
43 	case AF_INET:
44 		addr4 = (struct sockaddr_in *)&addr;
45 		addr4->sin_family = AF_INET;
46 		addr4->sin_addr.s_addr = htonl(INADDR_ANY);
47 		addr4->sin_port = htons(PORT);
48 		break;
49 	case AF_INET6:
50 		addr6 = (struct sockaddr_in6 *)&addr;
51 		addr6->sin6_family = AF_INET6;
52 		addr6->sin6_addr = in6addr_any;
53 		addr6->sin6_port = htons(PORT);
54 		break;
55 	default:
56 		error(1, 0, "Unsupported family %d", family);
57 	}
58 
59 	for (i = 0; i < len; ++i) {
60 		rcv_fd[i] = socket(family, proto, 0);
61 		if (rcv_fd[i] < 0)
62 			error(1, errno, "failed to create receive socket");
63 
64 		opt = 1;
65 		if (setsockopt(rcv_fd[i], SOL_SOCKET, SO_REUSEPORT, &opt,
66 			       sizeof(opt)))
67 			error(1, errno, "failed to set SO_REUSEPORT");
68 
69 		if (bind(rcv_fd[i], (struct sockaddr *)&addr, sizeof(addr)))
70 			error(1, errno, "failed to bind receive socket");
71 
72 		if (proto == SOCK_STREAM && listen(rcv_fd[i], len * 10))
73 			error(1, errno, "failed to listen on receive port");
74 	}
75 }
76 
attach_bpf(int fd)77 static void attach_bpf(int fd)
78 {
79 	struct sock_filter code[] = {
80 		/* A = raw_smp_processor_id() */
81 		{ BPF_LD  | BPF_W | BPF_ABS, 0, 0, SKF_AD_OFF + SKF_AD_CPU },
82 		/* return A */
83 		{ BPF_RET | BPF_A, 0, 0, 0 },
84 	};
85 	struct sock_fprog p = {
86 		.len = 2,
87 		.filter = code,
88 	};
89 
90 	if (setsockopt(fd, SOL_SOCKET, SO_ATTACH_REUSEPORT_CBPF, &p, sizeof(p)))
91 		error(1, errno, "failed to set SO_ATTACH_REUSEPORT_CBPF");
92 }
93 
send_from_cpu(int cpu_id,int family,int proto)94 static void send_from_cpu(int cpu_id, int family, int proto)
95 {
96 	struct sockaddr_storage saddr, daddr;
97 	struct sockaddr_in  *saddr4, *daddr4;
98 	struct sockaddr_in6 *saddr6, *daddr6;
99 	cpu_set_t cpu_set;
100 	int fd;
101 
102 	switch (family) {
103 	case AF_INET:
104 		saddr4 = (struct sockaddr_in *)&saddr;
105 		saddr4->sin_family = AF_INET;
106 		saddr4->sin_addr.s_addr = htonl(INADDR_ANY);
107 		saddr4->sin_port = 0;
108 
109 		daddr4 = (struct sockaddr_in *)&daddr;
110 		daddr4->sin_family = AF_INET;
111 		daddr4->sin_addr.s_addr = htonl(INADDR_LOOPBACK);
112 		daddr4->sin_port = htons(PORT);
113 		break;
114 	case AF_INET6:
115 		saddr6 = (struct sockaddr_in6 *)&saddr;
116 		saddr6->sin6_family = AF_INET6;
117 		saddr6->sin6_addr = in6addr_any;
118 		saddr6->sin6_port = 0;
119 
120 		daddr6 = (struct sockaddr_in6 *)&daddr;
121 		daddr6->sin6_family = AF_INET6;
122 		daddr6->sin6_addr = in6addr_loopback;
123 		daddr6->sin6_port = htons(PORT);
124 		break;
125 	default:
126 		error(1, 0, "Unsupported family %d", family);
127 	}
128 
129 	memset(&cpu_set, 0, sizeof(cpu_set));
130 	CPU_SET(cpu_id, &cpu_set);
131 	if (sched_setaffinity(0, sizeof(cpu_set), &cpu_set) < 0)
132 		error(1, errno, "failed to pin to cpu");
133 
134 	fd = socket(family, proto, 0);
135 	if (fd < 0)
136 		error(1, errno, "failed to create send socket");
137 
138 	if (bind(fd, (struct sockaddr *)&saddr, sizeof(saddr)))
139 		error(1, errno, "failed to bind send socket");
140 
141 	if (connect(fd, (struct sockaddr *)&daddr, sizeof(daddr)))
142 		error(1, errno, "failed to connect send socket");
143 
144 	if (send(fd, "a", 1, 0) < 0)
145 		error(1, errno, "failed to send message");
146 
147 	close(fd);
148 }
149 
150 static
receive_on_cpu(int * rcv_fd,int len,int epfd,int cpu_id,int proto)151 void receive_on_cpu(int *rcv_fd, int len, int epfd, int cpu_id, int proto)
152 {
153 	struct epoll_event ev;
154 	int i, fd;
155 	char buf[8];
156 
157 	i = epoll_wait(epfd, &ev, 1, -1);
158 	if (i < 0)
159 		error(1, errno, "epoll_wait failed");
160 
161 	if (proto == SOCK_STREAM) {
162 		fd = accept(ev.data.fd, NULL, NULL);
163 		if (fd < 0)
164 			error(1, errno, "failed to accept");
165 		i = recv(fd, buf, sizeof(buf), 0);
166 		close(fd);
167 	} else {
168 		i = recv(ev.data.fd, buf, sizeof(buf), 0);
169 	}
170 
171 	if (i < 0)
172 		error(1, errno, "failed to recv");
173 
174 	for (i = 0; i < len; ++i)
175 		if (ev.data.fd == rcv_fd[i])
176 			break;
177 	if (i == len)
178 		error(1, 0, "failed to find socket");
179 	fprintf(stderr, "send cpu %d, receive socket %d\n", cpu_id, i);
180 	if (cpu_id != i)
181 		error(1, 0, "cpu id/receive socket mismatch");
182 }
183 
test(int * rcv_fd,int len,int family,int proto)184 static void test(int *rcv_fd, int len, int family, int proto)
185 {
186 	struct epoll_event ev;
187 	int epfd, cpu;
188 
189 	build_rcv_group(rcv_fd, len, family, proto);
190 	attach_bpf(rcv_fd[0]);
191 
192 	epfd = epoll_create(1);
193 	if (epfd < 0)
194 		error(1, errno, "failed to create epoll");
195 	for (cpu = 0; cpu < len; ++cpu) {
196 		ev.events = EPOLLIN;
197 		ev.data.fd = rcv_fd[cpu];
198 		if (epoll_ctl(epfd, EPOLL_CTL_ADD, rcv_fd[cpu], &ev))
199 			error(1, errno, "failed to register sock epoll");
200 	}
201 
202 	/* Forward iterate */
203 	for (cpu = 0; cpu < len; ++cpu) {
204 		send_from_cpu(cpu, family, proto);
205 		receive_on_cpu(rcv_fd, len, epfd, cpu, proto);
206 	}
207 
208 	/* Reverse iterate */
209 	for (cpu = len - 1; cpu >= 0; --cpu) {
210 		send_from_cpu(cpu, family, proto);
211 		receive_on_cpu(rcv_fd, len, epfd, cpu, proto);
212 	}
213 
214 	/* Even cores */
215 	for (cpu = 0; cpu < len; cpu += 2) {
216 		send_from_cpu(cpu, family, proto);
217 		receive_on_cpu(rcv_fd, len, epfd, cpu, proto);
218 	}
219 
220 	/* Odd cores */
221 	for (cpu = 1; cpu < len; cpu += 2) {
222 		send_from_cpu(cpu, family, proto);
223 		receive_on_cpu(rcv_fd, len, epfd, cpu, proto);
224 	}
225 
226 	close(epfd);
227 	for (cpu = 0; cpu < len; ++cpu)
228 		close(rcv_fd[cpu]);
229 }
230 
main(void)231 int main(void)
232 {
233 	int *rcv_fd, cpus;
234 
235 	cpus = sysconf(_SC_NPROCESSORS_ONLN);
236 	if (cpus <= 0)
237 		error(1, errno, "failed counting cpus");
238 
239 	rcv_fd = calloc(cpus, sizeof(int));
240 	if (!rcv_fd)
241 		error(1, 0, "failed to allocate array");
242 
243 	fprintf(stderr, "---- IPv4 UDP ----\n");
244 	test(rcv_fd, cpus, AF_INET, SOCK_DGRAM);
245 
246 	fprintf(stderr, "---- IPv6 UDP ----\n");
247 	test(rcv_fd, cpus, AF_INET6, SOCK_DGRAM);
248 
249 	fprintf(stderr, "---- IPv4 TCP ----\n");
250 	test(rcv_fd, cpus, AF_INET, SOCK_STREAM);
251 
252 	fprintf(stderr, "---- IPv6 TCP ----\n");
253 	test(rcv_fd, cpus, AF_INET6, SOCK_STREAM);
254 
255 	free(rcv_fd);
256 
257 	fprintf(stderr, "SUCCESS\n");
258 	return 0;
259 }
260