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 	const struct linger lopt = { 1, 0 };
58 
59 	/*
60 	 * We expect each listening socket to accept roughly nconns/nsds
61 	 * connections, but allow for some error.
62 	 */
63 	excnt = nconns / nsds / 8;
64 	acceptcnt = calloc(nsds, sizeof(*acceptcnt));
65 	ATF_REQUIRE_MSG(acceptcnt != NULL, "calloc() failed: %s",
66 	    strerror(errno));
67 
68 	while (nconns-- > 0) {
69 		sd = socket(domain, SOCK_STREAM, 0);
70 		ATF_REQUIRE_MSG(sd >= 0, "socket() failed: %s",
71 		    strerror(errno));
72 
73 		error = connect(sd, addr, addr->sa_len);
74 		ATF_REQUIRE_MSG(error == 0, "connect() failed: %s",
75 		    strerror(errno));
76 
77 		error = setsockopt(sd, SOL_SOCKET, SO_LINGER, &lopt, sizeof(lopt));
78 		ATF_REQUIRE_MSG(error == 0, "Setting linger failed: %s",
79 		    strerror(errno));
80 
81 		/*
82 		 * Poll the listening sockets.
83 		 */
84 		do {
85 			for (i = 0; i < nsds; i++) {
86 				csd = accept(sds[i], NULL, NULL);
87 				if (csd < 0) {
88 					ATF_REQUIRE_MSG(errno == EWOULDBLOCK ||
89 					    errno == EAGAIN,
90 					    "accept() failed: %s",
91 					    strerror(errno));
92 					continue;
93 				}
94 
95 				error = close(csd);
96 				ATF_REQUIRE_MSG(error == 0,
97 				    "close() failed: %s", strerror(errno));
98 
99 				acceptcnt[i]++;
100 				break;
101 			}
102 		} while (i == nsds);
103 
104 		error = close(sd);
105 		ATF_REQUIRE_MSG(error == 0, "close() failed: %s",
106 		    strerror(errno));
107 	}
108 
109 	for (i = 0; i < nsds; i++)
110 		ATF_REQUIRE_MSG(acceptcnt[i] > excnt, "uneven balancing");
111 }
112 
113 static int
114 lb_listen_socket(int domain, int flags)
115 {
116 	int one;
117 	int error, sd;
118 
119 	sd = socket(domain, SOCK_STREAM | flags, 0);
120 	ATF_REQUIRE_MSG(sd >= 0, "socket() failed: %s", strerror(errno));
121 
122 	one = 1;
123 	error = setsockopt(sd, SOL_SOCKET, SO_REUSEPORT_LB, &one, sizeof(one));
124 	ATF_REQUIRE_MSG(error == 0, "setsockopt(SO_REUSEPORT_LB) failed: %s",
125 	    strerror(errno));
126 
127 	return (sd);
128 }
129 
130 ATF_TC_WITHOUT_HEAD(basic_ipv4);
131 ATF_TC_BODY(basic_ipv4, tc)
132 {
133 	struct sockaddr_in addr;
134 	socklen_t slen;
135 	size_t i;
136 	const int nconns = 16384;
137 	int error, sds[16];
138 	uint16_t port;
139 
140 	sds[0] = lb_listen_socket(PF_INET, SOCK_NONBLOCK);
141 
142 	memset(&addr, 0, sizeof(addr));
143 	addr.sin_len = sizeof(addr);
144 	addr.sin_family = AF_INET;
145 	addr.sin_port = htons(0);
146 	addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
147 	error = bind(sds[0], (const struct sockaddr *)&addr, sizeof(addr));
148 	ATF_REQUIRE_MSG(error == 0, "bind() failed: %s", strerror(errno));
149 	error = listen(sds[0], 1);
150 	ATF_REQUIRE_MSG(error == 0, "listen() failed: %s", strerror(errno));
151 
152 	slen = sizeof(addr);
153 	error = getsockname(sds[0], (struct sockaddr *)&addr, &slen);
154 	ATF_REQUIRE_MSG(error == 0, "getsockname() failed: %s",
155 	    strerror(errno));
156 	ATF_REQUIRE_MSG(slen == sizeof(addr), "sockaddr size changed");
157 	port = addr.sin_port;
158 
159 	memset(&addr, 0, sizeof(addr));
160 	addr.sin_len = sizeof(addr);
161 	addr.sin_family = AF_INET;
162 	addr.sin_port = port;
163 	addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
164 	for (i = 1; i < nitems(sds); i++) {
165 		sds[i] = lb_listen_socket(PF_INET, SOCK_NONBLOCK);
166 
167 		error = bind(sds[i], (const struct sockaddr *)&addr,
168 		    sizeof(addr));
169 		ATF_REQUIRE_MSG(error == 0, "bind() failed: %s",
170 		    strerror(errno));
171 		error = listen(sds[i], 1);
172 		ATF_REQUIRE_MSG(error == 0, "listen() failed: %s",
173 		    strerror(errno));
174 	}
175 
176 	lb_simple_accept_loop(PF_INET, (struct sockaddr *)&addr, sds,
177 	    nitems(sds), nconns);
178 	for (i = 0; i < nitems(sds); i++) {
179 		error = close(sds[i]);
180 		ATF_REQUIRE_MSG(error == 0, "close() failed: %s",
181 		    strerror(errno));
182 	}
183 }
184 
185 ATF_TC_WITHOUT_HEAD(basic_ipv6);
186 ATF_TC_BODY(basic_ipv6, tc)
187 {
188 	const struct in6_addr loopback6 = IN6ADDR_LOOPBACK_INIT;
189 	struct sockaddr_in6 addr;
190 	socklen_t slen;
191 	size_t i;
192 	const int nconns = 16384;
193 	int error, sds[16];
194 	uint16_t port;
195 
196 	sds[0] = lb_listen_socket(PF_INET6, SOCK_NONBLOCK);
197 
198 	memset(&addr, 0, sizeof(addr));
199 	addr.sin6_len = sizeof(addr);
200 	addr.sin6_family = AF_INET6;
201 	addr.sin6_port = htons(0);
202 	addr.sin6_addr = loopback6;
203 	error = bind(sds[0], (const struct sockaddr *)&addr, sizeof(addr));
204 	ATF_REQUIRE_MSG(error == 0, "bind() failed: %s", strerror(errno));
205 	error = listen(sds[0], 1);
206 	ATF_REQUIRE_MSG(error == 0, "listen() failed: %s", strerror(errno));
207 
208 	slen = sizeof(addr);
209 	error = getsockname(sds[0], (struct sockaddr *)&addr, &slen);
210 	ATF_REQUIRE_MSG(error == 0, "getsockname() failed: %s",
211 	    strerror(errno));
212 	ATF_REQUIRE_MSG(slen == sizeof(addr), "sockaddr size changed");
213 	port = addr.sin6_port;
214 
215 	memset(&addr, 0, sizeof(addr));
216 	addr.sin6_len = sizeof(addr);
217 	addr.sin6_family = AF_INET6;
218 	addr.sin6_port = port;
219 	addr.sin6_addr = loopback6;
220 	for (i = 1; i < nitems(sds); i++) {
221 		sds[i] = lb_listen_socket(PF_INET6, SOCK_NONBLOCK);
222 
223 		error = bind(sds[i], (const struct sockaddr *)&addr,
224 		    sizeof(addr));
225 		ATF_REQUIRE_MSG(error == 0, "bind() failed: %s",
226 		    strerror(errno));
227 		error = listen(sds[i], 1);
228 		ATF_REQUIRE_MSG(error == 0, "listen() failed: %s",
229 		    strerror(errno));
230 	}
231 
232 	lb_simple_accept_loop(PF_INET6, (struct sockaddr *)&addr, sds,
233 	    nitems(sds), nconns);
234 	for (i = 0; i < nitems(sds); i++) {
235 		error = close(sds[i]);
236 		ATF_REQUIRE_MSG(error == 0, "close() failed: %s",
237 		    strerror(errno));
238 	}
239 }
240 
241 ATF_TP_ADD_TCS(tp)
242 {
243 	ATF_TP_ADD_TC(tp, basic_ipv4);
244 	ATF_TP_ADD_TC(tp, basic_ipv6);
245 
246 	return (atf_no_error());
247 }
248