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