| #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 <arpa/inet.h> |
| #include <netinet/in.h> |
| #include <netinet/tcp.h> |
| |
| #include <ctype.h> |
| #include <errno.h> |
| #include <fcntl.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> |
| |
| |
| 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]; |
| |
| /* display the message and exit with the code */ |
| __attribute__((noreturn)) void die(int code, const char *format, ...) |
| { |
| va_list args; |
| |
| 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 [<ip>:]port [action*]\n", arg0); |
| } |
| |
| 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; |
| } |
| |
| |
| /* 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(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, SOL_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) |
| { |
| /* warning: do not use during connect if nothing is to be sent! */ |
| return setsockopt(sock, SOL_TCP, TCP_QUICKACK, &zero, sizeof(zero)); |
| } |
| |
| /* Try to listen to address <sa>. Return the fd or -1 in case of error */ |
| int tcp_listen(const struct sockaddr_storage *sa, const char *arg) |
| { |
| int sock; |
| int backlog; |
| |
| 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); |
| return -1; |
| } |
| |
| sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP); |
| if (sock < 0) { |
| perror("socket()"); |
| return -1; |
| } |
| |
| 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; |
| } |
| |
| 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(const struct sockaddr_storage *sa, const char *arg) |
| { |
| int sock; |
| |
| sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP); |
| if (sock < 0) |
| return -1; |
| |
| 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(*sa)) < 0) { |
| if (errno != EINPROGRESS) |
| goto fail; |
| } |
| |
| return sock; |
| fail: |
| close(sock); |
| return -1; |
| } |
| |
| /* receives N bytes from the socket and returns 0 (or -1 in case of 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; |
| |
| if (arg[1]) { |
| count = atoi(arg + 1); |
| if (count < 0) { |
| fprintf(stderr, "recv count must be >= 0 or unset (was %d)\n", count); |
| return -1; |
| } |
| } |
| |
| while (1) { |
| ret = recv(sock, NULL, (count > 0) ? count : INT_MAX, MSG_NOSIGNAL | MSG_TRUNC); |
| if (ret < 0) { |
| if (errno == EINTR) |
| continue; |
| if (errno != EAGAIN) |
| return -1; |
| while (!wait_on_fd(sock, POLLIN)); |
| continue; |
| } |
| 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 error). If not |
| * set, sends only one block. Sending zero means try to send forever. |
| */ |
| int tcp_send(int sock, const char *arg) |
| { |
| int count = -1; // stop after first block |
| int ret; |
| |
| 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; |
| } |
| } |
| |
| 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) |
| return -1; |
| while (!wait_on_fd(sock, POLLOUT)); |
| continue; |
| } |
| 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) |
| return -1; |
| while (!wait_on_fd(sock, POLLIN)); |
| continue; |
| } |
| 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) |
| return -1; |
| while (!wait_on_fd(sock, POLLOUT)); |
| continue; |
| } |
| 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); |
| |
| 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); |
| 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); |
| |
| /* 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) |
| { |
| struct pollfd pollfd; |
| int delay = 100; |
| 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; |
| } |
| } |
| |
| usleep(delay * 1000); |
| return 0; |
| } |
| |
| int main(int argc, char **argv) |
| { |
| struct sockaddr_storage ss; |
| struct err_msg err; |
| const char *arg0; |
| int arg; |
| int sock; |
| |
| arg0 = argv[0]; |
| if (argc < 2) |
| usage(1, arg0); |
| |
| if (addr_to_ss(argv[1], &ss, &err) < 0) |
| die(1, "%s\n", err.msg); |
| |
| sock = -1; |
| for (arg = 2; arg < argc; arg++) { |
| switch (argv[arg][0]) { |
| case 'L': |
| /* silently ignore existing connections */ |
| if (sock == -1) |
| sock = tcp_listen(&ss, argv[arg]); |
| if (sock < 0) |
| die(1, "Fatal: tcp_listen() failed.\n"); |
| break; |
| |
| case 'C': |
| /* silently ignore existing connections */ |
| if (sock == -1) |
| sock = tcp_connect(&ss, argv[arg]); |
| if (sock < 0) |
| die(1, "Fatal: tcp_connect() failed.\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"); |
| 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"); |
| if (tcp_recv(sock, argv[arg]) < 0) |
| die(1, "Fatal: tcp_recv() failed.\n"); |
| break; |
| |
| case 'S': |
| if (sock < 0) |
| die(1, "Fatal: tcp_send() on non-socket.\n"); |
| if (tcp_send(sock, argv[arg]) < 0) |
| 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"); |
| 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"); |
| 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"); |
| break; |
| |
| case 'K': |
| if (sock < 0 || close(sock) < 0) |
| die(1, "Fatal: close() on non-socket.\n"); |
| sock = -1; |
| break; |
| |
| case 'F': |
| /* ignore errors on shutdown() as they are common */ |
| if (sock >= 0) |
| shutdown(sock, SHUT_WR); |
| break; |
| |
| default: |
| usage(1, arg0); |
| } |
| } |
| return 0; |
| } |