From 9d2c14b3357db88eca5a67b1302398bbe53b8bcf Mon Sep 17 00:00:00 2001
From: Max Lv <max.c.lv@gmail.com>
Date: Thu, 3 Sep 2015 12:29:14 +0800
Subject: [PATCH] encode auth flag in atyp

---
 src/encrypt.h | 4 ++++
 src/local.c   | 6 +++---
 src/redir.c   | 1 +
 src/server.c  | 8 ++++----
 src/tunnel.c  | 1 +
 5 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/encrypt.h b/src/encrypt.h
index 9ee37527..77bf5a5a 100644
--- a/src/encrypt.h
+++ b/src/encrypt.h
@@ -143,6 +143,10 @@ typedef struct {
 #define ONETIMEAUTH_BYTES 16U
 #define ONETIMEAUTH_KEYBYTES 32U
 
+#define ONETIMEAUTH_FLAG 0x10
+#define ONETIMEAUTH_MASK 0xF0
+#define ADDRTYPE_MASK 0xF
+
 #define min(a, b) (((a) < (b)) ? (a) : (b))
 #define max(a, b) (((a) > (b)) ? (a) : (b))
 
diff --git a/src/local.c b/src/local.c
index 6ca85baa..127d3cd5 100644
--- a/src/local.c
+++ b/src/local.c
@@ -471,13 +471,13 @@ static void server_recv_cb(EV_P_ ev_io *w, int revents)
                 }
 
                 if (!remote->direct) {
-                    memcpy(remote->buf, ss_addr_to_send, addr_len);
-
                     if (auth) {
-                        ss_onetimeauth(remote->buf + addr_len, ss_addr_to_send, addr_len);
+                        ss_addr_to_send[0] |= ONETIMEAUTH_FLAG;
+                        ss_onetimeauth(ss_addr_to_send + addr_len, ss_addr_to_send, addr_len);
                         addr_len += ONETIMEAUTH_BYTES;
                     }
 
+                    memcpy(remote->buf, ss_addr_to_send, addr_len);
                     if (r > 0) {
                         memcpy(remote->buf + addr_len, buf, r);
                     }
diff --git a/src/redir.c b/src/redir.c
index 0fcfd0c3..f0e042fc 100644
--- a/src/redir.c
+++ b/src/redir.c
@@ -380,6 +380,7 @@ static void remote_send_cb(EV_P_ ev_io *w, int revents)
             if (auth) {
                 ss_onetimeauth(ss_addr_to_send + addr_len, ss_addr_to_send, addr_len);
                 addr_len += ONETIMEAUTH_BYTES;
+                ss_addr_to_send[0] |= ONETIMEAUTH_FLAG;
             }
 
             int s = send(remote->fd, ss_addr_to_send, addr_len, 0);
diff --git a/src/server.c b/src/server.c
index c10b8665..dd298ab0 100644
--- a/src/server.c
+++ b/src/server.c
@@ -539,7 +539,7 @@ static void server_recv_cb(EV_P_ ev_io *w, int revents)
         memset(&storage, 0, sizeof(struct sockaddr_storage));
 
         // get remote addr and port
-        if (atyp == 1) {
+        if ((atyp & ADDRTYPE_MASK) == 1) {
             // IP V4
             struct sockaddr_in *addr = (struct sockaddr_in *)&storage;
             size_t in_addr_len = sizeof(struct in_addr);
@@ -561,7 +561,7 @@ static void server_recv_cb(EV_P_ ev_io *w, int revents)
             info.ai_protocol = IPPROTO_TCP;
             info.ai_addrlen = sizeof(struct sockaddr_in);
             info.ai_addr = (struct sockaddr *)addr;
-        } else if (atyp == 3) {
+        } else if ((atyp & ADDRTYPE_MASK) == 3) {
             // Domain name
             uint8_t name_len = *(uint8_t *)(server->buf + offset);
             if (name_len < r) {
@@ -597,7 +597,7 @@ static void server_recv_cb(EV_P_ ev_io *w, int revents)
             } else {
                 need_query = 1;
             }
-        } else if (atyp == 4) {
+        } else if ((atyp & ADDRTYPE_MASK) == 4) {
             // IP V6
             struct sockaddr_in6 *addr = (struct sockaddr_in6 *)&storage;
             size_t in6_addr_len = sizeof(struct in6_addr);
@@ -640,7 +640,7 @@ static void server_recv_cb(EV_P_ ev_io *w, int revents)
 
         offset += 2;
 
-        if (auth) {
+        if (auth || (atyp & ONETIMEAUTH_MASK)) {
             if (ss_onetimeauth_verify(server->buf + offset, server->buf, offset)) {
                 LOGE("authentication error %d", atyp);
                 report_addr(server->fd);
diff --git a/src/tunnel.c b/src/tunnel.c
index d44a0e00..b66f591d 100644
--- a/src/tunnel.c
+++ b/src/tunnel.c
@@ -422,6 +422,7 @@ static void remote_send_cb(EV_P_ ev_io *w, int revents)
             if (auth) {
                 ss_onetimeauth(ss_addr_to_send + addr_len, ss_addr_to_send, addr_len);
                 addr_len += ONETIMEAUTH_BYTES;
+                ss_addr_to_send[0] |= ONETIMEAUTH_FLAG;
             }
 
             int s = send(remote->fd, ss_addr_to_send, addr_len, 0);