/*** connection forwarder and logger
 * (c) 2005, Luciano Rocha <strange@nsk.pt>
 *
 *      compilation:
 *                   cc -o plug plug.c
 *      if it doesn't work, try:
 *                   cc -o plug plug.c -lnsl -lsocket
 *
 */

#include <stdio.h>
#include <assert.h>
#include <errno.h>
#include <unistd.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/un.h>
#include <netinet/in.h>
#include <sys/time.h>
#include <signal.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <ctype.h>
#include <fcntl.h>

static int gethost(const char *h, struct in_addr *addr);
static unsigned short getport(const char *p);
static int parse_tcp(char *s, struct sockaddr_in *sa);
static int parse_unix(char *s, struct sockaddr_un *su);
static int parse_server(char *);
static int parse_client(char *);
static int accpt_tcp(void);
static int accpt_unix(void);
static void loop(void);
static void do_write1(int);
static void do_write2(int);
static void do_read1(int);
static void do_read2(int);
static void check_cid(int);
static void openlog(int);

#define MAX_CON 100
#define BUF_SIZ (4<<10)

static struct {
    int cid, fdlog;
    int fd1, fd2;
    int len1, len2;
    char buf1[BUF_SIZ], buf2[BUF_SIZ];
    struct {
        int c1_r:1;
        int c1_w:1;
        int c2_r:1;
        int c2_w:1;
    };
} conns[MAX_CON];

#define C(a, b) conns[a].b

static struct {
    union {
        struct sockaddr_in sa;
        struct sockaddr_un su;
    } sa;
    int salen;
} connd;

static int dfd, null;
static int dcid;
static int cid_next;
static int ssock;
static int (*accpt)(void);
static char *dfn;

int main(int ac, char *av[]) {

    dfd = 0;
    if (ac == 4) {
        char *p;
        
        if ((p = strrchr(av[3], '%'))) {
            int i;

            for (i = 1; isdigit(p[i]); ++i);
            if (p[i] != 'd') {
                fprintf(stderr, "file '%s' has invalid %% directive,"
                        " assuming static file name\n", av[3]);
            } else if (strchr(av[3], '%') != p) {
                fprintf(stderr, "file pattern '%s' must define only one"
                        " %% directive\n", av[3]);
                return 1;
            } else
                dfn = av[3];
        }
        if (!dfn && (dfd = open(av[3], O_WRONLY|O_CREAT|O_APPEND, 0666)) < 0) {
            perror(av[3]);
            return 1;
        }
        ac = 3;
    }
    if (ac != 3) {
        fprintf(stderr, "Usage: plug <from> <to> [<file>]\n"
                "\n"
                "from/to: [tcp://][host:]port\n"
                "          unix://file\n"
                "          unix://@anon\n"
                "\n"
                "Transfer logs will be stored in file, if specified.\n"
                "If file has a %%[0-9]*d pattern, one file will be created"
                " per connection.\n"
                "Example of a file pattern: log.%%03d\n"
                "\n");
        return 1;
    }

    if ((null = open("/dev/null", O_WRONLY|O_APPEND)) < 0) {
        perror("/dev/null");
        return 1;
    }

	signal(SIGPIPE, SIG_IGN);

    if (parse_server(av[1]) && parse_client(av[2]))
        loop();
    return 1;
}

