From c18ae5f6717e40081594083b4ae0044a737f7af0 Mon Sep 17 00:00:00 2001
From: Max Lv <mlv@nvidia.com>
Date: Fri, 25 Dec 2015 12:40:46 +0800
Subject: [PATCH] Fix #494

---
 src/server.c | 84 ++++++++++++++++++++++++++++++++++------------------
 1 file changed, 55 insertions(+), 29 deletions(-)

diff --git a/src/server.c b/src/server.c
index 600b516e..0a51c09c 100644
--- a/src/server.c
+++ b/src/server.c
@@ -102,9 +102,10 @@ static void free_remote(remote_t *remote);
 static void close_and_free_remote(EV_P_ remote_t *remote);
 static void free_server(server_t *server);
 static void close_and_free_server(EV_P_ server_t *server);
-
 static void server_resolve_cb(struct sockaddr *addr, void *data);
 
+static size_t parse_header_len(const char atyp, const char *data, size_t offset);
+
 int verbose = 0;
 
 static int white_list = 0;
@@ -215,6 +216,24 @@ static void free_connections(struct ev_loop *loop)
     }
 }
 
+static size_t parse_header_len(const char atyp, const char *data, size_t offset)
+{
+    size_t len = 0;
+    if ((atyp & ADDRTYPE_MASK) == 1) {
+        // IP V4
+        len += sizeof(struct in_addr);
+    } else if ((atyp & ADDRTYPE_MASK) == 3) {
+        // Domain name
+        uint8_t name_len = *(uint8_t *)(data + offset);
+        len += name_len + 1;
+    } else if ((atyp & ADDRTYPE_MASK) == 4) {
+        // IP V6
+        len += sizeof(struct in6_addr);
+    }
+    len += 2;
+    return len;
+}
+
 static char *get_peer_name(int fd)
 {
     static char peer_name[INET6_ADDRSTRLEN] = { 0 };
@@ -603,6 +622,39 @@ static void server_recv_cb(EV_P_ ev_io *w, int revents)
         memset(&info, 0, sizeof(struct addrinfo));
         memset(&storage, 0, sizeof(struct sockaddr_storage));
 
+        if (auth || (atyp & ONETIMEAUTH_FLAG)) {
+            size_t header_len = parse_header_len(atyp, server->buf->array, offset);
+            size_t len        = server->buf->len;
+
+            if (len < offset + header_len + ONETIMEAUTH_BYTES) {
+                report_addr(server->fd);
+                close_and_free_server(EV_A_ server);
+                return;
+            }
+
+            server->buf->len = offset + header_len + ONETIMEAUTH_BYTES;
+            if (ss_onetimeauth_verify(server->buf, server->d_ctx->evp.iv)) {
+                char *peer_name = get_peer_name(server->fd);
+                if (peer_name) {
+                    LOGE("authentication error from %s", peer_name);
+                    if (acl) {
+                        if (acl_get_mode() == BLACK_LIST) {
+                            acl_add_ip(peer_name);
+                            LOGE("add %s to the black list", peer_name);
+                        } else {
+                            acl_remove_ip(peer_name);
+                            LOGE("remove %s from the white list", peer_name);
+                        }
+                    }
+                }
+                close_and_free_server(EV_A_ server);
+                return;
+            }
+
+            server->buf->len = len;
+            server->auth     = 1;
+        }
+
         // get remote addr and port
         if ((atyp & ADDRTYPE_MASK) == 1) {
             // IP V4
@@ -697,34 +749,8 @@ static void server_recv_cb(EV_P_ ev_io *w, int revents)
 
         offset += 2;
 
-        if (auth || (atyp & ONETIMEAUTH_FLAG)) {
-            if (server->buf->len < offset + ONETIMEAUTH_BYTES) {
-                report_addr(server->fd);
-                close_and_free_server(EV_A_ server);
-                return;
-            }
-            size_t len = server->buf->len;
-            server->buf->len = offset + ONETIMEAUTH_BYTES;
-            if (ss_onetimeauth_verify(server->buf, server->d_ctx->evp.iv)) {
-                char *peer_name = get_peer_name(server->fd);
-                if (peer_name) {
-                    LOGE("authentication error from %s", peer_name);
-                    if (acl) {
-                        if (acl_get_mode() == BLACK_LIST) {
-                            acl_add_ip(peer_name);
-                            LOGE("add %s to the black list", peer_name);
-                        } else {
-                            acl_remove_ip(peer_name);
-                            LOGE("remove %s from the white list", peer_name);
-                        }
-                    }
-                }
-                close_and_free_server(EV_A_ server);
-                return;
-            }
-            server->buf->len = len;
-            offset          += ONETIMEAUTH_BYTES;
-            server->auth     = 1;
+        if (server->auth) {
+            offset += ONETIMEAUTH_BYTES;
         }
 
         if (server->buf->len < offset) {