1 /*-
2  * Copyright (c) 2018 The FreeBSD Foundation
3  *
4  * This software was developed by Mark Johnston under sponsorship from
5  * the FreeBSD Foundation.
6  *
7  * Redistribution and use in source and binary forms, with or without
8  * modification, are permitted provided that the following conditions are
9  * met:
10  * 1. Redistributions of source code must retain the above copyright
11  *    notice, this list of conditions and the following disclaimer.
12  * 2. Redistributions in binary form must reproduce the above copyright
13  *    notice, this list of conditions and the following disclaimer in
14  *    the documentation and/or other materials provided with the
15  *    distribution.
16  *
17  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
18  * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
19  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
20  * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
21  * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
22  * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
23  * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
24  * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
25  * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
26  * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
27  * SUCH DAMAGE.
28  */
29 
30 #include <sys/cdefs.h>
31 __FBSDID("$FreeBSD$");
32 
33 #include <sys/param.h>
34 #include <sys/socket.h>
35 
36 #include <netinet/in.h>
37 
38 #include <err.h>
39 #include <errno.h>
40 #include <stdlib.h>
41 #include <unistd.h>
42 
43 #include <atf-c.h>
44 
45 /*
46  * Given an array of non-blocking listening sockets configured in a LB group
47  * for "addr", try connecting to "addr" in a loop and verify that connections
48  * are roughly balanced across the sockets.
49  */
50 static void
51 lb_simple_accept_loop(int domain, const struct sockaddr *addr, int sds[],
52     size_t nsds, int nconns)
53 {
54 	size_t i;
55 	int *acceptcnt;
56 	int csd, error, excnt, sd;
57 
58 	/*
59 	 * We expect each listening socket to accept roughly nconns/nsds
60 	 * connections, but allow for some error.
61 	 */
62 	excnt = nconns / nsds / 8;
63 	acceptcnt = calloc(nsds, sizeof(*acceptcnt));
64 	ATF_REQUIRE_MSG(acceptcnt != NULL, "calloc() failed: %s",
65 	    strerror(errno));
66 
67 	while (nconns-- > 0) {
68 		sd = socket(domain, SOCK_STREAM, 0);
69 		ATF_REQUIRE_MSG(sd >= 0, "socket() failed: %s",
70 		    strerror(errno));
71 
72 		error = connect(sd, addr, addr->sa_len);
73 		ATF_REQUIRE_MSG(error == 0, "connect() failed: %s",
74 		    strerror(errno));
75 
76 		/*
77 		 * Poll the listening sockets.
78 		 */
79 		do {
80 			for (i = 0; i < nsds; i++) {
81 				csd = accept(sds[i], NULL, NULL);
82 				if (csd < 0) {
83 					ATF_REQUIRE_MSG(errno == EWOULDBLOCK ||
84 					    errno == EAGAIN,
85 					    "accept() failed: %s",
86 					    strerror(errno));
87 					continue;
88 				}
89 
90 				error = close(csd);
91 				ATF_REQUIRE_MSG(error == 0,
92 				    "close() failed: %s", strerror(errno));
93 
94 				acceptcnt[i]++;
95 				break;
96 			}
97 		} while (i == nsds);
98 
99 		error = close(sd);
100 		ATF_REQUIRE_MSG(error == 0, "close() failed: %s",
101 		    strerror(errno));
102 	}
103 
104 	for (i = 0; i < nsds; i++)
105 		ATF_REQUIRE_MSG(acceptcnt[i] > excnt, "uneven balancing");
106 }
107 
108 static int
109 lb_listen_socket(int domain, int flags)
110 {
111 	int one;
112 	int error, sd;
113 
114 	sd = socket(domain, SOCK_STREAM | flags, 0);
115 	ATF_REQUIRE_MSG(sd >= 0, "socket() failed: %s", strerror(errno));
116 
117 	one = 1;
118 	error = setsockopt(sd, SOL_SOCKET, SO_REUSEPORT_LB, &one, sizeof(one));
119 	ATF_REQUIRE_MSG(error == 0, "setsockopt(SO_REUSEPORT_LB) failed: %s",
120 	    strerror(errno));
121 
122 	return (sd);
123 }
124 
125 ATF_TC_WITHOUT_HEAD(basic_ipv4);
126 ATF_TC_BODY(basic_ipv4, tc)
127 {
128 	struct sockaddr_in addr;
129 	socklen_t slen;
130 	size_t i;
131 	const int nconns = 16384;
132 	int error, sds[16];
133 	uint16_t port;
134 
135 	sds[0] = lb_listen_socket(PF_INET, SOCK_NONBLOCK);
136 
137 	memset(&addr, 0, sizeof(addr));
138 	addr.sin_len = sizeof(addr);
139 	addr.sin_family = AF_INET;
140 	addr.sin_port = htons(0);
141 	addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
142 	error = bind(sds[0], (const struct sockaddr *)&addr, sizeof(addr));
143 	ATF_REQUIRE_MSG(error == 0, "bind() failed: %s", strerror(errno));
144 	error = listen(sds[0], 1);
145 	ATF_REQUIRE_MSG(error == 0, "listen() failed: %s", strerror(errno));
146 
147 	slen = sizeof(addr);
148 	error = getsockname(sds[0], (struct sockaddr *)&addr, &slen);
149 	ATF_REQUIRE_MSG(error == 0, "getsockname() failed: %s",
150 	    strerror(errno));
151 	ATF_REQUIRE_MSG(slen == sizeof(addr), "sockaddr size changed");
152 	port = addr.sin_port;
153 
154 	memset(&addr, 0, sizeof(addr));
155 	addr.sin_len = sizeof(addr);
156 	addr.sin_family = AF_INET;
157 	addr.sin_port = port;
158 	addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
159 	for (i = 1; i < nitems(sds); i++) {
160 		sds[i] = lb_listen_socket(PF_INET, SOCK_NONBLOCK);
161 
162 		error = bind(sds[i], (const struct sockaddr *)&addr,
163 		    sizeof(addr));
164 		ATF_REQUIRE_MSG(error == 0, "bind() failed: %s",
165 		    strerror(errno));
166 		error = listen(sds[i], 1);
167 		ATF_REQUIRE_MSG(error == 0, "listen() failed: %s",
168 		    strerror(errno));
169 	}
170 
171 	lb_simple_accept_loop(PF_INET, (struct sockaddr *)&addr, sds,
172 	    nitems(sds), nconns);
173 	for (i = 0; i < nitems(sds); i++) {
174 		error = close(sds[i]);
175 		ATF_REQUIRE_MSG(error == 0, "close() failed: %s",
176 		    strerror(errno));
177 	}
178 }
179 
180 ATF_TC_WITHOUT_HEAD(basic_ipv6);
181 ATF_TC_BODY(basic_ipv6, tc)
182 {
183 	const struct in6_addr loopback6 = IN6ADDR_LOOPBACK_INIT;
184 	struct sockaddr_in6 addr;
185 	socklen_t slen;
186 	size_t i;
187 	const int nconns = 16384;
188 	int error, sds[16];
189 	uint16_t port;
190 
191 	sds[0] = lb_listen_socket(PF_INET6, SOCK_NONBLOCK);
192 
193 	memset(&addr, 0, sizeof(addr));
194 	addr.sin6_len = sizeof(addr);
195 	addr.sin6_family = AF_INET6;
196 	addr.sin6_port = htons(0);
197 	addr.sin6_addr = loopback6;
198 	error = bind(sds[0], (const struct sockaddr *)&addr, sizeof(addr));
199 	ATF_REQUIRE_MSG(error == 0, "bind() failed: %s", strerror(errno));
200 	error = listen(sds[0], 1);
201 	ATF_REQUIRE_MSG(error == 0, "listen() failed: %s", strerror(errno));
202 
203 	slen = sizeof(addr);
204 	error = getsockname(sds[0], (struct sockaddr *)&addr, &slen);
205 	ATF_REQUIRE_MSG(error == 0, "getsockname() failed: %s",
206 	    strerror(errno));
207 	ATF_REQUIRE_MSG(slen == sizeof(addr), "sockaddr size changed");
208 	port = addr.sin6_port;
209 
210 	memset(&addr, 0, sizeof(addr));
211 	addr.sin6_len = sizeof(addr);
212 	addr.sin6_family = AF_INET6;
213 	addr.sin6_port = port;
214 	addr.sin6_addr = loopback6;
215 	for (i = 1; i < nitems(sds); i++) {
216 		sds[i] = lb_listen_socket(PF_INET6, SOCK_NONBLOCK);
217 
218 		error = bind(sds[i], (const struct sockaddr *)&addr,
219 		    sizeof(addr));
220 		ATF_REQUIRE_MSG(error == 0, "bind() failed: %s",
221 		    strerror(errno));
222 		error = listen(sds[i], 1);
223 		ATF_REQUIRE_MSG(error == 0, "listen() failed: %s",
224 		    strerror(errno));
225 	}
226 
227 	lb_simple_accept_loop(PF_INET6, (struct sockaddr *)&addr, sds,
228 	    nitems(sds), nconns);
229 	for (i = 0; i < nitems(sds); i++) {
230 		error = close(sds[i]);
231 		ATF_REQUIRE_MSG(error == 0, "close() failed: %s",
232 		    strerror(errno));
233 	}
234 }
235 
236 ATF_TP_ADD_TCS(tp)
237 {
238 	ATF_TP_ADD_TC(tp, basic_ipv4);
239 	ATF_TP_ADD_TC(tp, basic_ipv6);
240 
241 	return (atf_no_error());
242 }
243