static void loop(void) {
    fd_set inset, outset;
    int smax, i, r, cfr, t;
    struct sockaddr *sa = (struct sockaddr *) &connd.sa;
    

	smax = ssock + 1;

    for (;;) {
		FD_ZERO(&inset);
		FD_ZERO(&outset);

        cfr = -1;

        for (i = 0; i < MAX_CON; ++i) {
            if (C(i,cid) > 0) {
                if (!C(i,c1_r) && !C(i,c2_w) && C(i,len1) < BUF_SIZ) {
                        FD_SET(C(i,fd1), &inset);
                }
                if (!C(i,c2_r) && !C(i,c1_w) && C(i,len2) < BUF_SIZ) {
                        FD_SET(C(i,fd2), &inset);
                }
                if (!C(i,c1_w) && C(i,len2) > 0) {
                    FD_SET(C(i,fd1), &outset);
                }
                if (!C(i,c2_w) && C(i,len1) > 0) {
                    FD_SET(C(i,fd2), &outset);
                }
            } else if (C(i,cid) < 0) {
                FD_SET(C(i,fd2), &outset);
            } else if (cfr < 0) cfr = i;
        }

		if (cfr >= 0) FD_SET(ssock, &inset);

		r = select(smax, &inset, &outset, NULL, NULL);
        if (r == 0)
            continue;
        if (r < 0) {
			if (errno == EINTR || errno == EAGAIN) continue;
			perror("select");
			break;
		}

        for (i = 0; i < MAX_CON; ++i) {
            if (C(i,cid) > 0) {
                if (!C(i,c1_w) && FD_ISSET(C(i,fd1), &outset))
                        do_write1(i);
                if (!C(i,c2_w) && FD_ISSET(C(i,fd2), &outset))
                        do_write2(i);
                if (!C(i,c1_r) && FD_ISSET(C(i,fd1), &inset))
                        do_read1(i);
                if (!C(i,c2_r) && FD_ISSET(C(i,fd2), &inset))
                        do_read2(i);
            }
            if (C(i,cid) < 0 && FD_ISSET(C(i,fd2), &outset)) {
                socklen_t t = sizeof r;
                if (getsockopt(C(i,fd2), SOL_SOCKET, SO_ERROR, &r, &t) < 0) {
                    perror("getsockopt(2)");
                    r = 1;
                } else if (r != 0) {
                    fprintf(stderr, "connect(2): %s\n",
                            strerror(r));
                }
                if (r) {
                    close(C(i,fd1));
                    close(C(i,fd2));
                    C(i,cid) = 0;
                    continue;
                } else C(i,cid) = ++cid_next;
            }
        }

        if (FD_ISSET(ssock, &inset) && (r = accpt()) > 0) {
            memset(&conns[cfr], '\0', sizeof conns[0]);
            if (fcntl(r, F_SETFL, O_NONBLOCK) < 0) {
                perror("client: fcntl(O_NONBLOCK)");
                close(r);
                continue;
            }
            if ((t = socket(sa->sa_family, SOCK_STREAM, 0)) < 0) {
                perror("client(2): socket");
                close(r);
                continue;
            }
            if (fcntl(t, F_SETFL, O_NONBLOCK) < 0) {
                perror("client(2): fcntl(O_NONBLOCK)");
                close(r);
                close(t);
                continue;
            }
            if (connect(t, sa, connd.salen) < 0) {
                if (errno != EINPROGRESS) {
                    perror("client(2): connect");
                    close(r);
                    close(t);
                    continue;
                }
                C(cfr,cid) = -1;
            } else C(cfr,cid) = ++cid_next;
            C(cfr,fd1) = r;
            C(cfr,fd2) = t;
            if (r >= smax) smax = r + 1;
            if (t >= smax) smax = t + 1;
        }
	}
}

static void do_read2(int i) {
    int r;

    r = recv(C(i,fd2), C(i,buf2) + C(i,len2),
            BUF_SIZ - C(i,len2), 0);
    if (r < 0 && (errno == EAGAIN || errno == EINTR)) return;
    if (r <= 0) {
        if (r < 0)
            fprintf(stderr, "recv: %d/2:%s", C(i,cid), strerror(errno));
        if (C(i,c2_w)) {
            close(C(i,fd2));
            C(i,fd2) = -1;
        } else {
            shutdown(C(i,fd2), SHUT_RD);
        }
        C(i,c2_r) = 1;
        // fd2 > 0, else wouldn't read
        if (C(i,len2) <= 0) {
            if (C(i,c1_r)) {
                close(C(i,fd1));
                C(i,fd1) = -1;
            } else {
                shutdown(C(i,fd1), SHUT_WR);
            }
            C(i,c1_w) = 1;
        }
        if (C(i,fd1) == -1 && C(i,fd2) == -1) {
            C(i,cid) = 0;
            if (C(i,fdlog) > 0)
                close(C(i,fdlog));
        }
    } else {
        C(i,len2) += r;
    }
}

static void do_read1(int i) {
    int r;

    r = recv(C(i,fd1), C(i,buf1) + C(i,len1),
            BUF_SIZ - C(i,len1), 0);
    if (r < 0 && (errno == EAGAIN || errno == EINTR)) return;
    if (r <= 0) {
        if (r < 0)
            fprintf(stderr, "recv: %d/1:%s", C(i,cid), strerror(errno));
        if (C(i,c1_w)) {
            close(C(i,fd1));
            C(i,fd1) = -1;
        } else {
            shutdown(C(i,fd1), SHUT_RD);
        }
        C(i,c1_r) = 1;
        // fd2 > 0, else wouldn't read
        if (C(i,len1) <= 0) {
            if (C(i,c2_r)) {
                close(C(i,fd2));
                C(i,fd2) = -1;
            } else {
                shutdown(C(i,fd2), SHUT_WR);
            }
            C(i,c2_w) = 1;
        }
        if (C(i,fd1) == -1 && C(i,fd2) == -1) {
            C(i,cid) = 0;
            if (C(i,fdlog) > 0)
                close(C(i,fdlog));
        }
    } else {
        if (dfd > 0) {
            check_cid(C(i,cid));
            write(dfd, C(i,buf1) + C(i,len1), r);
        } else if (dfn) {
            if (C(i,fdlog) <= 0) openlog(i);
            write(C(i,fdlog),  C(i,buf1) + C(i,len1), r);
        }
        C(i,len1) += r;
    }
}

