From 3807b5766f3224fe5a9ce37a39b7d1a7ecf22c0b Mon Sep 17 00:00:00 2001
From: Max Lv <max.c.lv@gmail.com>
Date: Sat, 8 Nov 2014 18:07:10 +0800
Subject: [PATCH] refine fast open

---
 src/local.c | 59 ++++++++++++++++++++++++++---------------------------
 src/local.h |  2 --
 2 files changed, 29 insertions(+), 32 deletions(-)

diff --git a/src/local.c b/src/local.c
index e28db052..3eb14484 100644
--- a/src/local.c
+++ b/src/local.c
@@ -203,7 +203,7 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents)
         }
     }
 
-    while (r > 0)
+    while (1)
     {
         // local socks5 server
         if (server->stage == 5)
@@ -215,27 +215,9 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents)
                 return;
             }
 
-            // Copy to remote->buf
-            if (buf != remote->buf) {
-                memcpy(remote->buf, buf, r);
-            }
-
             // insert shadowsocks header
             if (!remote->direct)
             {
-                if (!remote->send_ctx->connected)
-                {
-                    char *tmp = malloc(max(BUF_SIZE, r + server->addr_len));
-
-                    memcpy(tmp, server->addr_to_send, server->addr_len);
-                    memcpy(tmp + server->addr_len, remote->buf, r);
-                    r += server->addr_len;
-
-                    // deallocate
-                    free(remote->buf);
-                    remote->buf = tmp;
-                }
-
                 remote->buf = ss_encrypt(BUF_SIZE, remote->buf, &r, server->e_ctx);
 
                 if (remote->buf == NULL)
@@ -342,7 +324,7 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents)
             }
 
             // all processed
-            r = 0;
+            return;
         }
         else if (server->stage == 0)
         {
@@ -352,8 +334,6 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents)
             char *send_buf = (char *)&response;
             send(server->fd, send_buf, sizeof(response), 0);
             server->stage = 1;
-            r -= 3;
-            buf += 3;
             return;
         }
         else if (server->stage == 1)
@@ -452,12 +432,10 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents)
                     return;
                 }
 
-                server->addr_to_send = ss_addr_to_send;
-                server->addr_len = addr_len;
                 server->stage = 5;
 
-                r -= (4 + addr_len);
-                buf += (4 + addr_len);
+                r -= (3 + addr_len);
+                buf += (3 + addr_len);
 
                 if (verbose)
                 {
@@ -486,6 +464,23 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents)
                     return;
                 }
 
+                if (!remote->direct)
+                {
+                    memcpy(remote->buf, ss_addr_to_send, addr_len);
+                    if (r > 0)
+                    {
+                        memcpy(remote->buf + addr_len, buf, r);
+                    }
+                    r += addr_len;
+                }
+                else
+                {
+                    if (r > 0)
+                    {
+                        memcpy(remote->buf, buf, r);
+                    }
+                }
+
                 server->remote = remote;
                 remote->server = server;
             }
@@ -675,6 +670,14 @@ static void remote_send_cb (EV_P_ ev_io *w, int revents)
             ev_timer_stop(EV_A_ &remote_send_ctx->watcher);
             ev_timer_start(EV_A_ &remote->recv_ctx->watcher);
             ev_io_start(EV_A_ &remote->recv_ctx->io);
+
+            // no need to send any data
+            if (remote->buf_len == 0)
+            {
+                ev_io_stop(EV_A_ &remote_send_ctx->io);
+                ev_io_start(EV_A_ &server->recv_ctx->io);
+                return;
+            }
         }
         else
         {
@@ -838,10 +841,6 @@ static void free_server(struct server *server)
         {
             free(server->buf);
         }
-        if (server->addr_to_send != NULL)
-        {
-            free(server->addr_to_send);
-        }
         free(server->recv_ctx);
         free(server->send_ctx);
         free(server);
diff --git a/src/local.h b/src/local.h
index 11198075..f213befa 100644
--- a/src/local.h
+++ b/src/local.h
@@ -54,8 +54,6 @@ struct server
     ssize_t buf_len;
     ssize_t buf_idx;
     char *buf; // server send from, remote recv into
-    int addr_len;
-    char *addr_to_send;
     char stage;
     struct enc_ctx *e_ctx;
     struct enc_ctx *d_ctx;