summaryrefslogtreecommitdiffstats
path: root/dev/tcploop/tcploop.c
diff options
context:
space:
mode:
Diffstat (limited to 'dev/tcploop/tcploop.c')
-rw-r--r--dev/tcploop/tcploop.c1055
1 files changed, 1055 insertions, 0 deletions
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;
+}
generated by cgit v1.2.3 (git 2.39.1) at 2024-11-12 11:27:55 +0000