tcpcat.c

/* Copyright (c) 2016 Ada Worcester <ada@oss.pikhq.com>
 * 
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
 * SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
 * OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
 * CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */
#define _XOPEN_SOURCE 700
#include <stdarg.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <pthread.h>
#include <locale.h>
#include <errno.h>

bool opt_listen;
char *opt_port;

char *argv0;

locale_t loc;

static const char *usage_msg =
	"Usage: tcpcat [-lph] [host] [port]\n";

static const char *help_msg =
	"Usage: tcpcat [-lph] [host] [port]\n"
	"tcpcat connects to or listens on TCP, and sends data between stdin/stdout and\n"
	"the socket.\n"
	"\n"
	" -l\t\tSet up a listening socket\n"
	" -p port\tSpecify a port to connect to or listen on\n"
	" -h\t\tPrint out this message\n";

void die(char *fmt, ...)
{
	va_list ap;
	fprintf(stderr, "%s: ", argv0);
	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);
	exit(1);
}

int get_socket(char *node, char *port)
{
	int status, fd;

	struct addrinfo hints = {
		.ai_family = AF_UNSPEC,
		.ai_socktype = SOCK_STREAM
	};
	struct addrinfo *res, *p;

	if((status = getaddrinfo(node, port, &hints, &res)) != 0) {
		die("%s\n", gai_strerror(status));
	}

	for(p = res; p; p = p->ai_next) {
		fd = socket(p->ai_family, p->ai_socktype, p->ai_protocol);
		if(fd < 0) continue;
		if(opt_listen) {
			if(bind(fd, p->ai_addr, p->ai_addrlen) < 0) {
				close(fd);
				continue;
			}
			break;
		} else {
			if(connect(fd, p->ai_addr, p->ai_addrlen) < 0) {
				close(fd);
				continue;
			}
			break;
		}
	}
	if(!p) {
		die("%s\n", strerror(errno));
	}
	if(opt_listen) {
		if(listen(fd, 1) < 0)
			die("%s\n", strerror(errno));
	}
	freeaddrinfo(res);

	return fd;
}

void *inout_fn(void *fd_ptr)
{
	ssize_t n, cnt;
	char buf[4096], *p;
	int fd_in, fd_out;

	fd_in = ((int*)fd_ptr)[0];
	fd_out = ((int*)fd_ptr)[1];

	for(;;) {
		n = read(fd_in, buf, sizeof buf);
		if(n < 0) die("%s\n", strerror_l(errno, loc));
		if(n == 0) {
			if(shutdown(fd_out, SHUT_WR) < 0) {
				if(errno != ENOTSOCK)
					die("%s\n", strerror_l(errno, loc));
			}
			return 0;
		}
		p = buf;
		while(n) {
			cnt = write(fd_out, p, n);
			if(cnt < 0) die("%s\n", strerror_l(errno, loc));
			n -= cnt;
			p += cnt;
		}
	}
}

int main(int argc, char **argv)
{
	int c;
	char **v;
	int o;
	int fd, listen_fd;
	int err;
	pthread_t stdin_thread;

	int stdin_eof = 0, server_eof = 0;

	argv0 = strrchr(argv[0], '/');
	if(!argv0) argv0 = argv[0]; else argv0++;

	while((o = getopt(argc, argv, "hlp:")) != -1) {
		switch(o) {
		case 'h':
			printf("%s", help_msg);
			return 0;
		case 'l':
			opt_listen = 1;
			break;
		case 'p':
			opt_port = optarg;
			break;
		default:
			fprintf(stderr, "%s", usage_msg);
			return 1;
		}
	}

	c = argc - optind;
	v = argv + optind;
	if(!opt_listen && c < 2) {
		die("missing operands\n%s", usage_msg);
	}
	if(c > 2) {
		die("too many operands\n%s", usage_msg);
	}

	if(!opt_port && c < 2) {
		die("missing port\n%s", usage_msg);
	}
	
	fd = get_socket(c > 1 ? v[0] : 0, c > 1 ? v[1] : opt_port);

	if(opt_listen) {
		listen_fd = fd;
		fd = accept(listen_fd, 0, 0);
		if(fd < 0)
			die("%s\n", strerror(errno));
	}

	loc = newlocale(LC_ALL_MASK, "", 0);
	if(loc == (locale_t)0) die("%s\n", strerror(errno));

	err = pthread_create(&stdin_thread, 0, inout_fn, (void*)(int[]){0, fd});
	if(err) {
		die("%s\n", strerror(err));
	}
	inout_fn((void*)(int[]){fd, 1});
	err = pthread_join(stdin_thread, 0);
	if(err) {
		die("%s\n", strerror(err));
	}
	return 0;
}