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 */
50static void
51lb_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
108static int
109lb_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
125ATF_TC_WITHOUT_HEAD(basic_ipv4);
126ATF_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
180ATF_TC_WITHOUT_HEAD(basic_ipv6);
181ATF_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
236ATF_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