From b46aad6df449445a9fc4aa7b32bd40005438e3f7 Mon Sep 17 00:00:00 2001
From: Daniel Baumann <daniel.baumann@progress-linux.org>
Date: Sat, 13 Apr 2024 14:18:05 +0200
Subject: Adding upstream version 2.9.5.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
---
 dev/tcploop/tcploop.c | 1055 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 1055 insertions(+)
 create mode 100644 dev/tcploop/tcploop.c

(limited to 'dev/tcploop/tcploop.c')

diff --git a/dev/tcploop/tcploop.c b/dev/tcploop/tcploop.c
new file mode 100644
index 0000000..c091f10
--- /dev/null
+++ b/dev/tcploop/tcploop.c
@@ -0,0 +1,1055 @@
+/*
+ * TCP client and server for bug hunting
+ *
+ * Copyright (C) 2016 Willy Tarreau <w@1wt.eu>
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining
+ * a copy of this software and associated documentation files (the
+ * "Software"), to deal in the Software without restriction, including
+ * without limitation the rights to use, copy, modify, merge, publish,
+ * distribute, sublicense, and/or sell copies of the Software, and to
+ * permit persons to whom the Software is furnished to do so, subject to
+ * the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be
+ * included in all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
+ * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+ * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
+ * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+ * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+ * OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+#define _GNU_SOURCE     // for POLLRDHUP
+#include <sys/resource.h>
+#include <sys/select.h>
+#include <sys/types.h>
+#include <sys/socket.h>
+#include <sys/stat.h>
+#include <sys/time.h>
+#include <sys/ioctl.h>
+#include <sys/un.h>
+#include <sys/wait.h>
+
+#ifdef __linux__
+#include <sys/epoll.h>
+#endif
+
+#include <arpa/inet.h>
+#include <netinet/in.h>
+#include <netinet/tcp.h>
+
+#include <ctype.h>
+#include <errno.h>
+#include <fcntl.h>
+#include <limits.h>
+#include <netdb.h>
+#include <poll.h>
+#include <signal.h>
+#include <stdarg.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <time.h>
+#include <unistd.h>
+
+/* for OSes which don't have it */
+#ifndef POLLRDHUP
+#define POLLRDHUP 0
+#endif
+
+#ifndef MSG_MORE
+#define MSG_MORE 0
+#endif
+
+struct err_msg {
+	int size;
+	int len;
+	char msg[0];
+};
+
+const int zero = 0;
+const int one = 1;
+const struct linger nolinger = { .l_onoff = 1, .l_linger = 0 };
+
+#define TRASH_SIZE 65536
+static char trash[TRASH_SIZE];
+
+volatile int nbproc = 0;
+static struct timeval start_time;
+static int showtime;
+static int verbose;
+static int use_epoll;
+static int pid;
+static int sock_type = SOCK_STREAM;
+static int sock_proto = IPPROTO_TCP;
+
+
+/* display the message and exit with the code */
+__attribute__((noreturn)) void die(int code, const char *format, ...)
+{
+	va_list args;
+
+	if (format) {
+		va_start(args, format);
+		vfprintf(stderr, format, args);
+		va_end(args);
+	}
+	exit(code);
+}
+
+/* display the usage message and exit with the code */
+__attribute__((noreturn)) void usage(int code, const char *arg0)
+{
+	die(code,
+	    "Usage : %s [options]* [<ip>:]port [<action>*]\n"
+	    "\n"
+	    "options :\n"
+	    "  -v           : verbose\n"
+	    "  -u           : use UDP instead of TCP (limited)\n"
+	    "  -U           : use UNIX instead of TCP (limited, addr must have one '/')\n"
+	    "  -t|-tt|-ttt  : show time (msec / relative / absolute)\n"
+	    "  -e           : use epoll instead of poll on Linux\n"
+	    "actions :\n"
+	    "  A[<count>]   : Accepts <count> incoming sockets and closes count-1\n"
+	    "                 Note: fd=accept(fd)\n"
+	    "  B[[ip]:port] : Bind a new socket to ip:port or default one if unspecified.\n"
+	    "                 Note: fd=socket,bind(fd)\n"
+	    "  C[[ip]:port] : Connects to ip:port or default ones if unspecified.\n"
+	    "                 Note: fd=socket,connect(fd)\n"
+	    "  D            : Disconnect (connect to AF_UNSPEC)\n"
+	    "  E[<size>]    : Echo this amount of bytes. 0=infinite. unset=any amount.\n"
+	    "  F            : FIN : shutdown(SHUT_WR)\n"
+	    "  G            : disable lingering\n"
+	    "  I            : wait for Input data to be present (POLLIN)\n"
+	    "  J            : Jump back to oldest post-fork/post-accept action\n"
+	    "  K            : kill the connection and go on with next operation\n"
+	    "  L[<backlog>] : Listens to ip:port and optionally sets backlog\n"
+	    "                 Note: fd=socket,bind(fd),listen(fd)\n"
+	    "  N<max>       : fork New process, limited to <max> concurrent (default 1)\n"
+	    "  O            : wait for Output queue to be empty (POLLOUT + TIOCOUTQ)\n"
+	    "  P[<time>]    : Pause for <time> ms (100 by default)\n"
+	    "  Q            : disable TCP Quick-ack\n"
+	    "  R[<size>]    : Read this amount of bytes. 0=infinite. unset=any amount.\n"
+	    "  S[<size>]    : Send this amount of bytes. 0=infinite. unset=any amount.\n"
+	    "  S:<string>   : Send this exact string. \\r, \\n, \\t, \\\\ supported.\n"
+	    "  T            : set TCP_NODELAY\n"
+	    "  W[<time>]    : Wait for any event on the socket, maximum <time> ms\n"
+	    "  X[i|o|e]* ** : execvp() next args passing socket as stdin/stdout/stderr.\n"
+	    "                 If i/o/e present, only stdin/out/err are mapped to socket.\n"
+	    "  r            : shutr : shutdown(SHUT_RD) (pauses a listener or ends recv)\n"
+	    "\n"
+	    "It's important to note that a single FD is used at once and that Accept\n"
+	    "replaces the listening FD with the accepted one. Thus always do it after\n"
+	    "a fork if other connections have to be accepted.\n"
+	    "\n"
+	    "After a fork, we loop back to the beginning and silently skip L/C if the\n"
+	    "main socket already exists.\n"
+	    "\n"
+	    "Example dummy HTTP request drain server :\n"
+	    "   tcploop 8001 L W N20 A R S10 [ F K ]\n"
+	    "\n"
+	    "Example large bandwidth HTTP request drain server :\n"
+	    "   tcploop 8001 L W N20 A R S0 [ F K ]\n"
+	    "\n"
+	    "Example TCP client with pauses at each step :\n"
+	    "   tcploop 8001 C T W P100 S10 O P100 R S10 O R G K\n"
+	    "\n"
+	    "Simple chargen server :\n"
+	    "   tcploop 8001 L A Xo cat /dev/zero\n"
+	    "\n"
+	    "Simple telnet server :\n"
+	    "   tcploop 8001 L W N A X /usr/sbin/in.telnetd\n"
+	    "", arg0);
+}
+
+void dolog(const char *format, ...)
+{
+	struct timeval date, tv;
+	int delay;
+	va_list args;
+
+	if (!verbose)
+		return;
+
+	if (showtime) {
+		gettimeofday(&date, NULL);
+		switch (showtime) {
+		case 1: // [msec] relative
+			delay = (date.tv_sec - start_time.tv_sec) * 1000000 + date.tv_usec - start_time.tv_usec;
+			fprintf(stderr, "[%d] ", delay / 1000);
+			break;
+		case 2: // [sec.usec] relative
+			tv.tv_usec = date.tv_usec - start_time.tv_usec;
+			tv.tv_sec  = date.tv_sec  - start_time.tv_sec;
+			if ((signed)tv.tv_sec > 0) {
+				if ((signed)tv.tv_usec < 0) {
+					tv.tv_usec += 1000000;
+					tv.tv_sec--;
+				}
+			} else if (tv.tv_sec == 0) {
+				if ((signed)tv.tv_usec < 0)
+					tv.tv_usec = 0;
+			} else {
+				tv.tv_sec = 0;
+				tv.tv_usec = 0;
+			}
+			fprintf(stderr, "[%d.%06d] ", (int)tv.tv_sec, (int)tv.tv_usec);
+			break;
+		default: // [sec.usec] absolute
+			fprintf(stderr, "[%d.%06d] ", (int)date.tv_sec, (int)date.tv_usec);
+			break;
+		}
+	}
+
+	fprintf(stderr, "%5d ", pid);
+
+	va_start(args, format);
+	vfprintf(stderr, format, args);
+	va_end(args);
+}
+
+/* convert '\n', '\t', '\r', '\\' to their respective characters */
+int unescape(char *out, int size, const char *in)
+{
+	int len;
+
+	for (len = 0; len < size && *in; in++, out++, len++) {
+		if (*in == '\\') {
+			switch (in[1]) {
+			case  'n' : *out = '\n'; in++; continue;
+			case  't' : *out = '\t'; in++; continue;
+			case  'r' : *out = '\r'; in++; continue;
+			case '\\' : *out = '\\'; in++; continue;
+			default   : break;
+			}
+		}
+		*out = *in;
+	}
+	return len;
+}
+
+struct err_msg *alloc_err_msg(int size)
+{
+	struct err_msg *err;
+
+	err = malloc(sizeof(*err) + size);
+	if (err) {
+		err->len = 0;
+		err->size = size;
+	}
+	return err;
+}
+
+void sig_handler(int sig)
+{
+	if (sig == SIGCHLD) {
+		while (waitpid(-1, NULL, WNOHANG) > 0)
+			__sync_sub_and_fetch(&nbproc, 1);
+	}
+}
+
+/* converts str in the form [[<ipv4>|<ipv6>|<hostname>]:]port to struct sockaddr_storage.
+ * Returns < 0 with err set in case of error.
+ */
+int addr_to_ss(const char *str, struct sockaddr_storage *ss, struct err_msg *err)
+{
+	char *port_str;
+	int port;
+
+	memset(ss, 0, sizeof(*ss));
+
+	/* if there's a slash it's a unix socket */
+	if (strchr(str, '/')) {
+		((struct sockaddr_un *)ss)->sun_family = AF_UNIX;
+		strncpy(((struct sockaddr_un *)ss)->sun_path, str, sizeof(((struct sockaddr_un *)ss)->sun_path) - 1);
+		((struct sockaddr_un *)ss)->sun_path[sizeof(((struct sockaddr_un *)ss)->sun_path)] = 0;
+		return 0;
+	}
+
+	/* look for the addr/port delimiter, it's the last colon. If there's no
+	 * colon, it's 0:<port>.
+	 */
+	if ((port_str = strrchr(str, ':')) == NULL) {
+		port = atoi(str);
+		if (port < 0 || port > 65535) {
+			err->len = snprintf(err->msg, err->size, "Missing/invalid port number: '%s'\n", str);
+			return -1;
+		}
+
+		ss->ss_family = AF_INET;
+		((struct sockaddr_in *)ss)->sin_port = htons(port);
+		((struct sockaddr_in *)ss)->sin_addr.s_addr = INADDR_ANY;
+		return 0;
+	}
+
+	*port_str++ = 0;
+
+	if (strrchr(str, ':') != NULL) {
+		/* IPv6 address contains ':' */
+		ss->ss_family = AF_INET6;
+		((struct sockaddr_in6 *)ss)->sin6_port = htons(atoi(port_str));
+
+		if (!inet_pton(ss->ss_family, str, &((struct sockaddr_in6 *)ss)->sin6_addr)) {
+			err->len = snprintf(err->msg, err->size, "Invalid server address: '%s'\n", str);
+			return -1;
+		}
+	}
+	else {
+		ss->ss_family = AF_INET;
+		((struct sockaddr_in *)ss)->sin_port = htons(atoi(port_str));
+
+		if (*str == '*' || *str == '\0') { /* INADDR_ANY */
+			((struct sockaddr_in *)ss)->sin_addr.s_addr = INADDR_ANY;
+			return 0;
+		}
+
+		if (!inet_pton(ss->ss_family, str, &((struct sockaddr_in *)ss)->sin_addr)) {
+			struct hostent *he = gethostbyname(str);
+
+			if (he == NULL) {
+				err->len = snprintf(err->msg, err->size, "Invalid server name: '%s'\n", str);
+				return -1;
+			}
+			((struct sockaddr_in *)ss)->sin_addr = *(struct in_addr *) *(he->h_addr_list);
+		}
+	}
+
+	return 0;
+}
+
+/* waits up to <ms> milliseconds on fd <fd> for events <events> (POLLIN|POLLRDHUP|POLLOUT).
+ * returns poll's status.
+ */
+int wait_on_fd(int fd, int events, int ms)
+{
+	struct pollfd pollfd;
+	int ret;
+
+#ifdef __linux__
+	while (use_epoll) {
+		struct epoll_event evt;
+		static int epoll_fd = -1;
+
+		if (epoll_fd == -1)
+			epoll_fd = epoll_create(1024);
+		if (epoll_fd == -1)
+			break;
+		evt.events = ((events & POLLIN) ? EPOLLIN : 0)      |
+		             ((events & POLLOUT) ? EPOLLOUT : 0)    |
+		             ((events & POLLRDHUP) ? EPOLLRDHUP : 0);
+		evt.data.fd = fd;
+		epoll_ctl(epoll_fd, EPOLL_CTL_ADD, fd, &evt);
+
+		do {
+			ret = epoll_wait(epoll_fd, &evt, 1, ms);
+		} while (ret == -1 && errno == EINTR);
+
+		evt.data.fd = fd;
+		epoll_ctl(epoll_fd, EPOLL_CTL_DEL, fd, &evt);
+		return ret;
+	}
+#endif
+
+	do {
+		pollfd.fd = fd;
+		pollfd.events = events;
+		ret = poll(&pollfd, 1, ms);
+	} while (ret == -1 && errno == EINTR);
+
+	return ret;
+}
+
+int tcp_set_nodelay(int sock, const char *arg)
+{
+	return setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));
+}
+
+int tcp_set_nolinger(int sock, const char *arg)
+{
+	return setsockopt(sock, SOL_SOCKET, SO_LINGER, (struct linger *) &nolinger, sizeof(struct linger));
+}
+
+int tcp_set_noquickack(int sock, const char *arg)
+{
+#ifdef TCP_QUICKACK
+	/* warning: do not use during connect if nothing is to be sent! */
+	return setsockopt(sock, IPPROTO_TCP, TCP_QUICKACK, &zero, sizeof(zero));
+#else
+	return 0;
+#endif
+}
+
+/* Create a new TCP socket for either listening or connecting */
+int tcp_socket(sa_family_t fam)
+{
+	int sock;
+
+	sock = socket(fam, sock_type, sock_proto);
+	if (sock < 0) {
+		perror("socket()");
+		return -1;
+	}
+
+	return sock;
+}
+
+/* Try to bind to local address <sa>. Return the fd or -1 in case of error.
+ * Supports being passed NULL for arg if none has to be passed.
+ */
+int tcp_bind(int sock, const struct sockaddr_storage *sa, const char *arg)
+{
+	struct sockaddr_storage conn_addr;
+
+	if (arg && arg[1]) {
+		struct err_msg err;
+
+		if (addr_to_ss(arg + 1, &conn_addr, &err) < 0)
+			die(1, "%s\n", err.msg);
+		sa = &conn_addr;
+	}
+
+
+	if (sock < 0) {
+		sock = tcp_socket(sa->ss_family);
+		if (sock < 0)
+			return sock;
+	}
+
+	if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) == -1) {
+		perror("setsockopt(SO_REUSEADDR)");
+		goto fail;
+	}
+
+#ifdef SO_REUSEPORT
+	if (setsockopt(sock, SOL_SOCKET, SO_REUSEPORT, (char *) &one, sizeof(one)) == -1) {
+		perror("setsockopt(SO_REUSEPORT)");
+		goto fail;
+	}
+#endif
+	if (bind(sock, (struct sockaddr *)sa, sa->ss_family == AF_INET6 ?
+		 sizeof(struct sockaddr_in6) : sizeof(struct sockaddr_in)) == -1) {
+		perror("bind");
+		goto fail;
+	}
+
+	return sock;
+ fail:
+	close(sock);
+	return -1;
+}
+
+/* Try to listen to address <sa>. Return the fd or -1 in case of error */
+int tcp_listen(int sock, const struct sockaddr_storage *sa, const char *arg)
+{
+	int backlog;
+
+	if (sock < 0) {
+		sock = tcp_bind(sock, sa, NULL);
+		if (sock < 0)
+			return sock;
+	}
+
+	if (arg[1])
+		backlog = atoi(arg + 1);
+	else
+		backlog = 1000;
+
+	if (backlog < 0 || backlog > 65535) {
+		fprintf(stderr, "backlog must be between 0 and 65535 inclusive (was %d)\n", backlog);
+		goto fail;
+	}
+
+	if (listen(sock, backlog) == -1) {
+		perror("listen");
+		goto fail;
+	}
+
+	return sock;
+ fail:
+	close(sock);
+	return -1;
+}
+
+/* accepts a socket from listening socket <sock>, and returns it (or -1 in case of error) */
+int tcp_accept(int sock, const char *arg)
+{
+	int count;
+	int newsock;
+
+	if (arg[1])
+		count = atoi(arg + 1);
+	else
+		count = 1;
+
+	if (count <= 0) {
+		fprintf(stderr, "accept count must be > 0 or unset (was %d)\n", count);
+		return -1;
+	}
+
+	do {
+		newsock = accept(sock, NULL, NULL);
+		if (newsock < 0) { // TODO: improve error handling
+			if (errno == EINTR || errno == EAGAIN || errno == ECONNABORTED)
+				continue;
+			perror("accept()");
+			break;
+		}
+
+		if (count > 1)
+			close(newsock);
+		count--;
+	} while (count > 0);
+
+	fcntl(newsock, F_SETFL, O_NONBLOCK);
+	return newsock;
+}
+
+/* Try to establish a new connection to <sa>. Return the fd or -1 in case of error */
+int tcp_connect(int sock, const struct sockaddr_storage *sa, const char *arg)
+{
+	struct sockaddr_storage conn_addr;
+
+	if (arg[1]) {
+		struct err_msg err;
+
+		if (addr_to_ss(arg + 1, &conn_addr, &err) < 0)
+			die(1, "%s\n", err.msg);
+		sa = &conn_addr;
+	}
+
+	if (sock < 0) {
+		sock = tcp_socket(sa->ss_family);
+		if (sock < 0)
+			return sock;
+	}
+
+	if (fcntl(sock, F_SETFL, O_NONBLOCK) == -1)
+		goto fail;
+
+	if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) == -1)
+		goto fail;
+
+	if (connect(sock, (const struct sockaddr *)sa, sizeof(struct sockaddr_in)) < 0) {
+		if (errno != EINPROGRESS)
+			goto fail;
+	}
+
+	return sock;
+ fail:
+	close(sock);
+	return -1;
+}
+
+/* Try to disconnect by connecting to AF_UNSPEC. Return >=0 on success, -1 in case of error */
+int tcp_disconnect(int sock)
+{
+	const struct sockaddr sa = { .sa_family = AF_UNSPEC };
+
+	return connect(sock, &sa, sizeof(sa));
+}
+
+/* receives N bytes from the socket and returns 0 (or -1 in case of a recv
+ * error, or -2 in case of an argument error). When no arg is passed, receives
+ * anything and stops. Otherwise reads the requested amount of data. 0 means
+ * read as much as possible.
+ */
+int tcp_recv(int sock, const char *arg)
+{
+	int count = -1; // stop at first read
+	int ret;
+	int max;
+
+	if (arg[1]) {
+		count = atoi(arg + 1);
+		if (count < 0) {
+			fprintf(stderr, "recv count must be >= 0 or unset (was %d)\n", count);
+			return -2;
+		}
+	}
+
+	while (1) {
+		max = (count > 0) ? count : INT_MAX;
+		if (max > sizeof(trash))
+			max = sizeof(trash);
+		ret = recv(sock, trash, max, MSG_NOSIGNAL | MSG_TRUNC);
+		if (ret < 0) {
+			if (errno == EINTR)
+				continue;
+			if (errno != EAGAIN) {
+				dolog("recv %d\n", ret);
+				return -1;
+			}
+			while (!wait_on_fd(sock, POLLIN | POLLRDHUP, 1000));
+			continue;
+		}
+		dolog("recv %d\n", ret);
+		if (!ret)
+			break;
+
+		if (!count)
+			continue;
+		else if (count > 0)
+			count -= ret;
+
+		if (count <= 0)
+			break;
+	}
+
+	return 0;
+}
+
+/* Sends N bytes to the socket and returns 0 (or -1 in case of send error, -2
+ * in case of an argument error. If the byte count is not set, sends only one
+ * block. Sending zero means try to send forever. If the argument starts with
+ * ':' then whatever follows is interpreted as the payload to be sent as-is.
+ * Escaped characters '\r', '\n', '\t' and '\\' are detected and converted. In
+ * this case, blocks must be small so that send() doesn't fragment them, as
+ * they will be put into the trash and expected to be sent at once.
+ */
+int tcp_send(int sock, const char *arg)
+{
+	int count = -1; // stop after first block
+	int ret;
+
+	if (arg[1] == ':') {
+		count = unescape(trash, sizeof(trash), arg + 2);
+	} else if (arg[1]) {
+		count = atoi(arg + 1);
+		if (count < 0) {
+			fprintf(stderr, "send count must be >= 0 or unset (was %d)\n", count);
+			return -2;
+		}
+	}
+
+	while (1) {
+		ret = send(sock, trash,
+		           (count > 0) && (count < sizeof(trash)) ? count : sizeof(trash),
+		           MSG_NOSIGNAL | ((count > sizeof(trash)) ? MSG_MORE : 0));
+		if (ret < 0) {
+			if (errno == EINTR)
+				continue;
+			if (errno != EAGAIN) {
+				dolog("send %d\n", ret);
+				return -1;
+			}
+			while (!wait_on_fd(sock, POLLOUT, 1000));
+			continue;
+		}
+		dolog("send %d\n", ret);
+		if (!count)
+			continue;
+		else if (count > 0)
+			count -= ret;
+
+		if (count <= 0)
+			break;
+	}
+
+	return 0;
+}
+
+/* echoes N bytes to the socket and returns 0 (or -1 in case of error). If not
+ * set, echoes only the first block. Zero means forward forever.
+ */
+int tcp_echo(int sock, const char *arg)
+{
+	int count = -1; // echo forever
+	int ret;
+	int rcvd;
+
+	if (arg[1]) {
+		count = atoi(arg + 1);
+		if (count < 0) {
+			fprintf(stderr, "send count must be >= 0 or unset (was %d)\n", count);
+			return -1;
+		}
+	}
+
+	rcvd = 0;
+	while (1) {
+		if (rcvd <= 0) {
+			/* no data pending */
+			rcvd = recv(sock, trash, (count > 0) && (count < sizeof(trash)) ? count : sizeof(trash), MSG_NOSIGNAL);
+			if (rcvd < 0) {
+				if (errno == EINTR)
+					continue;
+				if (errno != EAGAIN) {
+					dolog("recv %d\n", rcvd);
+					return -1;
+				}
+				while (!wait_on_fd(sock, POLLIN | POLLRDHUP, 1000));
+				continue;
+			}
+			dolog("recv %d\n", rcvd);
+			if (!rcvd)
+				break;
+		}
+		else {
+			/* some data still pending */
+			ret = send(sock, trash, rcvd, MSG_NOSIGNAL | ((count > rcvd) ? MSG_MORE : 0));
+			if (ret < 0) {
+				if (errno == EINTR)
+					continue;
+				if (errno != EAGAIN) {
+					dolog("send %d\n", ret);
+					return -1;
+				}
+				while (!wait_on_fd(sock, POLLOUT, 1000));
+				continue;
+			}
+			dolog("send %d\n", ret);
+			rcvd -= ret;
+			if (rcvd)
+				continue;
+
+			if (!count)
+				continue;
+			else if (count > 0)
+				count -= ret;
+
+			if (count <= 0)
+				break;
+		}
+	}
+	return 0;
+}
+
+/* waits for an event on the socket, usually indicates an accept for a
+ * listening socket and a connect for an outgoing socket.
+ */
+int tcp_wait(int sock, const char *arg)
+{
+	int delay = -1; // wait forever
+	int ret;
+
+	if (arg[1]) {
+		delay = atoi(arg + 1);
+		if (delay < 0) {
+			fprintf(stderr, "wait time must be >= 0 or unset (was %d)\n", delay);
+			return -1;
+		}
+	}
+
+	/* FIXME: this doesn't take into account delivered signals */
+	ret = wait_on_fd(sock, POLLIN | POLLRDHUP | POLLOUT, delay);
+	if (ret < 0)
+		return ret;
+
+	return 0;
+}
+
+/* waits for the input data to be present */
+int tcp_wait_in(int sock, const char *arg)
+{
+	int ret;
+
+	ret = wait_on_fd(sock, POLLIN | POLLRDHUP, 1000);
+	if (ret < 0)
+		return ret;
+
+	return 0;
+}
+
+/* waits for the output queue to be empty */
+int tcp_wait_out(int sock, const char *arg)
+{
+	int ret;
+
+	ret = wait_on_fd(sock, POLLOUT, 1000);
+	if (ret < 0)
+		return ret;
+
+	/* Now wait for data to leave the socket */
+	do {
+		if (ioctl(sock, TIOCOUTQ, &ret) < 0)
+			return -1;
+	} while (ret > 0);
+	return 0;
+}
+
+/* delays processing for <time> milliseconds, 100 by default */
+int tcp_pause(int sock, const char *arg)
+{
+	int delay = 100;
+
+	if (arg[1]) {
+		delay = atoi(arg + 1);
+		if (delay < 0) {
+			fprintf(stderr, "wait time must be >= 0 or unset (was %d)\n", delay);
+			return -1;
+		}
+	}
+
+	usleep(delay * 1000);
+	return 0;
+}
+
+/* forks another process while respecting the limit imposed in argument (1 by
+ * default). Will wait for another process to exit before creating a new one.
+ * Returns the value of the fork() syscall, ie 0 for the child, non-zero for
+ * the parent, -1 for an error.
+ */
+int tcp_fork(int sock, const char *arg)
+{
+	int max = 1;
+	int ret;
+
+	if (arg[1]) {
+		max = atoi(arg + 1);
+		if (max <= 0) {
+			fprintf(stderr, "max process must be > 0 or unset (was %d)\n", max);
+			return -1;
+		}
+	}
+
+	while (nbproc >= max)
+		poll(NULL, 0, 1000);
+
+	ret = fork();
+	if (ret > 0)
+		__sync_add_and_fetch(&nbproc, 1);
+	return ret;
+}
+
+int main(int argc, char **argv)
+{
+	struct sockaddr_storage default_addr;
+	struct err_msg err;
+	const char *arg0;
+	int loop_arg;
+	int arg;
+	int ret;
+	int sock;
+	int errfd;
+
+	arg0 = argv[0];
+
+	while (argc > 1 && argv[1][0] == '-') {
+		argc--; argv++;
+		if (strcmp(argv[0], "-t") == 0)
+			showtime++;
+		else if (strcmp(argv[0], "-tt") == 0)
+			showtime += 2;
+		else if (strcmp(argv[0], "-ttt") == 0)
+			showtime += 3;
+		else if (strcmp(argv[0], "-e") == 0)
+			use_epoll = 1;
+		else if (strcmp(argv[0], "-v") == 0)
+			verbose ++;
+		else if (strcmp(argv[0], "-u") == 0) {
+			sock_type = SOCK_DGRAM;
+			sock_proto = IPPROTO_UDP;
+		}
+		else if (strcmp(argv[0], "-U") == 0) {
+			sock_proto = 0;
+		}
+		else if (strcmp(argv[0], "--") == 0)
+			break;
+		else
+			usage(1, arg0);
+	}
+
+	if (argc < 2)
+		usage(1, arg0);
+
+	pid = getpid();
+	signal(SIGCHLD, sig_handler);
+
+	if (addr_to_ss(argv[1], &default_addr, &err) < 0)
+		die(1, "%s\n", err.msg);
+
+	gettimeofday(&start_time, NULL);
+
+	sock = -1;
+	loop_arg = 2;
+	for (arg = loop_arg; arg < argc; arg++) {
+		switch (argv[arg][0]) {
+		case 'L':
+			sock = tcp_listen(sock, &default_addr, argv[arg]);
+			if (sock < 0)
+				die(1, "Fatal: tcp_listen() failed.\n");
+			break;
+
+		case 'B':
+			/* silently ignore existing connections */
+			sock = tcp_bind(sock, &default_addr, argv[arg]);
+			if (sock < 0)
+				die(1, "Fatal: tcp_connect() failed.\n");
+			dolog("connect\n");
+			break;
+
+		case 'C':
+			sock = tcp_connect(sock, &default_addr, argv[arg]);
+			if (sock < 0)
+				die(1, "Fatal: tcp_connect() failed.\n");
+			dolog("connect\n");
+			break;
+
+		case 'D':
+			/* silently ignore non-existing connections */
+			if (sock >= 0 && tcp_disconnect(sock) < 0)
+				die(1, "Fatal: tcp_connect() failed.\n");
+			dolog("disconnect\n");
+			break;
+
+		case 'A':
+			if (sock < 0)
+				die(1, "Fatal: tcp_accept() on non-socket.\n");
+			sock = tcp_accept(sock, argv[arg]);
+			if (sock < 0)
+				die(1, "Fatal: tcp_accept() failed.\n");
+			dolog("accept\n");
+			loop_arg = arg + 1; // cannot loop before accept()
+			break;
+
+		case 'T':
+			if (sock < 0)
+				die(1, "Fatal: tcp_set_nodelay() on non-socket.\n");
+			if (tcp_set_nodelay(sock, argv[arg]) < 0)
+				die(1, "Fatal: tcp_set_nodelay() failed.\n");
+			break;
+
+		case 'G':
+			if (sock < 0)
+				die(1, "Fatal: tcp_set_nolinger() on non-socket.\n");
+			if (tcp_set_nolinger(sock, argv[arg]) < 0)
+				die(1, "Fatal: tcp_set_nolinger() failed.\n");
+			break;
+
+		case 'Q':
+			if (sock < 0)
+				die(1, "Fatal: tcp_set_noquickack() on non-socket.\n");
+			if (tcp_set_noquickack(sock, argv[arg]) < 0)
+				die(1, "Fatal: tcp_set_noquickack() failed.\n");
+			break;
+
+		case 'R':
+			if (sock < 0)
+				die(1, "Fatal: tcp_recv() on non-socket.\n");
+			ret = tcp_recv(sock, argv[arg]);
+			if (ret < 0) {
+				if (ret == -1) // usually ECONNRESET, silently exit
+					die(0, NULL);
+				die(1, "Fatal: tcp_recv() failed.\n");
+			}
+			break;
+
+		case 'S':
+			if (sock < 0)
+				die(1, "Fatal: tcp_send() on non-socket.\n");
+			ret = tcp_send(sock, argv[arg]);
+			if (ret < 0) {
+				if (ret == -1) // usually a broken pipe, silently exit
+					die(0, NULL);
+				die(1, "Fatal: tcp_send() failed.\n");
+			}
+			break;
+
+		case 'E':
+			if (sock < 0)
+				die(1, "Fatal: tcp_echo() on non-socket.\n");
+			if (tcp_echo(sock, argv[arg]) < 0)
+				die(1, "Fatal: tcp_echo() failed.\n");
+			break;
+
+		case 'P':
+			if (tcp_pause(sock, argv[arg]) < 0)
+				die(1, "Fatal: tcp_pause() failed.\n");
+			break;
+
+		case 'W':
+			if (sock < 0)
+				die(1, "Fatal: tcp_wait() on non-socket.\n");
+			if (tcp_wait(sock, argv[arg]) < 0)
+				die(1, "Fatal: tcp_wait() failed.\n");
+			dolog("ready_any\n");
+			break;
+
+		case 'I':
+			if (sock < 0)
+				die(1, "Fatal: tcp_wait_in() on non-socket.\n");
+			if (tcp_wait_in(sock, argv[arg]) < 0)
+				die(1, "Fatal: tcp_wait_in() failed.\n");
+			dolog("ready_in\n");
+			break;
+
+		case 'O':
+			if (sock < 0)
+				die(1, "Fatal: tcp_wait_out() on non-socket.\n");
+			if (tcp_wait_out(sock, argv[arg]) < 0)
+				die(1, "Fatal: tcp_wait_out() failed.\n");
+			dolog("ready_out\n");
+			break;
+
+		case 'K':
+			if (sock < 0 || close(sock) < 0)
+				die(1, "Fatal: close() on non-socket.\n");
+			dolog("close\n");
+			sock = -1;
+			break;
+
+		case 'F':
+			/* ignore errors on shutdown() as they are common */
+			if (sock >= 0)
+				shutdown(sock, SHUT_WR);
+			dolog("shutdown(w)\n");
+			break;
+
+		case 'r':
+			/* ignore errors on shutdown() as they are common */
+			if (sock >= 0)
+				shutdown(sock, SHUT_RD);
+			dolog("shutdown(r)\n");
+			break;
+
+		case 'N':
+			ret = tcp_fork(sock, argv[arg]);
+			if (ret < 0)
+				die(1, "Fatal: fork() failed.\n");
+			if (ret > 0) {
+				/* loop back to first arg */
+				arg = loop_arg - 1;
+				continue;
+			}
+			/* OK we're in the child, let's continue */
+			pid = getpid();
+			loop_arg = arg + 1;
+			break;
+
+		case 'J': // jump back to oldest post-fork action
+			arg = loop_arg - 1;
+			continue;
+
+		case 'X': // execute command. Optionally supports redirecting only i/o/e
+			if (arg + 1 >= argc)
+				die(1, "Fatal: missing argument after %s\n", argv[arg]);
+
+			errfd = dup(2);
+			fcntl(errfd, F_SETFD, fcntl(errfd, F_GETFD, FD_CLOEXEC) | FD_CLOEXEC);
+			fcntl(sock,  F_SETFL, fcntl(sock,  F_GETFL, O_NONBLOCK) & ~O_NONBLOCK);
+			if (!argv[arg][1] || strchr(argv[arg], 'i'))
+				dup2(sock, 0);
+			if (!argv[arg][1] || strchr(argv[arg], 'o'))
+				dup2(sock, 1);
+			if (!argv[arg][1] || strchr(argv[arg], 'e'))
+				dup2(sock, 2);
+			argv += arg + 1;
+			if (execvp(argv[0], argv) == -1) {
+				int e = errno;
+
+				dup2(errfd, 2); // restore original stderr
+				close(errfd);
+				die(1, "Fatal: execvp(%s) failed : %s\n", argv[0], strerror(e));
+			}
+			break;
+		default:
+			usage(1, arg0);
+		}
+	}
+	return 0;
+}
-- 
cgit v1.2.3