redir/redir.c

   1 /*
   2  * Simple TCP redirector
   3  *
   4  * Luiz Fernando N. Capitulino <lcapitulino@gmail.com>
   5  */
   6 #include <stdio.h>
   7 #include <stdlib.h>
   8 #include <string.h>
   9 #include <unistd.h>
  10 #include <signal.h>
  11 #include <sys/types.h>
  12 #include <sys/socket.h>
  13 #include <sys/select.h>
  14 #include <sys/wait.h>
  15 #include <netinet/in.h>
  16 #include <arpa/inet.h>
  17
  18 static void die(const char *s)
  19 {
  20     perror(s);
  21     exit(1);
  22 }
  23
  24 static void help(void)
  25 {
  26     printf("redir <source IP>:<source port> <destination IP>:<destination port>\n");
  27     exit(1);
  28 }
  29
  30 static void cleanup_child(int signr)
  31 {
  32     wait(NULL);
  33 }
  34
  35 static int read_user_addr(struct sockaddr_in *addr, const char *addr_str)
  36 {
  37     int ret = -1;
  38     char *p, *addr_str2;
  39
  40     addr_str2 = strdup(addr_str);
  41     p = strchr(addr_str2, ':');
  42     if (!p) {
  43         goto out;
  44     }
  45
  46     *p = '\0';
  47     if (!inet_aton(addr_str2, &addr->sin_addr)) {
  48         goto out;
  49     }
  50
  51     p++;
  52     if (p == '\0') {
  53         goto out;
  54     }
  55
  56     addr->sin_port = htons(strtol(p, NULL, 10));
  57     addr->sin_family = AF_INET;
  58     ret = 0;
  59
  60 out:
  61     free(addr_str2);
  62     return ret;
  63 }
  64
  65 static int redir(int sfd, int dfd)
  66 {
  67     unsigned char buf[4096];
  68     size_t count, total, idx;
  69
  70     memset(buf, 0, sizeof(buf));
  71
  72     count = read(sfd, buf, sizeof(buf));
  73     if (!count) {
  74         return 1;
  75     }
  76
  77     idx = 0;
  78     total = count;
  79     for (;;) {
  80         count = write(dfd, buf + idx, total);
  81         if (count < 0) {
  82             die("write()");
  83         }
  84
  85         total -= count;
  86         if (!total) {
  87             return 0;
  88         }
  89         idx += count;
  90     }
  91 }
  92
  93 static void handle_new_client(int cfd, struct sockaddr_in *client_add,
  94                               struct sockaddr_in *dest_addr)
  95 {
  96     int ret, dfd, nfds;
  97
  98     dfd = socket(AF_INET, SOCK_STREAM, 0);
  99     if (dfd < 0) {
 100         die("socket()");
 101     }
 102
 103     ret = connect(dfd, (struct sockaddr *) dest_addr,
 104                   sizeof(struct sockaddr_in));
 105     if (dfd < 0) {
 106         die("connect()");
 107     }
 108
 109     nfds = cfd >= dfd ? cfd : dfd;
 110     for (;;) {
 111         fd_set rfds;
 112
 113         FD_ZERO(&rfds);
 114         FD_SET(cfd, &rfds);
 115         FD_SET(dfd, &rfds);
 116
 117         ret = select(nfds + 1, &rfds, NULL, NULL, 0);
 118         if (ret < 0) {
 119             die("select()");
 120         } else if (FD_ISSET(cfd, &rfds)) {
 121             if (redir(cfd, dfd)) {
 122                 return;
 123             }
 124         } else if (FD_ISSET(dfd, &rfds)) {
 125             if (redir(dfd, cfd)) {
 126                 return;
 127             }
 128         } else {
 129             abort();
 130         }
 131     }
 132
 133     close(dfd);
 134 }
 135
 136 int main(int argc, const char *argv[])
 137 {
 138     struct sockaddr_in src_addr, dest_addr, client_addr;
 139     socklen_t dest_addr_sz;
 140     int err, sfd, cfd;
 141
 142     if (argc != 3) {
 143         help();
 144     }
 145
 146     err = read_user_addr(&src_addr, argv[1]);
 147     if (err) {
 148         help();
 149     }
 150
 151     err = read_user_addr(&dest_addr, argv[2]);
 152     if (err) {
 153         help();
 154     }
 155
 156     signal(SIGCHLD, cleanup_child);
 157
 158     sfd = socket(AF_INET, SOCK_STREAM, 0);
 159     if (sfd < 0) {
 160         die("socket()");
 161     }
 162
 163     err = bind(sfd, (struct sockaddr *) &src_addr, sizeof(struct sockaddr_in));
 164     if (err) {
 165         die("bind()");
 166     }
 167
 168     err = listen(sfd, 2);
 169     if (err) {
 170         die("listen()");
 171     }
 172
 173     for (;;) {
 174         pid_t pid;
 175
 176         dest_addr_sz = sizeof(struct sockaddr_in);
 177         cfd = accept(sfd, (struct sockaddr *) &client_addr, &dest_addr_sz);
 178         if (cfd < 0) {
 179             die("accept()");
 180         }
 181
 182         fprintf(stderr, "-> New connection from: %s:%d\n",
 183                 inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));
 184
 185         pid = fork();
 186         if (pid < 0) {
 187             perror("fork()");
 188         } else if (pid == 0) {
 189             /* child */
 190             handle_new_client(cfd, &client_addr, &dest_addr);
 191             exit(0);
 192         }
 193
 194         /* parent */
 195         close(cfd);
 196     }
 197
 198     return 0;
 199 }