static void do_write1(int i) {
    int r;

    r = send(C(i,fd1), C(i,buf2), C(i,len2), 0);
    if (r < 0 && (errno == EAGAIN || errno == EINTR)) return;
    if (r <= 0) {
        if (r < 0)
            fprintf(stderr, "send: %d/1:%s", C(i,cid), strerror(errno));
        if (C(i,c1_r)) {
            close(C(i,fd1));
            C(i,fd1) = -1;
        } else {
            shutdown(C(i,fd1), SHUT_WR);
        }
        C(i,c1_w) = 1;
        if (C(i,c2_w)) {
            close(C(i,fd2));
            C(i,fd2) = -1;
        } else {
            shutdown(C(i,fd2), SHUT_RD);
        }
        C(i,c2_r) = 1;
        if (C(i,fd1) == -1 && C(i,fd2) == -1) {
            C(i,cid) = 0;
            if (C(i,fdlog) > 0)
                close(C(i,fdlog));
        }
    } else {
        if (dfd > 0) {
            check_cid(C(i,cid));
            write(dfd, C(i,buf2), r);
        } else if (dfn) {
            if (C(i,fdlog) <= 0) openlog(i);
            write(C(i,fdlog), C(i,buf2), r);
        }
        memmove(C(i,buf2), C(i,buf2)+r, C(i,len2) - r);
        C(i,len2) -= r;
    }
}

static void do_write2(int i) {
    int r;

    r = send(C(i,fd2), C(i,buf1), C(i,len1), 0);
    if (r < 0 && (errno == EAGAIN || errno == EINTR)) return;
    if (r <= 0) {
        if (r < 0)
            fprintf(stderr, "send: %d/2:%s", C(i,cid), strerror(errno));
        if (C(i,c2_r)) {
            close(C(i,fd2));
            C(i,fd2) = -1;
        } else {
            shutdown(C(i,fd2), SHUT_WR);
        }
        C(i,c2_w) = 1;
        if (C(i,c1_w)) {
            close(C(i,fd1));
            C(i,fd1) = -1;
        } else {
            shutdown(C(i,fd1), SHUT_RD);
        }
        C(i,c1_r) = 1;
        if (C(i,fd1) == -1 && C(i,fd2) == -1) {
            C(i,cid) = 0;
            if (C(i,fdlog) > 0)
                close(C(i,fdlog));
        }
    } else {
        memmove(C(i,buf1), C(i,buf1)+r, C(i,len1) - r);
        C(i,len1) -= r;
    }
}

static unsigned short getport(const char *p) {
	struct servent *serv;
	int i;
		
	if (!isdigit(*p)) {
		if (!(serv = getservbyname(p, 0))) {
            fprintf(stderr, "unrecognized service %s\n", p);
			return 0;
        }
		return serv->s_port;
	}
	i = strtol(p, NULL, 0);
	if (i < 1 || i > 65535) {
        fprintf(stderr, "invalid port: %d\n", i);
		return 0;
    }
	return htons(i);
}

static int gethost(const char *h, struct in_addr *addr) {
    struct hostent *host;
    char b[40];

    if (inet_pton(AF_INET, h, addr) == 1)
        return 1;
    if (!(host = gethostbyname(h))) {
        herror(h);
        return 0;
    }
    if (host->h_addrtype != AF_INET) {
        fprintf(stderr, "%s doesn't resolve to an IPv4 address\n", h);
        return 0;
    }
    memcpy(addr, host->h_addr_list[0], host->h_length);
    printf("using %s for %s",
                inet_ntop(AF_INET, host->h_addr_list[0], b, sizeof b),
                h);
    if (host->h_addr_list[1]) {
        int i;

        printf(" (remaining ");
        for (i = 1; host->h_addr_list[i]; ++i) {
            printf("%s%s", i == 1 ? "" : ",",
                    inet_ntop(AF_INET, host->h_addr_list[i], b, sizeof b));
        }
        printf(")");
    }
    printf("\n");
    return 1;
}

