#include <unistd.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <stdio.h>
#include <sys/un.h>
#include <string.h>

unsigned char ibuf[65536];
unsigned char obuf[65536];

int main(int argc, unsigned char **argv)
{
    int fd, proto, type;
    struct sockaddr *name, *peer;
    int namelen;
    unsigned char *p;
    int ret, ilen, olen;
    fd_set rfds, wfds;
    int io;
    
    if(argc < 4)
    {
	fprintf(stderr, "usage: redirsock [io] proto type peer\n");
	exit(1);
    }
    io = 3;
    if(!strcmp(argv[1], "r"))
    {
	io = 1;
	argv++;
	argc--;
    }
    if(!strcmp(argv[1], "w"))
    {
	io = 2;
	argv++;
	argc--;
    }
    if(!strcmp(argv[1], "rw"))
    {
	io = 3;
	argv++;
	argc--;
    }
    proto = -1;
    if(!strcmp(argv[1], "un"))
    {
	proto = PF_UNIX;
	namelen = sizeof(struct sockaddr_un);
	peer = (struct sockaddr *)malloc(sizeof(namelen));
	((struct sockaddr_un *)peer)->sun_family = AF_UNIX;
	strcpy(((struct sockaddr_un *)peer)->sun_path, argv[3]);
    }
    if(!strcmp(argv[1], "in"))
    {
	if((p = strchr(argv[3], ':')) == NULL)
	{
	    fprintf(stderr, "no port in address\n");
	    return(1);
	}
	*(p++) = 0;
	proto = PF_INET;
	namelen = sizeof(struct sockaddr_in);
	peer = (struct sockaddr *)malloc(sizeof(namelen));
	((struct sockaddr_in *)peer)->sin_family = AF_INET;
	((struct sockaddr_in *)peer)->sin_port = htons(atoi(p));
	if(!inet_aton(argv[3], &((struct sockaddr_in *)peer)->sin_addr))
	{
	    fprintf(stderr, "invalid IP address\n");
	    return(1);
	}
    }
    if(proto == -1)
    {
	fprintf(stderr, "invalid protocol peer\n");
	return(1);
    }
    switch(*(argv[2]))
    {
    case 'd':
	type = SOCK_DGRAM;
	break;
    case 's':
	type = SOCK_STREAM;
	break;
    default:
	fprintf(stderr, "invalid socket type\n");
	return(1);
    }
    if((fd = socket(proto, type, 0)) < 0)
    {
	perror("socket");
	return(1);
    }
    if(connect(fd, peer, namelen) < 0)
    {
	perror("connect");
	return(1);
    }
    ilen = olen = 0;
    while(1)
    {
	FD_ZERO(&rfds);
	FD_ZERO(&wfds);
	if(io & 1)
	{
	    if(olen < 65536)
		FD_SET(0, &rfds);
	    if(olen > 0)
		FD_SET(fd, &wfds);
	}
	if(io & 2)
	{
	    if(ilen < 65536)
		FD_SET(fd, &rfds);
	    if(ilen > 0)
		FD_SET(1, &wfds);
	}
	if((ret = select(fd + 1, &rfds, &wfds, NULL, NULL)) < 0)
	{
	    perror("select");
	    return(1);
	} else if(ret > 0) {
	    if(FD_ISSET(0, &rfds))
	    {
		if((ret = read(0, obuf, 65536 - olen)) < 0)
		{
		    perror("read(0)");
		    return(1);
		} else if(ret == 0) {
		    break;
		} else {
		    olen += ret;
		}
	    }
	    if(FD_ISSET(fd, &rfds))
	    {
		if((ret = read(fd, ibuf, 65536 - ilen)) < 0)
		{
		    perror("read(socket)");
		    return(1);
		} else if(ret == 0) {
		    break;
		} else {
		    ilen += ret;
		}
	    }
	    if(FD_ISSET(1, &wfds))
	    {
		if((ret = write(1, ibuf, ilen)) < 0)
		{
		    perror("write(1)");
		    return(1);
		} else {
		    memmove(ibuf, ibuf + ret, ilen -= ret);
		}
	    }
	    if(FD_ISSET(fd, &wfds))
	    {
		if((ret = write(fd, obuf, olen)) < 0)
		{
		    perror("write(socket)");
		    return(1);
		} else {
		    memmove(obuf, obuf + ret, olen -= ret);
		}
	    }
	}
    }
    return(0);
}

