/*
 * Trampoline a TCP connection.  This is useful for mounting something
 * that serves 9p.
 *
 * Dan Cross <net!gajendra!cross>
 */

#include <sys/types.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <netinet/in.h>
#include <netinet/tcp.h>

#include <netdb.h>
#include <signal.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

void
nonagle(int sd)
{
	int one;

	one = 1;
	setsockopt(sd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));
}

int
announceon(int port)
{
	struct sockaddr_in srv;
	int sd;

	sd = socket(AF_INET, SOCK_STREAM, 0);
	if (sd < 0) {
		perror("socket");
		exit(EXIT_FAILURE);
	}
	memset(&srv, 0, sizeof(srv));
	srv.sin_family = AF_INET;
	srv.sin_port = htons(port);
	srv.sin_addr.s_addr = htonl(INADDR_ANY);
	if (bind(sd, (struct sockaddr *)&srv, sizeof(srv)) < 0) {
		perror("bind");
		exit(EXIT_FAILURE);
	}
	if (listen(sd, 5) < 0) {
		perror("listen");
		exit(EXIT_FAILURE);
	}
	return(sd);
}

int
listento(int sd)
{
	struct sockaddr_in cli;
	socklen_t csize;

	csize = sizeof(cli);
	memset(&cli, 0, csize);
	sd = accept(sd, (struct sockaddr *)&cli, &csize);
	if (sd < 0)
		return(sd);
	nonagle(sd);
	return(sd);
}

int
dialto(char *host, int port)
{
	struct hostent *h;
	struct sockaddr_in srv;
	int sd;

	h = gethostbyname(host);
	if (h == NULL) {
		return(-1);
	}
	sd = socket(AF_INET, SOCK_STREAM, 0);
	if (sd < 0) {
		return(-1);
	}
	memset(&srv, 0, sizeof(srv));
	srv.sin_family = AF_INET;
	srv.sin_port = htons(port);
	memmove(&srv.sin_addr.s_addr, h->h_addr, h->h_length);
	if (connect(sd, (struct sockaddr *)&srv, sizeof(srv)) < 0) {
		close(sd);
		return(-1);
	}
	nonagle(sd);
	return(sd);
}

void
xfer(int from, int to)
{
	int n;
	char buf[8192];

	for ( ; ; ) {
		n = read(from, buf, sizeof(buf));
		if (n <= 0)
			return;
		if (write(to, buf, n) != n)
			return;
	}
}

void
srv(int myport, char *host, int port)
{
	int ad, sd, fd;
	pid_t pid;

	ad = announceon(myport);
	for ( ; ; ) {
		sd = listento(ad);
		if (sd < 0) {
			perror("listento");
			continue;
		}
		fd = dialto(host, port);
		if (fd < 0) {
			perror("dialto");
			close(sd);
			continue;
		}
		pid = fork();
		if (pid != 0) {		/*  Parent or error.  */
			close(sd);
			close(fd);
			continue;
		}
		/*  This is the child.  */
		setpgid(0, getpid());
		pid = fork();
		if (pid < 0)
			_exit(EXIT_FAILURE);
		if (pid == 0)
			xfer(fd, sd);
		else
			xfer(sd, fd);
		killpg(getpgid(0), SIGTERM);
	}
}

void
reaper(int sig)
{
	wait(NULL);
}

int
main(int argc, char *argv[])
{

	if (argc != 4) {
		fprintf(stderr, "Usage: trampoline port host port\n");
		exit(EXIT_FAILURE);
	}
	/*  Handlers stay installed with signal() these days.  */
	signal(SIGCHLD, reaper);
	srv(atoi(argv[1]), argv[2], atoi(argv[3]));

	return(EXIT_SUCCESS);
}