static int parse_tcp(char *s, struct sockaddr_in *sa) {
    char *p;

    memset(sa, '\0', sizeof *sa);
    sa->sin_family = AF_INET;
    sa->sin_addr.s_addr = INADDR_ANY;
    if ((p = strrchr(s, ':'))) {
        *p = '\0';
        if (!gethost(s, &(sa->sin_addr)))
            return 0;
        *p = ':';
        s = p+1;
    }
    if (!(sa->sin_port = getport(s)))
        return 0;
    return 1;
}

static int parse_unix(char *s, struct sockaddr_un *su) {

    memset(su, '\0', sizeof *su);
    su->sun_family = AF_UNIX;
    strncpy(su->sun_path, s, sizeof(su->sun_path) - 1);
    if (su->sun_path[0] == '@')
        su->sun_path[0] = '\0';
    return 1;
}

static int parse_server(char *s) {
    union {
        struct sockaddr_in sa;
        struct sockaddr_un su;
    } _sa;
    struct sockaddr *sa = (struct sockaddr *) &_sa;
    int slen, r;

    if (!strstr(s, "://")) {
        r = parse_tcp(s, &_sa.sa);
        accpt = accpt_tcp;
        slen = sizeof(struct sockaddr_in);
    } else if (!strncmp(s, "tcp://", 6)) {
        r = parse_tcp(s+6, &_sa.sa);
        accpt = accpt_tcp;
        slen = sizeof(struct sockaddr_in);
    } else if (!strncmp(s, "unix://", 7)) {
        r = parse_unix(s+6, &_sa.su);
        accpt = accpt_unix;
        slen = sizeof(struct sockaddr_un);
    } else {
        fprintf(stderr, "unknown protocol in %s\n", s);
        r = 0;
    }
    if (!r) return 0;
    if ((ssock = socket(sa->sa_family, SOCK_STREAM, 0)) < 0) {
        perror("socket");
        return 0;
    }
	r = 1;
	if (setsockopt(ssock, SOL_SOCKET, SO_REUSEADDR, &r, sizeof r)) {
		perror("setsockopt(SO_REUSEADDR)");
        return 0;
    }
    if (fcntl(ssock, F_SETFL, O_NONBLOCK) < 0) {
        perror("fcntl(O_NONBLOCK)");
        return 0;
    }
	if (bind(ssock, sa, slen)) {
		perror("bind");
		return 0;
	}
	if (listen(ssock, 4)) {
		perror("listen");
		return 0;
	}
    return 1;
}

static int parse_client(char *s) {
    int r;

    if (!strstr(s, "://")) {
        r = parse_tcp(s, &connd.sa.sa);
        connd.salen = sizeof(struct sockaddr_in);
    } else if (!strncmp(s, "tcp://", 6)) {
        r = parse_tcp(s+6, &connd.sa.sa);
        connd.salen = sizeof(struct sockaddr_in);
    } else if (!strncmp(s, "unix://", 7)) {
        r = parse_unix(s+6, &connd.sa.su);
        connd.salen = sizeof(struct sockaddr_un);
    } else {
        fprintf(stderr, "unknown protocol in %s\n", s);
        r = 0;
    }
    return r;
}

static int accpt_tcp(void) {
    struct sockaddr_in sa;
    socklen_t salen;
    int fd;

    salen = sizeof sa;
    fd = accept(ssock, (struct sockaddr *) &sa, &salen);
    if (fd >= 0) {
        char b[40];

        printf("new connection from %s:%d\n",
                inet_ntop(AF_INET, &sa.sin_addr, b, sizeof b),
                ntohs(sa.sin_port));
    }
    return fd;
}

static int accpt_unix(void) {
    struct sockaddr_un su;
    socklen_t sulen;
    int fd;

    sulen = sizeof su;
    fd = accept(ssock, (struct sockaddr *) &su, &sulen);
    if (fd >= 0) {
        if (su.sun_path[0] == '\0')
            su.sun_path[0] = '@';
        printf("new connection from %.*s\n", (int) sizeof su.sun_path,
                su.sun_path);
    }
    return fd;
}

static void check_cid(int cid) {
    static char b[40];

    if (cid != dcid) {
        snprintf(b, sizeof b, "\n------------- %d -------------\n", cid);
        write(dfd, b, strlen(b));
        dcid = cid;
    }
}

static void openlog(int i) {
    static char b[256];
    int fd;

    snprintf(b, sizeof b, dfn, C(i,cid));
    if ((fd = creat(b, 0666)) < 0) {
        perror(b);
        C(i,fdlog) = null;
    } else {
        C(i,fdlog) = fd;
    }
}

// vim: ts=4 shiftwidth=4 expandtab