stdio-proxy.c

/*
 * stdio-proxy - bidirectional proxy between stdio and a TCP/UDP socket.
 *
 * Usage: stdio-proxy [-v] tcp|udp:host:port
 *
 * Designed for use as an SSH ProxyCommand.  Reads from stdin, writes to the
 * remote socket; reads from the remote socket, writes to stdout.  For TCP,
 * sends a half-close (SHUT_WR) when stdin reaches EOF so the server gets a
 * clean FIN without the proxy having to know when the exchange is over.
 */

#define _GNU_SOURCE

#include <arpa/inet.h>
#include <errno.h>
#include <netdb.h>
#include <poll.h>
#include <signal.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

#define BUF_SIZE 65536

static int verbose;
static volatile sig_atomic_t g_interrupted;
static int g_sockfd = -1;

/* -------------------------------------------------------------------------
 * Logging
 * ---------------------------------------------------------------------- */

static void vlog(const char *fmt, ...)
{
	if (!verbose)
		return;
	va_list ap;
	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);
	fputc('\n', stderr);
}

/* -------------------------------------------------------------------------
 * Helpers
 * ---------------------------------------------------------------------- */

/* Write exactly len bytes, retrying on short writes. */
static int write_all(int fd, const char *buf, size_t len)
{
	while (len > 0) {
		ssize_t w = write(fd, buf, len);
		if (w <= 0)
			return -1;
		buf += w;
		len -= (size_t)w;
	}
	return 0;
}

/* -------------------------------------------------------------------------
 * Signal handling
 * ---------------------------------------------------------------------- */

static void sig_handler(int sig)
{
	(void)sig;
	g_interrupted = 1;
	/* Shut down the socket so the poll() loop unblocks. */
	if (g_sockfd >= 0)
		shutdown(g_sockfd, SHUT_RDWR);
}

/* -------------------------------------------------------------------------
 * main
 * ---------------------------------------------------------------------- */

int main(int argc, char *argv[])
{
	/* --- Option parsing ------------------------------------------------- */
	int opt;
	while ((opt = getopt(argc, argv, "v")) != -1) {
		switch (opt) {
		case 'v':
			verbose = 1;
			break;
		default:
			goto usage;
		}
	}

	if (argc - optind != 1)
		goto usage;

	/* --- Argument parsing: protocol:host:port --------------------------- */
	char arg[512];
	snprintf(arg, sizeof(arg), "%s", argv[optind]);

	/* First colon separates protocol from host:port. */
	char *p1 = strchr(arg, ':');
	if (!p1) {
		fprintf(stderr,
			"Invalid format. Expected [protocol]:[host]:[port]\n");
		return 1;
	}
	*p1 = '\0';
	char *proto = arg;
	char *hostport = p1 + 1;

	/* Last colon separates host from port (safe for bare IPv6 addresses). */
	char *p2 = strrchr(hostport, ':');
	if (!p2) {
		fprintf(stderr,
			"Invalid format. Expected [protocol]:[host]:[port]\n");
		return 1;
	}
	*p2 = '\0';
	char *host = hostport;
	char *port = p2 + 1;

	/* --- Protocol selection --------------------------------------------- */
	int is_tcp, sock_type;
	if (strcasecmp(proto, "tcp") == 0) {
		is_tcp = 1;
		sock_type = SOCK_STREAM;
	} else if (strcasecmp(proto, "udp") == 0) {
		is_tcp = 0;
		sock_type = SOCK_DGRAM;
	} else {
		fprintf(stderr,
			"Unsupported protocol: %s. Use 'tcp' or 'udp'.\n",
			proto);
		return 1;
	}

	/* --- Name resolution ------------------------------------------------ */
	struct addrinfo hints = {
		.ai_family = AF_UNSPEC,
		.ai_socktype = sock_type,
	};
	struct addrinfo *res;
	int rc = getaddrinfo(host, port, &hints, &res);
	if (rc != 0) {
		fprintf(stderr, "Error resolving %s:%s: %s\n", host, port,
			gai_strerror(rc));
		return 1;
	}

	/* --- Connect -------------------------------------------------------- */
	vlog("Connecting to %s %s:%s...", proto, host, port);

	int fd = -1;
	for (struct addrinfo * r = res; r; r = r->ai_next) {
		fd = socket(r->ai_family, r->ai_socktype, r->ai_protocol);
		if (fd < 0)
			continue;
		if (connect(fd, r->ai_addr, r->ai_addrlen) == 0)
			break;
		close(fd);
		fd = -1;
	}
	freeaddrinfo(res);

	if (fd < 0) {
		fprintf(stderr, "Error connecting to %s://%s:%s: %s\n",
			proto, host, port, strerror(errno));
		return 1;
	}

	g_sockfd = fd;
	vlog("Connected to %s://%s:%s", proto, host, port);

	/* --- Signal handlers ------------------------------------------------ */
	struct sigaction sa = {.sa_handler = sig_handler };
	sigemptyset(&sa.sa_mask);
	sigaction(SIGINT, &sa, NULL);
	sigaction(SIGTERM, &sa, NULL);

	/* --- Bidirectional proxy loop (poll-based, no threads) -------------- */
	char buf[BUF_SIZE];
	int stdin_eof = 0;
	int sock_eof = 0;

	struct pollfd fds[2];
	fds[0].fd = STDIN_FILENO;
	fds[1].fd = fd;

	while (!g_interrupted && (!stdin_eof || !sock_eof)) {
		fds[0].events = stdin_eof ? 0 : POLLIN;
		fds[1].events = sock_eof ? 0 : POLLIN;

		int ready = poll(fds, 2, -1);
		if (ready < 0) {
			if (errno == EINTR)
				continue;
			break;
		}

		/* stdin → socket */
		if (!stdin_eof
		    && (fds[0].revents & (POLLIN | POLLHUP | POLLERR))) {
			ssize_t n = read(STDIN_FILENO, buf, sizeof(buf));
			if (n <= 0) {
				if (n < 0 && errno != EINTR)
					vlog("stdin read error: %s",
					     strerror(errno));
				stdin_eof = 1;
				/* TCP half-close: tell the server we're done sending. */
				if (is_tcp)
					shutdown(fd, SHUT_WR);
				vlog("stdin EOF; half-closed write side");
			} else {
				if (write_all(fd, buf, (size_t)n) < 0) {
					vlog("socket write error: %s",
					     strerror(errno));
					break;
				}
			}
		}

		/* socket → stdout */
		if (!sock_eof
		    && (fds[1].revents & (POLLIN | POLLHUP | POLLERR))) {
			ssize_t n = read(fd, buf, sizeof(buf));
			if (n <= 0) {
				if (n < 0 && errno != EINTR)
					vlog("socket read error: %s",
					     strerror(errno));
				sock_eof = 1;
				vlog("remote closed connection");
			} else {
				if (write_all(STDOUT_FILENO, buf, (size_t)n) <
				    0) {
					vlog("stdout write error: %s",
					     strerror(errno));
					break;
				}
			}
		}

		/* Once the server closes its side we're done. */
		if (sock_eof)
			break;
	}

	vlog("Shutting down connection to %s://%s:%s", proto, host, port);
	close(fd);
	return 0;

 usage:
	fprintf(stderr, "Usage:    %s [-v] [protocol]:[host]:[port]\n",
		argv[0]);
	fprintf(stderr, "Examples: %s tcp:example.com:80\n", argv[0]);
	fprintf(stderr, "          %s udp:dns.server:53\n", argv[0]);
	return 1;
}