added async connect

git-svn-id: https://svn.spreadspace.org/tcpproxy/trunk@46 e61f0598-a718-4e21-a8f0-0aadfa62ad6b

2 files changed, 91 insertions, 50 deletions

diff --git a/src/clients.c b/src/clients.c
index dec47ae..a1c8840 100644
--- a/src/clients.c
+++ b/src/clients.c
@@ -1,14 +1,14 @@
 /*
  *  tcpproxy
  *
- *  tcpproxy is a simple tcp connection proxy which combines the 
- *  features of rinetd and 6tunnel. tcpproxy supports IPv4 and 
- *  IPv6 and also supports connections from IPv6 to IPv4 
+ *  tcpproxy is a simple tcp connection proxy which combines the
+ *  features of rinetd and 6tunnel. tcpproxy supports IPv4 and
+ *  IPv6 and also supports connections from IPv6 to IPv4
  *  endpoints and vice versa.
- *  
+ *
  *
  *  Copyright (C) 2010-2011 Christian Pointner <equinox@spreadspace.org>
- *                         
+ *
  *  This file is part of tcpproxy.
  *
  *  tcpproxy is free software: you can redistribute it and/or modify
@@ -37,6 +37,7 @@
 #include <sys/select.h>
 #include <netinet/in.h>
 #include <netinet/tcp.h>
+#include <fcntl.h>
 
 #include "clients.h"
 #include "tcp.h"
@@ -46,7 +47,7 @@ void clients_delete_element(void* e)
 {
   if(!e)
     return;
-  
+
   client_t* element = (client_t*)e;
   close(element->fd_[0]);
   close(element->fd_[1]);
@@ -69,6 +70,35 @@ void clients_clear(clients_t* list)
   slist_clear(&(list->list_));
 }
 
+static int handle_connect(client_t* c, int32_t buffer_size_)
+{
+  if(!c || c->state_ != CONNECTING)
+    return -1;
+
+  int error = 0;
+  int len = sizeof(error);
+  if(getsockopt(c->fd_[1], SOL_SOCKET, SO_ERROR, &error, &len)==-1) {
+    log_printf(ERROR, "Error on getsockopt(): %s", strerror(errno));
+    return -1;
+  }
+  if(error) {
+    log_printf(ERROR, "Error on connect(): %s, not adding client %d", strerror(error), c->fd_[0]);
+    return -1;
+  }
+
+  int i;
+  for(i = 0; i < 2; ++i) {
+    c->write_buf_[i].buf_ = malloc(buffer_size_);
+    if(!c->write_buf_[i].buf_) return -2;
+    c->write_buf_[i].length_ = buffer_size_;
+    c->write_buf_offset_[i] = 0;
+  }
+
+  log_printf(INFO, "successfully added client %d", c->fd_[0]);
+  c->state_ = CONNECTED;
+  return 0;
+}
+
 int clients_add(clients_t* list, int fd, const tcp_endpoint_t remote_end, const tcp_endpoint_t source_end)
 {
   if(!list)
@@ -80,9 +110,16 @@ int clients_add(clients_t* list, int fd, const tcp_endpoint_t remote_end, const
     return -2;
   }
 
+  int i;
+  for(i = 0; i < 2; ++i) {
+    element->write_buf_[i].buf_ = NULL;
+    element->write_buf_[i].length_ = 0;
+    element->write_buf_offset_[i] = 0;
+  }
+  element->state_ = CONNECTING;
   element->fd_[0] = fd;
   element->fd_[1] = socket(remote_end.addr_.ss_family, SOCK_STREAM, 0);
-  if(element->fd_[1] < 0) { 
+  if(element->fd_[1] < 0) {
     log_printf(INFO, "Error on socket(): %s, not adding client %d", strerror(errno), element->fd_[0]);
     close(element->fd_[0]);
     free(element);
@@ -90,7 +127,7 @@ int clients_add(clients_t* list, int fd, const tcp_endpoint_t remote_end, const
   }
 
   int on = 1;
-  if(setsockopt(element->fd_[0], IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on)) || 
+  if(setsockopt(element->fd_[0], IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on)) ||
      setsockopt(element->fd_[1], IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on))) {
     log_printf(ERROR, "Error on setsockopt(): %s", strerror(errno));
     close(element->fd_[0]);
@@ -99,8 +136,17 @@ int clients_add(clients_t* list, int fd, const tcp_endpoint_t remote_end, const
     return -1;
   }
 
+  if(fcntl(element->fd_[0], F_SETFL, O_NONBLOCK) ||
+     fcntl(element->fd_[1], F_SETFL, O_NONBLOCK)) {
+    log_printf(ERROR, "Error on fcntl(): %s", strerror(errno));
+    close(element->fd_[0]);
+    close(element->fd_[1]);
+    free(element);
+    return -1;
+  }
+
   if(source_end.addr_.ss_family != AF_UNSPEC) {
-    if(bind(element->fd_[1], (struct sockaddr *)&(source_end.addr_), source_end.len_)==-1) { 
+    if(bind(element->fd_[1], (struct sockaddr *)&(source_end.addr_), source_end.len_)==-1) {
       log_printf(INFO, "Error on bind(): %s, not adding client %d", strerror(errno), element->fd_[0]);
       close(element->fd_[0]);
       close(element->fd_[1]);
@@ -109,45 +155,29 @@ int clients_add(clients_t* list, int fd, const tcp_endpoint_t remote_end, const
     }
   }
 
-  if(connect(element->fd_[1], (struct sockaddr *)&(remote_end.addr_), remote_end.len_)==-1) { 
-    log_printf(INFO, "Error on connect(): %s, not adding client %d", strerror(errno), element->fd_[0]);
-    close(element->fd_[0]);
-    close(element->fd_[1]);
-    free(element);
-    return -1;
-  }
-
-  element->write_buf_[0].buf_ = malloc(list->buffer_size_);
-  if(!element->write_buf_[0].buf_) {
+  if(slist_add(&(list->list_), element) == NULL) {
     close(element->fd_[0]);
     close(element->fd_[1]);
     free(element);
     return -2;
   }
-  element->write_buf_[0].length_ = list->buffer_size_;
-  element->write_buf_offset_[0] = 0;
 
-  element->write_buf_[1].buf_ = malloc(list->buffer_size_);
-  if(!element->write_buf_[1].buf_) {
-    free(element->write_buf_[0].buf_);
-    close(element->fd_[0]);
-    close(element->fd_[1]);
-    free(element);
-    return -2;
-  }
-  element->write_buf_[1].length_ = list->buffer_size_;
-  element->write_buf_offset_[1] = 0;
+  if(connect(element->fd_[1], (struct sockaddr *)&(remote_end.addr_), remote_end.len_)==-1) {
+    if(errno == EINPROGRESS)
+      return 0;
 
-  if(slist_add(&(list->list_), element) == NULL) {
-    free(element->write_buf_[0].buf_);
-    free(element->write_buf_[1].buf_);
-    close(element->fd_[0]);
-    close(element->fd_[1]);
-    free(element);
-    return -2;
+    log_printf(INFO, "Error on connect(): %s, not adding client %d", strerror(errno), element->fd_[0]);
+    slist_remove(&(list->list_), element);
+    return -1;
   }
 
-  return 0;
+  log_printf(DEBUG, "connect() for client %d returned immediatly", element->fd_[0]);
+
+  int ret = handle_connect(element, list->buffer_size_);
+  if(ret)
+    slist_remove(&(list->list_), element);
+
+  return ret;
 }
 
 void clients_remove(clients_t* list, int fd)
@@ -175,12 +205,12 @@ void clients_print(clients_t* list)
 {
   if(!list)
     return;
-  
+
   slist_element_t* tmp = list->list_.first_;
   while(tmp) {
     client_t* c = (client_t*)tmp->data_;
     if(c) {
-          // print useful info
+          // TODO: print useful info
       printf("client %d <-> %d: tba...\n", c->fd_[0], c->fd_[1]);
     }
     tmp = tmp->next_;
@@ -195,7 +225,7 @@ void clients_read_fds(clients_t* list, fd_set* set, int* max_fd)
   slist_element_t* tmp = list->list_.first_;
   while(tmp) {
     client_t* c = (client_t*)tmp->data_;
-    if(c) {
+    if(c && c->state_ == CONNECTED) {
       if(c->write_buf_offset_[1] < c->write_buf_[1].length_) {
         FD_SET(c->fd_[0], set);
         *max_fd = *max_fd > c->fd_[0] ? *max_fd : c->fd_[0];
@@ -217,7 +247,7 @@ void clients_write_fds(clients_t* list, fd_set* set, int* max_fd)
   slist_element_t* tmp = list->list_.first_;
   while(tmp) {
     client_t* c = (client_t*)tmp->data_;
-    if(c) {
+    if(c && c->state_ == CONNECTED) {
       if(c->write_buf_offset_[0]) {
         FD_SET(c->fd_[0], set);
         *max_fd = *max_fd > c->fd_[0] ? *max_fd : c->fd_[0];
@@ -226,6 +256,9 @@ void clients_write_fds(clients_t* list, fd_set* set, int* max_fd)
         FD_SET(c->fd_[1], set);
         *max_fd = *max_fd > c->fd_[1] ? *max_fd : c->fd_[1];
       }
+    } else if(c && c->state_ == CONNECTING) {
+      FD_SET(c->fd_[1], set);
+      *max_fd = *max_fd > c->fd_[1] ? *max_fd : c->fd_[1];
     }
     tmp = tmp->next_;
   }
@@ -235,12 +268,12 @@ int clients_read(clients_t* list, fd_set* set)
 {
   if(!list)
     return -1;
-  
+
   slist_element_t* tmp = list->list_.first_;
   while(tmp) {
     client_t* c = (client_t*)tmp->data_;
     tmp = tmp->next_;
-    if(c) {
+    if(c && c->state_ == CONNECTED) {
       int i;
       for(i=0; i<2; ++i) {
         int in, out;
@@ -249,7 +282,7 @@ int clients_read(clients_t* list, fd_set* set)
           out = i ^ 1;
         }
         else continue;
-        
+
         int len = recv(c->fd_[in], &(c->write_buf_[out].buf_[c->write_buf_offset_[out]]),  c->write_buf_[out].length_ - c->write_buf_offset_[out], 0);
         if(len < 0) {
           log_printf(INFO, "Error on recv(): %s, removing client %d", strerror(errno), c->fd_[0]);
@@ -260,13 +293,13 @@ int clients_read(clients_t* list, fd_set* set)
           log_printf(INFO, "client %d closed connection, removing it", c->fd_[0]);
           slist_remove(&(list->list_), c);
           break;
-        }       
+        }
         else
           c->write_buf_offset_[out] += len;
       }
     }
   }
-  
+
   return 0;
 }
 
@@ -274,12 +307,12 @@ int clients_write(clients_t* list, fd_set* set)
 {
   if(!list)
     return -1;
-  
+
   slist_element_t* tmp = list->list_.first_;
   while(tmp) {
     client_t* c = (client_t*)tmp->data_;
     tmp = tmp->next_;
-    if(c) {
+    if(c && c->state_ == CONNECTED) {
       int i;
       for(i=0; i<2; ++i) {
         if(FD_ISSET(c->fd_[i], set)) {
@@ -299,6 +332,10 @@ int clients_write(clients_t* list, fd_set* set)
           }
         }
       }
+    } else if(c && c->state_ == CONNECTING && FD_ISSET(c->fd_[1], set)) {
+      int ret = handle_connect(c, list->buffer_size_);
+      if(ret)
+        slist_remove(&(list->list_), c);
     }
   }
 
diff --git a/src/clients.h b/src/clients.h
index a6337c3..a24c044 100644
--- a/src/clients.h
+++ b/src/clients.h
@@ -35,10 +35,14 @@
 
 #define BUFFER_LENGTH 10240
 
+enum client_state_enum { CONNECTING, CONNECTED };
+typedef enum client_state_enum client_state_t;
+
 typedef struct {
   int fd_[2];
   buffer_t write_buf_[2];
   u_int32_t write_buf_offset_[2];
+  client_state_t state_;
 } client_t;
 
 void clients_delete_element(void* e);