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