From 60135461ea9859c39413c94dc7aa30be4010153c Mon Sep 17 00:00:00 2001
From: Max Lv <max.c.lv@gmail.com>
Date: Tue, 25 Jun 2013 13:23:29 +0800
Subject: [PATCH] handle incomplete header

---
 src/encrypt.c |  4 ++++
 src/encrypt.h |  1 +
 src/server.c  | 30 +++++++++++++++++++++++-------
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/src/encrypt.c b/src/encrypt.c
index 41c30d5c..419788c2 100644
--- a/src/encrypt.c
+++ b/src/encrypt.c
@@ -117,6 +117,10 @@ static void merge_sort(uint8_t array[], int length,
 	merge(left, llength, right, middle, salt, key);
 }
 
+int enc_get_iv_len() {
+    return enc_iv_len;
+}
+
 void enc_table_init(const char *pass) {
     uint32_t i;
     uint32_t salt;
diff --git a/src/encrypt.h b/src/encrypt.h
index 2a8d6755..cbc8273b 100644
--- a/src/encrypt.h
+++ b/src/encrypt.h
@@ -48,5 +48,6 @@ char* ss_encrypt(char *plaintext, ssize_t *len, struct enc_ctx *ctx);
 char* ss_decrypt(char *ciphertext, ssize_t *len, struct enc_ctx *ctx);
 void enc_ctx_init(int method, struct enc_ctx *ctx, int enc);
 int enc_init(const char *pass, const char *method);
+int enc_get_iv_len(void);
 
 #endif // _ENCRYPT_H
diff --git a/src/server.c b/src/server.c
index 516f5edf..ef747ca3 100644
--- a/src/server.c
+++ b/src/server.c
@@ -140,6 +140,7 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents) {
     struct server *server = server_recv_ctx->server;
     struct remote *remote = NULL;
 
+    int len = server->buf_len;
     char **buf = &server->buf;
 
     ev_timer_again(EV_A_ &server->recv_ctx->watcher);
@@ -147,9 +148,10 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents) {
     if (server->stage != 0) {
         remote = server->remote;
         buf = &remote->buf;
+        len = 0;
     }
 
-    ssize_t r = recv(server->fd, *buf, BUF_SIZE, 0);
+    ssize_t r = recv(server->fd, *buf + len, BUF_SIZE - len, 0);
 
     if (r == 0) {
         // connection closed
@@ -172,6 +174,21 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents) {
         }
     }
 
+    // handle incomplete header
+    if (server->stage == 0) {
+        r += server->buf_len;
+        if (r <= enc_get_iv_len()) {
+            // wait for more
+            if (verbose) {
+                LOGD("imcomplete header: %zu", r);
+            }
+            server->buf_len = r;
+            return;
+        } else {
+            server->buf_len = 0;
+        }
+    }
+
     *buf = ss_decrypt(*buf, &r, server->d_ctx);
 
     // handshake and transmit data
@@ -196,6 +213,7 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents) {
             ev_io_start(EV_A_ &remote->send_ctx->io);
         }
         return;
+
     } else if (server->stage == 0) {
 
         /*
@@ -224,7 +242,6 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents) {
                         host, INET_ADDRSTRLEN);
                 offset += in_addr_len;
             }
-
         } else if (atyp == 3) {
             // Domain name
             uint8_t name_len = *(uint8_t *)(server->buf + offset);
@@ -234,7 +251,6 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents) {
             }
             memcpy(host, server->buf + offset + 1, name_len);
             offset += name_len + 1;
-
         } else if (atyp == 4) {
             // IP V6
             size_t in6_addr_len = sizeof(struct in6_addr);
@@ -243,10 +259,10 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents) {
                         host, INET6_ADDRSTRLEN);
                 offset += in6_addr_len;
             }
-
         }
         
         if (offset == 0) {
+            LOGE("incorrect header with length %zu", r);
             close_and_free_server(EV_A_ server);
             return;
         }
@@ -816,17 +832,17 @@ int main (int argc, char **argv) {
     signal(SIGPIPE, SIG_IGN);
     signal(SIGCHLD, SIG_IGN);
 
-    // Setup asyncns
+    // setup asyncns
     asyncns_t *asyncns;
     if (!(asyncns = asyncns_new(dns_thread_num))) {
         FATAL("asyncns failed");
     }
 
-    // Setup keys
+    // setup keys
     LOGD("initialize cihpers... %s", method);
     int m = enc_init(password, method);
 
-    // Inilitialize ev loop
+    // inilitialize ev loop
     struct ev_loop *loop = EV_DEFAULT;
 
     // bind to each interface