From fa82628dd938ee5285cbce481cc722539d8ade2a Mon Sep 17 00:00:00 2001
From: Max Lv <max.c.lv@gmail.com>
Date: Fri, 17 Mar 2017 09:54:03 +0800
Subject: [PATCH] Add delayed connect to ss-local as well

---
 src/common.h |  3 ++-
 src/local.c  | 54 ++++++++++++++++++++++++++++++++++++----------------
 src/local.h  |  2 ++
 3 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/src/common.h b/src/common.h
index 44b9cf1a8..f7ab3d956 100644
--- a/src/common.h
+++ b/src/common.h
@@ -55,7 +55,8 @@ int send_traffic_stat(uint64_t tx, uint64_t rx);
 #define STAGE_HANDSHAKE  1  /* Handshake with client            */
 #define STAGE_PARSE      2  /* Parse the header                 */
 #define STAGE_RESOLVE    4  /* Resolve the hostname             */
-#define STAGE_STREAM     5  /* Stream between client and server */
+#define STAGE_WAIT       5  /* Wait for more data               */
+#define STAGE_STREAM     6  /* Stream between client and server */
 
 /* Vals for long options */
 enum {
diff --git a/src/local.c b/src/local.c
index 0553e7105..475a18e79 100644
--- a/src/local.c
+++ b/src/local.c
@@ -247,6 +247,16 @@ free_connections(struct ev_loop *loop)
     }
 }
 
+static void
+delayed_connect_cb(EV_P_ ev_timer *watcher, int revents)
+{
+    server_t *server = cork_container_of(watcher, server_t,
+                                         delayed_connect_watcher);
+
+    server->stage = STAGE_WAIT;
+    server_recv_cb(EV_A_ & server->recv_ctx->io, revents);
+}
+
 static void
 server_recv_cb(EV_P_ ev_io *w, int revents)
 {
@@ -262,29 +272,32 @@ server_recv_cb(EV_P_ ev_io *w, int revents)
         buf = remote->buf;
     }
 
-    r = recv(server->fd, buf->data + buf->len, BUF_SIZE - buf->len, 0);
+    if (server->stage != STAGE_WAIT) {
+        r = recv(server->fd, buf->data + buf->len, BUF_SIZE - buf->len, 0);
 
-    if (r == 0) {
-        // connection closed
-        close_and_free_remote(EV_A_ remote);
-        close_and_free_server(EV_A_ server);
-        return;
-    } else if (r == -1) {
-        if (errno == EAGAIN || errno == EWOULDBLOCK) {
-            // no data
-            // continue to wait for recv
-            return;
-        } else {
-            if (verbose)
-                ERROR("server_recv_cb_recv");
+        if (r == 0) {
+            // connection closed
             close_and_free_remote(EV_A_ remote);
             close_and_free_server(EV_A_ server);
             return;
+        } else if (r == -1) {
+            if (errno == EAGAIN || errno == EWOULDBLOCK) {
+                // no data
+                // continue to wait for recv
+                return;
+            } else {
+                if (verbose)
+                    ERROR("server_recv_cb_recv");
+                close_and_free_remote(EV_A_ remote);
+                close_and_free_server(EV_A_ server);
+                return;
+            }
         }
+        buf->len += r;
+    } else {
+        server->stage = STAGE_STREAM;
     }
 
-    buf->len += r;
-
     while (1) {
         // local socks5 server
         if (server->stage == STAGE_STREAM) {
@@ -294,6 +307,8 @@ server_recv_cb(EV_P_ ev_io *w, int revents)
                 return;
             }
 
+            ev_timer_stop(EV_A_ & server->delayed_connect_watcher);
+
             // insert shadowsocks header
             if (!remote->direct) {
 #ifdef ANDROID
@@ -743,6 +758,10 @@ server_recv_cb(EV_P_ ev_io *w, int revents)
 
             server->remote = remote;
             remote->server = server;
+
+            ev_timer_start(EV_A_ & server->delayed_connect_watcher);
+
+            return;
         }
     }
 }
@@ -1048,6 +1067,9 @@ new_server(int fd)
     ev_io_init(&server->recv_ctx->io, server_recv_cb, fd, EV_READ);
     ev_io_init(&server->send_ctx->io, server_send_cb, fd, EV_WRITE);
 
+    ev_timer_init(&server->delayed_connect_watcher,
+            delayed_connect_cb, 0.05, 0);
+
     cork_dllist_add(&connections, &server->entries);
 
     return server;
diff --git a/src/local.h b/src/local.h
index 61eb97b41..5c5109070 100644
--- a/src/local.h
+++ b/src/local.h
@@ -67,6 +67,8 @@ typedef struct server {
     buffer_t *buf;
     buffer_t *abuf;
 
+    ev_timer delayed_connect_watcher;
+
     struct cork_dllist_item entries;
 } server_t;