1 /*
2  * TCP client and server for bug hunting
3  *
4  * Copyright (C) 2016 Willy Tarreau <w@1wt.eu>
5  *
6  * Permission is hereby granted, free of charge, to any person obtaining
7  * a copy of this software and associated documentation files (the
8  * "Software"), to deal in the Software without restriction, including
9  * without limitation the rights to use, copy, modify, merge, publish,
10  * distribute, sublicense, and/or sell copies of the Software, and to
11  * permit persons to whom the Software is furnished to do so, subject to
12  * the following conditions:
13  *
14  * The above copyright notice and this permission notice shall be
15  * included in all copies or substantial portions of the Software.
16  *
17  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
18  * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
19  * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
20  * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
21  * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
22  * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
23  * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
24  * OTHER DEALINGS IN THE SOFTWARE.
25  */
26 
27 #include <sys/resource.h>
28 #include <sys/select.h>
29 #include <sys/types.h>
30 #include <sys/socket.h>
31 #include <sys/stat.h>
32 #include <sys/time.h>
33 #include <sys/ioctl.h>
34 #include <sys/wait.h>
35 #include <arpa/inet.h>
36 #include <netinet/in.h>
37 #include <netinet/tcp.h>
38 
39 #include <ctype.h>
40 #include <errno.h>
41 #include <fcntl.h>
42 #include <limits.h>
43 #include <netdb.h>
44 #include <poll.h>
45 #include <signal.h>
46 #include <stdarg.h>
47 #include <stdio.h>
48 #include <stdlib.h>
49 #include <string.h>
50 #include <time.h>
51 #include <unistd.h>
52 
53 #ifndef SOL_TCP
54 #define SOL_TCP IPPROTO_TCP
55 #endif
56 
57 #ifndef MSG_MORE
58 #define MSG_MORE 0
59 #endif
60 
61 struct err_msg {
62 	int size;
63 	int len;
64 	char msg[0];
65 };
66 
67 const int zero = 0;
68 const int one = 1;
69 const struct linger nolinger = { .l_onoff = 1, .l_linger = 0 };
70 
71 #define TRASH_SIZE 65536
72 static char trash[TRASH_SIZE];
73 
74 volatile int nbproc = 0;
75 static struct timeval start_time;
76 static int showtime;
77 static int verbose;
78 static int pid;
79 
80 
81 /* display the message and exit with the code */
die(int code,const char * format,...)82 __attribute__((noreturn)) void die(int code, const char *format, ...)
83 {
84 	va_list args;
85 
86 	if (format) {
87 		va_start(args, format);
88 		vfprintf(stderr, format, args);
89 		va_end(args);
90 	}
91 	exit(code);
92 }
93 
94 /* display the usage message and exit with the code */
usage(int code,const char * arg0)95 __attribute__((noreturn)) void usage(int code, const char *arg0)
96 {
97 	die(code,
98 	    "Usage : %s [options]* [<ip>:]port [<action>*]\n"
99 	    "\n"
100 	    "options :\n"
101 	    "  -v           : verbose\n"
102 	    "  -t|-tt|-ttt  : show time (msec / relative / absolute)\n"
103 	    "actions :\n"
104 	    "  L[<backlog>] : Listens to ip:port and optionally sets backlog\n"
105 	    "                 Note: fd=socket,bind(fd),listen(fd)\n"
106 	    "  C            : Connects to ip:port\n"
107 	    "                 Note: fd=socket,connect(fd)\n"
108 	    "  A[<count>]   : Accepts <count> incoming sockets and closes count-1\n"
109 	    "                 Note: fd=accept(fd)\n"
110 	    "  J            : Jump back to oldest post-fork/post-accept action\n"
111 	    "  K            : kill the connection and go on with next operation\n"
112 	    "  G            : disable lingering\n"
113 	    "  T            : set TCP_NODELAY\n"
114 	    "  Q            : disable TCP Quick-ack\n"
115 	    "  R[<size>]    : Read this amount of bytes. 0=infinite. unset=any amount.\n"
116 	    "  S[<size>]    : Send this amount of bytes. 0=infinite. unset=any amount.\n"
117 	    "  S:<string>   : Send this exact string. \\r, \\n, \\t, \\\\ supported.\n"
118 	    "  E[<size>]    : Echo this amount of bytes. 0=infinite. unset=any amount.\n"
119 	    "  W[<time>]    : Wait for any event on the socket, maximum <time> ms\n"
120 	    "  P[<time>]    : Pause for <time> ms (100 by default)\n"
121 	    "  I            : wait for Input data to be present (POLLIN)\n"
122 	    "  O            : wait for Output queue to be empty (POLLOUT + TIOCOUTQ)\n"
123 	    "  F            : FIN : shutdown(SHUT_WR)\n"
124 	    "  N<max>       : fork New process, limited to <max> concurrent (default 1)\n"
125 	    "\n"
126 	    "It's important to note that a single FD is used at once and that Accept\n"
127 	    "replaces the listening FD with the accepted one. Thus always do it after\n"
128 	    "a fork if other connections have to be accepted.\n"
129 	    "\n"
130 	    "After a fork, we loop back to the beginning and silently skip L/C if the\n"
131 	    "main socket already exists.\n"
132 	    "\n"
133 	    "Example dummy HTTP request drain server :\n"
134 	    "   tcploop 8001 L W N20 A R S10 [ F K ]\n"
135 	    "\n"
136 	    "Example large bandwidth HTTP request drain server :\n"
137 	    "   tcploop 8001 L W N20 A R S0 [ F K ]\n"
138 	    "\n"
139 	    "Example TCP client with pauses at each step :\n"
140 	    "   tcploop 8001 C T W P100 S10 O P100 R S10 O R G K\n"
141 	    "", arg0);
142 }
143 
dolog(const char * format,...)144 void dolog(const char *format, ...)
145 {
146 	struct timeval date, tv;
147 	int delay;
148 	va_list args;
149 
150 	if (!verbose)
151 		return;
152 
153 	if (showtime) {
154 		gettimeofday(&date, NULL);
155 		switch (showtime) {
156 		case 1: // [msec] relative
157 			delay = (date.tv_sec - start_time.tv_sec) * 1000000 + date.tv_usec - start_time.tv_usec;
158 			fprintf(stderr, "[%d] ", delay / 1000);
159 			break;
160 		case 2: // [sec.usec] relative
161 			tv.tv_usec = date.tv_usec - start_time.tv_usec;
162 			tv.tv_sec  = date.tv_sec  - start_time.tv_sec;
163 			if ((signed)tv.tv_sec > 0) {
164 				if ((signed)tv.tv_usec < 0) {
165 					tv.tv_usec += 1000000;
166 					tv.tv_sec--;
167 				}
168 			} else if (tv.tv_sec == 0) {
169 				if ((signed)tv.tv_usec < 0)
170 					tv.tv_usec = 0;
171 			} else {
172 				tv.tv_sec = 0;
173 				tv.tv_usec = 0;
174 			}
175 			fprintf(stderr, "[%d.%06d] ", (int)tv.tv_sec, (int)tv.tv_usec);
176 			break;
177 		default: // [sec.usec] absolute
178 			fprintf(stderr, "[%d.%06d] ", (int)date.tv_sec, (int)date.tv_usec);
179 			break;
180 		}
181 	}
182 
183 	fprintf(stderr, "%5d ", pid);
184 
185 	va_start(args, format);
186 	vfprintf(stderr, format, args);
187 	va_end(args);
188 }
189 
190 /* convert '\n', '\t', '\r', '\\' to their respective characters */
unescape(char * out,int size,const char * in)191 int unescape(char *out, int size, const char *in)
192 {
193 	int len;
194 
195 	for (len = 0; len < size && *in; in++, out++, len++) {
196 		if (*in == '\\') {
197 			switch (in[1]) {
198 			case  'n' : *out = '\n'; in++; continue;
199 			case  't' : *out = '\t'; in++; continue;
200 			case  'r' : *out = '\r'; in++; continue;
201 			case '\\' : *out = '\\'; in++; continue;
202 			default   : break;
203 			}
204 		}
205 		*out = *in;
206 	}
207 	return len;
208 }
209 
alloc_err_msg(int size)210 struct err_msg *alloc_err_msg(int size)
211 {
212 	struct err_msg *err;
213 
214 	err = malloc(sizeof(*err) + size);
215 	if (err) {
216 		err->len = 0;
217 		err->size = size;
218 	}
219 	return err;
220 }
221 
sig_handler(int sig)222 void sig_handler(int sig)
223 {
224 	if (sig == SIGCHLD) {
225 		while (waitpid(-1, NULL, WNOHANG) > 0)
226 			__sync_sub_and_fetch(&nbproc, 1);
227 	}
228 }
229 
230 /* converts str in the form [[<ipv4>|<ipv6>|<hostname>]:]port to struct sockaddr_storage.
231  * Returns < 0 with err set in case of error.
232  */
addr_to_ss(char * str,struct sockaddr_storage * ss,struct err_msg * err)233 int addr_to_ss(char *str, struct sockaddr_storage *ss, struct err_msg *err)
234 {
235 	char *port_str;
236 	int port;
237 
238 	memset(ss, 0, sizeof(*ss));
239 
240 	/* look for the addr/port delimiter, it's the last colon. If there's no
241 	 * colon, it's 0:<port>.
242 	 */
243 	if ((port_str = strrchr(str, ':')) == NULL) {
244 		port = atoi(str);
245 		if (port <= 0 || port > 65535) {
246 			err->len = snprintf(err->msg, err->size, "Missing/invalid port number: '%s'\n", str);
247 			return -1;
248 		}
249 
250 		ss->ss_family = AF_INET;
251 		((struct sockaddr_in *)ss)->sin_port = htons(port);
252 		((struct sockaddr_in *)ss)->sin_addr.s_addr = INADDR_ANY;
253 		return 0;
254 	}
255 
256 	*port_str++ = 0;
257 
258 	if (strrchr(str, ':') != NULL) {
259 		/* IPv6 address contains ':' */
260 		ss->ss_family = AF_INET6;
261 		((struct sockaddr_in6 *)ss)->sin6_port = htons(atoi(port_str));
262 
263 		if (!inet_pton(ss->ss_family, str, &((struct sockaddr_in6 *)ss)->sin6_addr)) {
264 			err->len = snprintf(err->msg, err->size, "Invalid server address: '%s'\n", str);
265 			return -1;
266 		}
267 	}
268 	else {
269 		ss->ss_family = AF_INET;
270 		((struct sockaddr_in *)ss)->sin_port = htons(atoi(port_str));
271 
272 		if (*str == '*' || *str == '\0') { /* INADDR_ANY */
273 			((struct sockaddr_in *)ss)->sin_addr.s_addr = INADDR_ANY;
274 			return 0;
275 		}
276 
277 		if (!inet_pton(ss->ss_family, str, &((struct sockaddr_in *)ss)->sin_addr)) {
278 			struct hostent *he = gethostbyname(str);
279 
280 			if (he == NULL) {
281 				err->len = snprintf(err->msg, err->size, "Invalid server name: '%s'\n", str);
282 				return -1;
283 			}
284 			((struct sockaddr_in *)ss)->sin_addr = *(struct in_addr *) *(he->h_addr_list);
285 		}
286 	}
287 
288 	return 0;
289 }
290 
291 /* waits up to one second on fd <fd> for events <events> (POLLIN|POLLOUT).
292  * returns poll's status.
293  */
wait_on_fd(int fd,int events)294 int wait_on_fd(int fd, int events)
295 {
296 	struct pollfd pollfd;
297 	int ret;
298 
299 	do {
300 		pollfd.fd = fd;
301 		pollfd.events = events;
302 		ret = poll(&pollfd, 1, 1000);
303 	} while (ret == -1 && errno == EINTR);
304 
305 	return ret;
306 }
307 
tcp_set_nodelay(int sock,const char * arg)308 int tcp_set_nodelay(int sock, const char *arg)
309 {
310 	return setsockopt(sock, SOL_TCP, TCP_NODELAY, &one, sizeof(one));
311 }
312 
tcp_set_nolinger(int sock,const char * arg)313 int tcp_set_nolinger(int sock, const char *arg)
314 {
315 	return setsockopt(sock, SOL_SOCKET, SO_LINGER, (struct linger *) &nolinger, sizeof(struct linger));
316 }
317 
tcp_set_noquickack(int sock,const char * arg)318 int tcp_set_noquickack(int sock, const char *arg)
319 {
320 #ifdef TCP_QUICKACK
321 	/* warning: do not use during connect if nothing is to be sent! */
322 	return setsockopt(sock, SOL_TCP, TCP_QUICKACK, &zero, sizeof(zero));
323 #else
324 	return 0;
325 #endif
326 }
327 
328 /* Try to listen to address <sa>. Return the fd or -1 in case of error */
tcp_listen(const struct sockaddr_storage * sa,const char * arg)329 int tcp_listen(const struct sockaddr_storage *sa, const char *arg)
330 {
331 	int sock;
332 	int backlog;
333 
334 	if (arg[1])
335 		backlog = atoi(arg + 1);
336 	else
337 		backlog = 1000;
338 
339 	if (backlog < 0 || backlog > 65535) {
340 		fprintf(stderr, "backlog must be between 0 and 65535 inclusive (was %d)\n", backlog);
341 		return -1;
342 	}
343 
344 	sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
345 	if (sock < 0) {
346 		perror("socket()");
347 		return -1;
348 	}
349 
350 	if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) == -1) {
351 		perror("setsockopt(SO_REUSEADDR)");
352 		goto fail;
353 	}
354 
355 #ifdef SO_REUSEPORT
356 	if (setsockopt(sock, SOL_SOCKET, SO_REUSEPORT, (char *) &one, sizeof(one)) == -1) {
357 		perror("setsockopt(SO_REUSEPORT)");
358 		goto fail;
359 	}
360 #endif
361 	if (bind(sock, (struct sockaddr *)sa, sa->ss_family == AF_INET6 ?
362 		 sizeof(struct sockaddr_in6) : sizeof(struct sockaddr_in)) == -1) {
363 		perror("bind");
364 		goto fail;
365 	}
366 
367 	if (listen(sock, backlog) == -1) {
368 		perror("listen");
369 		goto fail;
370 	}
371 
372 	return sock;
373  fail:
374 	close(sock);
375 	return -1;
376 }
377 
378 /* accepts a socket from listening socket <sock>, and returns it (or -1 in case of error) */
tcp_accept(int sock,const char * arg)379 int tcp_accept(int sock, const char *arg)
380 {
381 	int count;
382 	int newsock;
383 
384 	if (arg[1])
385 		count = atoi(arg + 1);
386 	else
387 		count = 1;
388 
389 	if (count <= 0) {
390 		fprintf(stderr, "accept count must be > 0 or unset (was %d)\n", count);
391 		return -1;
392 	}
393 
394 	do {
395 		newsock = accept(sock, NULL, NULL);
396 		if (newsock < 0) { // TODO: improve error handling
397 			if (errno == EINTR || errno == EAGAIN || errno == ECONNABORTED)
398 				continue;
399 			perror("accept()");
400 			break;
401 		}
402 
403 		if (count > 1)
404 			close(newsock);
405 		count--;
406 	} while (count > 0);
407 
408 	fcntl(newsock, F_SETFL, O_NONBLOCK);
409 	return newsock;
410 }
411 
412 /* Try to establish a new connection to <sa>. Return the fd or -1 in case of error */
tcp_connect(const struct sockaddr_storage * sa,const char * arg)413 int tcp_connect(const struct sockaddr_storage *sa, const char *arg)
414 {
415 	int sock;
416 
417 	sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
418 	if (sock < 0)
419 		return -1;
420 
421 	if (fcntl(sock, F_SETFL, O_NONBLOCK) == -1)
422 		goto fail;
423 
424 	if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) == -1)
425 		goto fail;
426 
427 	if (connect(sock, (const struct sockaddr *)sa, sizeof(struct sockaddr_in)) < 0) {
428 		if (errno != EINPROGRESS)
429 			goto fail;
430 	}
431 
432 	return sock;
433  fail:
434 	close(sock);
435 	return -1;
436 }
437 
438 /* receives N bytes from the socket and returns 0 (or -1 in case of a recv
439  * error, or -2 in case of an argument error). When no arg is passed, receives
440  * anything and stops. Otherwise reads the requested amount of data. 0 means
441  * read as much as possible.
442  */
tcp_recv(int sock,const char * arg)443 int tcp_recv(int sock, const char *arg)
444 {
445 	int count = -1; // stop at first read
446 	int ret;
447 	int max;
448 
449 	if (arg[1]) {
450 		count = atoi(arg + 1);
451 		if (count < 0) {
452 			fprintf(stderr, "recv count must be >= 0 or unset (was %d)\n", count);
453 			return -2;
454 		}
455 	}
456 
457 	while (1) {
458 		max = (count > 0) ? count : INT_MAX;
459 		if (max > sizeof(trash))
460 			max = sizeof(trash);
461 		ret = recv(sock, trash, max, MSG_NOSIGNAL | MSG_TRUNC);
462 		if (ret < 0) {
463 			if (errno == EINTR)
464 				continue;
465 			if (errno != EAGAIN) {
466 				dolog("recv %d\n", ret);
467 				return -1;
468 			}
469 			while (!wait_on_fd(sock, POLLIN));
470 			continue;
471 		}
472 		dolog("recv %d\n", ret);
473 		if (!ret)
474 			break;
475 
476 		if (!count)
477 			continue;
478 		else if (count > 0)
479 			count -= ret;
480 
481 		if (count <= 0)
482 			break;
483 	}
484 
485 	return 0;
486 }
487 
488 /* Sends N bytes to the socket and returns 0 (or -1 in case of send error, -2
489  * in case of an argument error. If the byte count is not set, sends only one
490  * block. Sending zero means try to send forever. If the argument starts with
491  * ':' then whatever follows is interpreted as the payload to be sent as-is.
492  * Escaped characters '\r', '\n', '\t' and '\\' are detected and converted. In
493  * this case, blocks must be small so that send() doesn't fragment them, as
494  * they will be put into the trash and expected to be sent at once.
495  */
tcp_send(int sock,const char * arg)496 int tcp_send(int sock, const char *arg)
497 {
498 	int count = -1; // stop after first block
499 	int ret;
500 
501 	if (arg[1] == ':') {
502 		count = unescape(trash, sizeof(trash), arg + 2);
503 	} else if (arg[1]) {
504 		count = atoi(arg + 1);
505 		if (count < 0) {
506 			fprintf(stderr, "send count must be >= 0 or unset (was %d)\n", count);
507 			return -2;
508 		}
509 	}
510 
511 	while (1) {
512 		ret = send(sock, trash,
513 		           (count > 0) && (count < sizeof(trash)) ? count : sizeof(trash),
514 		           MSG_NOSIGNAL | ((count > sizeof(trash)) ? MSG_MORE : 0));
515 		if (ret < 0) {
516 			if (errno == EINTR)
517 				continue;
518 			if (errno != EAGAIN) {
519 				dolog("send %d\n", ret);
520 				return -1;
521 			}
522 			while (!wait_on_fd(sock, POLLOUT));
523 			continue;
524 		}
525 		dolog("send %d\n", ret);
526 		if (!count)
527 			continue;
528 		else if (count > 0)
529 			count -= ret;
530 
531 		if (count <= 0)
532 			break;
533 	}
534 
535 	return 0;
536 }
537 
538 /* echoes N bytes to the socket and returns 0 (or -1 in case of error). If not
539  * set, echoes only the first block. Zero means forward forever.
540  */
tcp_echo(int sock,const char * arg)541 int tcp_echo(int sock, const char *arg)
542 {
543 	int count = -1; // echo forever
544 	int ret;
545 	int rcvd;
546 
547 	if (arg[1]) {
548 		count = atoi(arg + 1);
549 		if (count < 0) {
550 			fprintf(stderr, "send count must be >= 0 or unset (was %d)\n", count);
551 			return -1;
552 		}
553 	}
554 
555 	rcvd = 0;
556 	while (1) {
557 		if (rcvd <= 0) {
558 			/* no data pending */
559 			rcvd = recv(sock, trash, (count > 0) && (count < sizeof(trash)) ? count : sizeof(trash), MSG_NOSIGNAL);
560 			if (rcvd < 0) {
561 				if (errno == EINTR)
562 					continue;
563 				if (errno != EAGAIN) {
564 					dolog("recv %d\n", rcvd);
565 					return -1;
566 				}
567 				while (!wait_on_fd(sock, POLLIN));
568 				continue;
569 			}
570 			dolog("recv %d\n", rcvd);
571 			if (!rcvd)
572 				break;
573 		}
574 		else {
575 			/* some data still pending */
576 			ret = send(sock, trash, rcvd, MSG_NOSIGNAL | ((count > rcvd) ? MSG_MORE : 0));
577 			if (ret < 0) {
578 				if (errno == EINTR)
579 					continue;
580 				if (errno != EAGAIN) {
581 					dolog("send %d\n", ret);
582 					return -1;
583 				}
584 				while (!wait_on_fd(sock, POLLOUT));
585 				continue;
586 			}
587 			dolog("send %d\n", ret);
588 			rcvd -= ret;
589 			if (rcvd)
590 				continue;
591 
592 			if (!count)
593 				continue;
594 			else if (count > 0)
595 				count -= ret;
596 
597 			if (count <= 0)
598 				break;
599 		}
600 	}
601 	return 0;
602 }
603 
604 /* waits for an event on the socket, usually indicates an accept for a
605  * listening socket and a connect for an outgoing socket.
606  */
tcp_wait(int sock,const char * arg)607 int tcp_wait(int sock, const char *arg)
608 {
609 	struct pollfd pollfd;
610 	int delay = -1; // wait forever
611 	int ret;
612 
613 	if (arg[1]) {
614 		delay = atoi(arg + 1);
615 		if (delay < 0) {
616 			fprintf(stderr, "wait time must be >= 0 or unset (was %d)\n", delay);
617 			return -1;
618 		}
619 	}
620 
621 	/* FIXME: this doesn't take into account delivered signals */
622 	do {
623 		pollfd.fd = sock;
624 		pollfd.events = POLLIN | POLLOUT;
625 		ret = poll(&pollfd, 1, delay);
626 	} while (ret == -1 && errno == EINTR);
627 
628 	if (ret > 0 && pollfd.revents & POLLERR)
629 		return -1;
630 
631 	return 0;
632 }
633 
634 /* waits for the input data to be present */
tcp_wait_in(int sock,const char * arg)635 int tcp_wait_in(int sock, const char *arg)
636 {
637 	struct pollfd pollfd;
638 	int ret;
639 
640 	do {
641 		pollfd.fd = sock;
642 		pollfd.events = POLLIN;
643 		ret = poll(&pollfd, 1, 1000);
644 	} while (ret == -1 && errno == EINTR);
645 
646 	if (ret > 0 && pollfd.revents & POLLERR)
647 		return -1;
648 
649 	return 0;
650 }
651 
652 /* waits for the output queue to be empty */
tcp_wait_out(int sock,const char * arg)653 int tcp_wait_out(int sock, const char *arg)
654 {
655 	struct pollfd pollfd;
656 	int ret;
657 
658 	do {
659 		pollfd.fd = sock;
660 		pollfd.events = POLLOUT;
661 		ret = poll(&pollfd, 1, 1000);
662 	} while (ret == -1 && errno == EINTR);
663 
664 	if (ret > 0 && pollfd.revents & POLLERR)
665 		return -1;
666 
667 	/* Now wait for data to leave the socket */
668 	do {
669 		if (ioctl(sock, TIOCOUTQ, &ret) < 0)
670 			return -1;
671 	} while (ret > 0);
672 	return 0;
673 }
674 
675 /* delays processing for <time> milliseconds, 100 by default */
tcp_pause(int sock,const char * arg)676 int tcp_pause(int sock, const char *arg)
677 {
678 	struct pollfd pollfd;
679 	int delay = 100;
680 	int ret;
681 
682 	if (arg[1]) {
683 		delay = atoi(arg + 1);
684 		if (delay < 0) {
685 			fprintf(stderr, "wait time must be >= 0 or unset (was %d)\n", delay);
686 			return -1;
687 		}
688 	}
689 
690 	usleep(delay * 1000);
691 	return 0;
692 }
693 
694 /* forks another process while respecting the limit imposed in argument (1 by
695  * default). Will wait for another process to exit before creating a new one.
696  * Returns the value of the fork() syscall, ie 0 for the child, non-zero for
697  * the parent, -1 for an error.
698  */
tcp_fork(int sock,const char * arg)699 int tcp_fork(int sock, const char *arg)
700 {
701 	int max = 1;
702 	int ret;
703 
704 	if (arg[1]) {
705 		max = atoi(arg + 1);
706 		if (max <= 0) {
707 			fprintf(stderr, "max process must be > 0 or unset (was %d)\n", max);
708 			return -1;
709 		}
710 	}
711 
712 	while (nbproc >= max)
713 		poll(NULL, 0, 1000);
714 
715 	ret = fork();
716 	if (ret > 0)
717 		__sync_add_and_fetch(&nbproc, 1);
718 	return ret;
719 }
720 
main(int argc,char ** argv)721 int main(int argc, char **argv)
722 {
723 	struct sockaddr_storage ss;
724 	struct err_msg err;
725 	const char *arg0;
726 	int loop_arg;
727 	int arg;
728 	int ret;
729 	int sock;
730 
731 	arg0 = argv[0];
732 
733 	while (argc > 1 && argv[1][0] == '-') {
734 		argc--; argv++;
735 		if (strcmp(argv[0], "-t") == 0)
736 			showtime++;
737 		else if (strcmp(argv[0], "-tt") == 0)
738 			showtime += 2;
739 		else if (strcmp(argv[0], "-ttt") == 0)
740 			showtime += 3;
741 		else if (strcmp(argv[0], "-v") == 0)
742 			verbose ++;
743 		else if (strcmp(argv[0], "--") == 0)
744 			break;
745 		else
746 			usage(1, arg0);
747 	}
748 
749 	if (argc < 2)
750 		usage(1, arg0);
751 
752 	pid = getpid();
753 	signal(SIGCHLD, sig_handler);
754 
755 	if (addr_to_ss(argv[1], &ss, &err) < 0)
756 		die(1, "%s\n", err.msg);
757 
758 	gettimeofday(&start_time, NULL);
759 
760 	sock = -1;
761 	loop_arg = 2;
762 	for (arg = loop_arg; arg < argc; arg++) {
763 		switch (argv[arg][0]) {
764 		case 'L':
765 			/* silently ignore existing connections */
766 			if (sock == -1)
767 				sock = tcp_listen(&ss, argv[arg]);
768 			if (sock < 0)
769 				die(1, "Fatal: tcp_listen() failed.\n");
770 			break;
771 
772 		case 'C':
773 			/* silently ignore existing connections */
774 			if (sock == -1)
775 				sock = tcp_connect(&ss, argv[arg]);
776 			if (sock < 0)
777 				die(1, "Fatal: tcp_connect() failed.\n");
778 			dolog("connect\n");
779 			break;
780 
781 		case 'A':
782 			if (sock < 0)
783 				die(1, "Fatal: tcp_accept() on non-socket.\n");
784 			sock = tcp_accept(sock, argv[arg]);
785 			if (sock < 0)
786 				die(1, "Fatal: tcp_accept() failed.\n");
787 			dolog("accept\n");
788 			loop_arg = arg + 1; // cannot loop before accept()
789 			break;
790 
791 		case 'T':
792 			if (sock < 0)
793 				die(1, "Fatal: tcp_set_nodelay() on non-socket.\n");
794 			if (tcp_set_nodelay(sock, argv[arg]) < 0)
795 				die(1, "Fatal: tcp_set_nodelay() failed.\n");
796 			break;
797 
798 		case 'G':
799 			if (sock < 0)
800 				die(1, "Fatal: tcp_set_nolinger() on non-socket.\n");
801 			if (tcp_set_nolinger(sock, argv[arg]) < 0)
802 				die(1, "Fatal: tcp_set_nolinger() failed.\n");
803 			break;
804 
805 		case 'Q':
806 			if (sock < 0)
807 				die(1, "Fatal: tcp_set_noquickack() on non-socket.\n");
808 			if (tcp_set_noquickack(sock, argv[arg]) < 0)
809 				die(1, "Fatal: tcp_set_noquickack() failed.\n");
810 			break;
811 
812 		case 'R':
813 			if (sock < 0)
814 				die(1, "Fatal: tcp_recv() on non-socket.\n");
815 			ret = tcp_recv(sock, argv[arg]);
816 			if (ret < 0) {
817 				if (ret == -1) // usually ECONNRESET, silently exit
818 					die(0, NULL);
819 				die(1, "Fatal: tcp_recv() failed.\n");
820 			}
821 			break;
822 
823 		case 'S':
824 			if (sock < 0)
825 				die(1, "Fatal: tcp_send() on non-socket.\n");
826 			ret = tcp_send(sock, argv[arg]);
827 			if (ret < 0) {
828 				if (ret == -1) // usually a broken pipe, silently exit
829 					die(0, NULL);
830 				die(1, "Fatal: tcp_send() failed.\n");
831 			}
832 			break;
833 
834 		case 'E':
835 			if (sock < 0)
836 				die(1, "Fatal: tcp_echo() on non-socket.\n");
837 			if (tcp_echo(sock, argv[arg]) < 0)
838 				die(1, "Fatal: tcp_echo() failed.\n");
839 			break;
840 
841 		case 'P':
842 			if (tcp_pause(sock, argv[arg]) < 0)
843 				die(1, "Fatal: tcp_pause() failed.\n");
844 			break;
845 
846 		case 'W':
847 			if (sock < 0)
848 				die(1, "Fatal: tcp_wait() on non-socket.\n");
849 			if (tcp_wait(sock, argv[arg]) < 0)
850 				die(1, "Fatal: tcp_wait() failed.\n");
851 			dolog("ready_any\n");
852 			break;
853 
854 		case 'I':
855 			if (sock < 0)
856 				die(1, "Fatal: tcp_wait_in() on non-socket.\n");
857 			if (tcp_wait_in(sock, argv[arg]) < 0)
858 				die(1, "Fatal: tcp_wait_in() failed.\n");
859 			dolog("ready_in\n");
860 			break;
861 
862 		case 'O':
863 			if (sock < 0)
864 				die(1, "Fatal: tcp_wait_out() on non-socket.\n");
865 			if (tcp_wait_out(sock, argv[arg]) < 0)
866 				die(1, "Fatal: tcp_wait_out() failed.\n");
867 			dolog("ready_out\n");
868 			break;
869 
870 		case 'K':
871 			if (sock < 0 || close(sock) < 0)
872 				die(1, "Fatal: close() on non-socket.\n");
873 			dolog("close\n");
874 			sock = -1;
875 			break;
876 
877 		case 'F':
878 			/* ignore errors on shutdown() as they are common */
879 			if (sock >= 0)
880 				shutdown(sock, SHUT_WR);
881 			dolog("shutdown\n");
882 			break;
883 
884 		case 'N':
885 			ret = tcp_fork(sock, argv[arg]);
886 			if (ret < 0)
887 				die(1, "Fatal: fork() failed.\n");
888 			if (ret > 0) {
889 				/* loop back to first arg */
890 				arg = loop_arg - 1;
891 				continue;
892 			}
893 			/* OK we're in the child, let's continue */
894 			pid = getpid();
895 			loop_arg = arg + 1;
896 			break;
897 
898 		case 'J': // jump back to oldest post-fork action
899 			arg = loop_arg - 1;
900 			continue;
901 
902 		default:
903 			usage(1, arg0);
904 		}
905 	}
906 	return 0;
907 }
908