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 */
47static void
48lb_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
110static int
111lb_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
127ATF_TC_WITHOUT_HEAD(basic_ipv4);
128ATF_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
182ATF_TC_WITHOUT_HEAD(basic_ipv6);
183ATF_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
238ATF_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