1/*	$NetBSD: test_client.c,v 1.2 2024/02/21 22:51:13 christos Exp $	*/
2
3/*
4 * Copyright (C) Internet Systems Consortium, Inc. ("ISC")
5 *
6 * SPDX-License-Identifier: MPL-2.0
7 *
8 * This Source Code Form is subject to the terms of the Mozilla Public
9 * License, v. 2.0. If a copy of the MPL was not distributed with this
10 * file, you can obtain one at https://mozilla.org/MPL/2.0/.
11 *
12 * See the COPYRIGHT file distributed with this work for additional
13 * information regarding copyright ownership.
14 */
15
16#include <fcntl.h>
17#include <getopt.h>
18#include <netdb.h>
19#include <netinet/in.h>
20#include <signal.h>
21#include <stdbool.h>
22#include <stdio.h>
23#include <stdlib.h>
24#include <strings.h>
25#include <sys/socket.h>
26#include <sys/stat.h>
27#include <sys/types.h>
28#include <unistd.h>
29
30#include <isc/managers.h>
31#include <isc/mem.h>
32#include <isc/netaddr.h>
33#include <isc/netmgr.h>
34#include <isc/os.h>
35#include <isc/print.h>
36#include <isc/sockaddr.h>
37#include <isc/string.h>
38#include <isc/util.h>
39
40typedef enum {
41	UDP,
42	TCP,
43	DOT,
44	HTTPS_POST,
45	HTTPS_GET,
46	HTTP_POST,
47	HTTP_GET
48} protocol_t;
49
50static const char *protocols[] = { "udp",	    "tcp",
51				   "dot",	    "https-post",
52				   "https-get",	    "http-plain-post",
53				   "http-plain-get" };
54
55static isc_mem_t *mctx = NULL;
56static isc_nm_t *netmgr = NULL;
57
58static protocol_t protocol;
59static const char *address;
60static const char *port;
61static int family = AF_UNSPEC;
62static isc_sockaddr_t sockaddr_local;
63static isc_sockaddr_t sockaddr_remote;
64static int workers;
65static int timeout;
66static uint8_t messagebuf[2 * 65536];
67static isc_region_t message = { .length = 0, .base = messagebuf };
68static int out = -1;
69
70static isc_tlsctx_t *tls_ctx = NULL;
71
72static isc_result_t
73parse_port(const char *input) {
74	char *endptr = NULL;
75	long val = strtol(input, &endptr, 10);
76
77	if ((*endptr != '\0') || (val <= 0) || (val >= 65536)) {
78		return (ISC_R_BADNUMBER);
79	}
80
81	port = input;
82
83	return (ISC_R_SUCCESS);
84}
85
86static isc_result_t
87parse_protocol(const char *input) {
88	for (size_t i = 0; i < ARRAY_SIZE(protocols); i++) {
89		if (!strcasecmp(input, protocols[i])) {
90			protocol = i;
91			return (ISC_R_SUCCESS);
92		}
93	}
94
95	return (ISC_R_BADNUMBER);
96}
97
98static isc_result_t
99parse_address(const char *input) {
100	struct in6_addr in6;
101	struct in_addr in;
102
103	if (inet_pton(AF_INET6, input, &in6) == 1) {
104		family = AF_INET6;
105		address = input;
106		return (ISC_R_SUCCESS);
107	}
108
109	if (inet_pton(AF_INET, input, &in) == 1) {
110		family = AF_INET;
111		address = input;
112		return (ISC_R_SUCCESS);
113	}
114
115	return (ISC_R_BADADDRESSFORM);
116}
117
118static int
119parse_workers(const char *input) {
120	char *endptr = NULL;
121	long val = strtol(input, &endptr, 10);
122
123	if ((*endptr != '\0') || (val <= 0) || (val >= 128)) {
124		return (ISC_R_BADNUMBER);
125	}
126
127	workers = val;
128
129	return (ISC_R_SUCCESS);
130}
131
132static isc_result_t
133parse_timeout(const char *input) {
134	char *endptr = NULL;
135	long val = strtol(input, &endptr, 10);
136
137	if ((*endptr != '\0') || (val <= 0) || (val >= 120)) {
138		return (ISC_R_BADNUMBER);
139	}
140
141	timeout = (in_port_t)val * 1000;
142
143	return (ISC_R_SUCCESS);
144}
145
146static isc_result_t
147parse_input(const char *input) {
148	int in = -1;
149
150	if (!strcmp(input, "-")) {
151		in = 0;
152	} else {
153		in = open(input, O_RDONLY);
154	}
155	RUNTIME_CHECK(in >= 0);
156
157	message.length = read(in, message.base, sizeof(messagebuf));
158
159	close(in);
160
161	return (ISC_R_SUCCESS);
162}
163
164static isc_result_t
165parse_output(const char *input) {
166	if (!strcmp(input, "-")) {
167		out = 1;
168	} else {
169		out = open(input, O_WRONLY | O_CREAT,
170			   S_IRUSR | S_IRGRP | S_IROTH);
171	}
172	RUNTIME_CHECK(out >= 0);
173
174	return (ISC_R_SUCCESS);
175}
176
177static void
178parse_options(int argc, char **argv) {
179	char buf[ISC_NETADDR_FORMATSIZE];
180
181	/* Set defaults */
182	RUNTIME_CHECK(parse_protocol("UDP") == ISC_R_SUCCESS);
183	RUNTIME_CHECK(parse_port("53000") == ISC_R_SUCCESS);
184	RUNTIME_CHECK(parse_address("::0") == ISC_R_SUCCESS);
185	workers = isc_os_ncpus();
186
187	while (true) {
188		int c;
189		int option_index = 0;
190		static struct option long_options[] = {
191			{ "port", required_argument, NULL, 'p' },
192			{ "address", required_argument, NULL, 'a' },
193			{ "protocol", required_argument, NULL, 'P' },
194			{ "workers", required_argument, NULL, 'w' },
195			{ "timeout", required_argument, NULL, 't' },
196			{ "input", required_argument, NULL, 'i' },
197			{ "output", required_argument, NULL, 'o' },
198			{ 0, 0, NULL, 0 }
199		};
200
201		c = getopt_long(argc, argv, "a:p:P:w:t:i:o:", long_options,
202				&option_index);
203		if (c == -1) {
204			break;
205		}
206
207		switch (c) {
208		case 'a':
209			RUNTIME_CHECK(parse_address(optarg) == ISC_R_SUCCESS);
210			break;
211
212		case 'p':
213			RUNTIME_CHECK(parse_port(optarg) == ISC_R_SUCCESS);
214			break;
215
216		case 'P':
217			RUNTIME_CHECK(parse_protocol(optarg) == ISC_R_SUCCESS);
218			break;
219
220		case 'w':
221			RUNTIME_CHECK(parse_workers(optarg) == ISC_R_SUCCESS);
222			break;
223
224		case 't':
225			RUNTIME_CHECK(parse_timeout(optarg) == ISC_R_SUCCESS);
226			break;
227
228		case 'i':
229			RUNTIME_CHECK(parse_input(optarg) == ISC_R_SUCCESS);
230			break;
231
232		case 'o':
233			RUNTIME_CHECK(parse_output(optarg) == ISC_R_SUCCESS);
234			break;
235
236		default:
237			UNREACHABLE();
238		}
239	}
240
241	{
242		struct addrinfo hints = {
243			.ai_family = family,
244			.ai_socktype = (protocol == UDP) ? SOCK_DGRAM
245							 : SOCK_STREAM,
246		};
247		struct addrinfo *result = NULL;
248		int r = getaddrinfo(address, NULL, &hints, &result);
249		RUNTIME_CHECK(r == 0);
250
251		for (struct addrinfo *rp = result; rp != NULL; rp = rp->ai_next)
252		{
253			RUNTIME_CHECK(isc_sockaddr_fromsockaddr(&sockaddr_local,
254								rp->ai_addr) ==
255				      ISC_R_SUCCESS);
256		}
257		freeaddrinfo(result);
258	}
259
260	{
261		struct addrinfo hints = {
262			.ai_family = family,
263			.ai_socktype = (protocol == UDP) ? SOCK_DGRAM
264							 : SOCK_STREAM,
265		};
266		struct addrinfo *result = NULL;
267		int r = getaddrinfo(argv[optind], port, &hints, &result);
268		RUNTIME_CHECK(r == 0);
269
270		for (struct addrinfo *rp = result; rp != NULL; rp = rp->ai_next)
271		{
272			RUNTIME_CHECK(isc_sockaddr_fromsockaddr(
273					      &sockaddr_remote, rp->ai_addr) ==
274				      ISC_R_SUCCESS);
275		}
276		freeaddrinfo(result);
277	}
278
279	isc_sockaddr_format(&sockaddr_local, buf, sizeof(buf));
280
281	printf("Will connect from %s://%s", protocols[protocol], buf);
282
283	isc_sockaddr_format(&sockaddr_remote, buf, sizeof(buf));
284
285	printf(" to %s, %d workers\n", buf, workers);
286}
287
288static void
289_signal(int sig, void (*handler)(int)) {
290	struct sigaction sa = { .sa_handler = handler };
291
292	RUNTIME_CHECK(sigfillset(&sa.sa_mask) == 0);
293	RUNTIME_CHECK(sigaction(sig, &sa, NULL) >= 0);
294}
295
296static void
297setup(void) {
298	sigset_t sset;
299
300	_signal(SIGPIPE, SIG_IGN);
301	_signal(SIGHUP, SIG_DFL);
302	_signal(SIGTERM, SIG_DFL);
303	_signal(SIGINT, SIG_DFL);
304
305	RUNTIME_CHECK(sigemptyset(&sset) == 0);
306	RUNTIME_CHECK(sigaddset(&sset, SIGHUP) == 0);
307	RUNTIME_CHECK(sigaddset(&sset, SIGINT) == 0);
308	RUNTIME_CHECK(sigaddset(&sset, SIGTERM) == 0);
309	RUNTIME_CHECK(pthread_sigmask(SIG_BLOCK, &sset, NULL) == 0);
310
311	isc_mem_create(&mctx);
312
313	isc_managers_create(mctx, workers, 0, &netmgr, NULL, NULL);
314}
315
316static void
317teardown(void) {
318	if (out > 0) {
319		close(out);
320	}
321
322	isc_managers_destroy(&netmgr, NULL, NULL);
323	isc_mem_destroy(&mctx);
324	if (tls_ctx) {
325		isc_tlsctx_free(&tls_ctx);
326	}
327}
328
329static void
330waitforsignal(void) {
331	sigset_t sset;
332	int sig;
333
334	RUNTIME_CHECK(sigemptyset(&sset) == 0);
335	RUNTIME_CHECK(sigaddset(&sset, SIGHUP) == 0);
336	RUNTIME_CHECK(sigaddset(&sset, SIGINT) == 0);
337	RUNTIME_CHECK(sigaddset(&sset, SIGTERM) == 0);
338	RUNTIME_CHECK(sigwait(&sset, &sig) == 0);
339
340	fprintf(stderr, "Shutting down...\n");
341}
342
343static void
344read_cb(isc_nmhandle_t *handle, isc_result_t eresult, isc_region_t *region,
345	void *cbarg) {
346	isc_nmhandle_t *readhandle = cbarg;
347
348	REQUIRE(handle != NULL);
349	REQUIRE(eresult == ISC_R_SUCCESS || eresult == ISC_R_CANCELED ||
350		eresult == ISC_R_EOF);
351	REQUIRE(cbarg != NULL);
352
353	fprintf(stderr, "%s(..., %s, ...)\n", __func__,
354		isc_result_totext(eresult));
355
356	if (eresult == ISC_R_SUCCESS) {
357		printf("RECEIVED %u bytes\n", region->length);
358		if (out >= 0) {
359			ssize_t len = write(out, region->base, region->length);
360			close(out);
361			REQUIRE((size_t)len == region->length);
362		}
363	}
364
365	isc_nmhandle_detach(&readhandle);
366	kill(getpid(), SIGTERM);
367}
368
369static void
370send_cb(isc_nmhandle_t *handle, isc_result_t eresult, void *cbarg) {
371	REQUIRE(handle != NULL);
372	REQUIRE(eresult == ISC_R_SUCCESS || eresult == ISC_R_CANCELED ||
373		eresult == ISC_R_EOF);
374	REQUIRE(cbarg == NULL);
375}
376
377static void
378connect_cb(isc_nmhandle_t *handle, isc_result_t eresult, void *cbarg) {
379	isc_nmhandle_t *readhandle = NULL;
380
381	REQUIRE(handle != NULL);
382	UNUSED(cbarg);
383
384	fprintf(stderr, "ECHO_CLIENT:%s:%s\n", __func__,
385		isc_result_totext(eresult));
386
387	if (eresult != ISC_R_SUCCESS) {
388		kill(getpid(), SIGTERM);
389		return;
390	}
391
392	isc_nmhandle_attach(handle, &readhandle);
393	isc_nm_read(handle, read_cb, readhandle);
394	isc_nm_send(handle, &message, send_cb, NULL);
395}
396
397static void
398run(void) {
399	switch (protocol) {
400	case UDP:
401		isc_nm_udpconnect(netmgr, &sockaddr_local, &sockaddr_remote,
402				  connect_cb, NULL, timeout, 0);
403		break;
404	case TCP:
405		isc_nm_tcpdnsconnect(netmgr, &sockaddr_local, &sockaddr_remote,
406				     connect_cb, NULL, timeout, 0);
407		break;
408	case DOT: {
409		isc_tlsctx_createclient(&tls_ctx);
410
411		isc_nm_tlsdnsconnect(netmgr, &sockaddr_local, &sockaddr_remote,
412				     connect_cb, NULL, timeout, 0, tls_ctx,
413				     NULL);
414		break;
415	}
416#if HAVE_LIBNGHTTP2
417	case HTTP_GET:
418	case HTTPS_GET:
419	case HTTPS_POST:
420	case HTTP_POST: {
421		bool is_https = (protocol == HTTPS_POST ||
422				 protocol == HTTPS_GET);
423		bool is_post = (protocol == HTTPS_POST ||
424				protocol == HTTP_POST);
425		char req_url[256];
426		isc_nm_http_makeuri(is_https, &sockaddr_remote, NULL, 0,
427				    ISC_NM_HTTP_DEFAULT_PATH, req_url,
428				    sizeof(req_url));
429		if (is_https) {
430			isc_tlsctx_createclient(&tls_ctx);
431		}
432		isc_nm_httpconnect(netmgr, &sockaddr_local, &sockaddr_remote,
433				   req_url, is_post, connect_cb, NULL, tls_ctx,
434				   NULL, timeout, 0);
435	} break;
436#endif
437	default:
438		UNREACHABLE();
439	}
440
441	waitforsignal();
442}
443
444int
445main(int argc, char **argv) {
446	parse_options(argc, argv);
447
448	setup();
449
450	run();
451
452	teardown();
453
454	exit(EXIT_SUCCESS);
455}
456