1021 lines
24 KiB
C
1021 lines
24 KiB
C
/*
|
|
* 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.
|
|
*/
|
|
|
|
#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/wait.h>
|
|
#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>
|
|
|
|
#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 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"
|
|
" -t|-tt|-ttt : show time (msec / relative / absolute)\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));
|
|
|
|
/* 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 one second on fd <fd> for events <events> (POLLIN|POLLOUT).
|
|
* returns poll's status.
|
|
*/
|
|
int wait_on_fd(int fd, int events)
|
|
{
|
|
struct pollfd pollfd;
|
|
int ret;
|
|
|
|
do {
|
|
pollfd.fd = fd;
|
|
pollfd.events = events;
|
|
ret = poll(&pollfd, 1, 1000);
|
|
} 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()
|
|
{
|
|
int sock;
|
|
|
|
sock = socket(AF_INET, 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();
|
|
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();
|
|
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));
|
|
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));
|
|
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));
|
|
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));
|
|
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)
|
|
{
|
|
struct pollfd pollfd;
|
|
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 */
|
|
do {
|
|
pollfd.fd = sock;
|
|
pollfd.events = POLLIN | POLLOUT;
|
|
ret = poll(&pollfd, 1, delay);
|
|
} while (ret == -1 && errno == EINTR);
|
|
|
|
if (ret > 0 && pollfd.revents & POLLERR)
|
|
return -1;
|
|
|
|
return 0;
|
|
}
|
|
|
|
/* waits for the input data to be present */
|
|
int tcp_wait_in(int sock, const char *arg)
|
|
{
|
|
struct pollfd pollfd;
|
|
int ret;
|
|
|
|
do {
|
|
pollfd.fd = sock;
|
|
pollfd.events = POLLIN;
|
|
ret = poll(&pollfd, 1, 1000);
|
|
} while (ret == -1 && errno == EINTR);
|
|
|
|
if (ret > 0 && pollfd.revents & POLLERR)
|
|
return -1;
|
|
|
|
return 0;
|
|
}
|
|
|
|
/* waits for the output queue to be empty */
|
|
int tcp_wait_out(int sock, const char *arg)
|
|
{
|
|
struct pollfd pollfd;
|
|
int ret;
|
|
|
|
do {
|
|
pollfd.fd = sock;
|
|
pollfd.events = POLLOUT;
|
|
ret = poll(&pollfd, 1, 1000);
|
|
} while (ret == -1 && errno == EINTR);
|
|
|
|
if (ret > 0 && pollfd.revents & POLLERR)
|
|
return -1;
|
|
|
|
/* 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], "-v") == 0)
|
|
verbose ++;
|
|
else if (strcmp(argv[0], "-u") == 0) {
|
|
sock_type = SOCK_DGRAM;
|
|
sock_proto = IPPROTO_UDP;
|
|
}
|
|
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;
|
|
}
|