From 018c5f8cc713f04df34adfceec9edda2ac10d040 Mon Sep 17 00:00:00 2001
From: Max Lv <max.c.lv@gmail.com>
Date: Sat, 19 Apr 2014 21:34:06 +0800
Subject: [PATCH] fix udprelay

---
 src/include.h  |  2 +-
 src/server.c   |  2 +-
 src/udprelay.c | 44 ++++++++++++++++++++++++++++----------------
 src/udprelay.h |  9 +++++++--
 4 files changed, 37 insertions(+), 20 deletions(-)

diff --git a/src/include.h b/src/include.h
index f2c3ad09..c5809626 100644
--- a/src/include.h
+++ b/src/include.h
@@ -9,7 +9,7 @@ int udprelay_init(const char *server_host, const char *server_port,
 #endif
 #endif
 #ifdef UDPRELAY_REMOTE
-                  asyncns_t *asyncns,
+                  int dns_thread_num,
 #endif
                   int method, int timeout, const char *iface);
 
diff --git a/src/server.c b/src/server.c
index 6d09d554..3536a01b 100644
--- a/src/server.c
+++ b/src/server.c
@@ -1069,7 +1069,7 @@ int main (int argc, char **argv)
     if (udprelay)
     {
         LOGD("udprelay enabled.");
-        udprelay_init(server_host[0], server_port, asyncns, m, listen_ctx->timeout, iface);
+        udprelay_init(server_host[0], server_port, dns_thread_num, m, listen_ctx->timeout, iface);
     }
 
     // setuid
diff --git a/src/udprelay.c b/src/udprelay.c
index 8e3618ae..c0b71373 100644
--- a/src/udprelay.c
+++ b/src/udprelay.c
@@ -316,7 +316,6 @@ struct query_ctx *new_query_ctx(asyncns_query_t *query,
     ctx->buf_len = buf_len;
     memcpy(ctx->buf, buf, buf_len);
     ctx->query = query;
-    ev_timer_init(&ctx->watcher, query_resolve_cb, 0.1, 0.2);
     return ctx;
 }
 
@@ -324,7 +323,6 @@ void close_and_free_query(EV_P_ struct query_ctx *ctx)
 {
     if (ctx != NULL)
     {
-        ev_timer_stop(EV_A_ &ctx->watcher);
         if (ctx->buf != NULL)
         {
             free(ctx->buf);
@@ -362,27 +360,28 @@ static void remote_timeout_cb(EV_P_ ev_timer *watcher, int revents)
 }
 
 #ifdef UDPRELAY_REMOTE
-static void query_resolve_cb(EV_P_ ev_timer *watcher, int revents)
+static void query_resolve_cb(EV_P_ ev_io *w, int revents)
 {
     int err;
     struct addrinfo *result, *rp;
-    struct query_ctx *query_ctx = (struct query_ctx *)((void*)watcher);
-    asyncns_t *asyncns = query_ctx->server_ctx->asyncns;
-    asyncns_query_t *query = query_ctx->query;
+    struct resolve_ctx *resolve_ctx = (struct resolve_ctx*)w;
+    asyncns_t *asyncns = resolve_ctx->asyncns;
 
-    if (asyncns == NULL || query == NULL)
+    err = asyncns_handle(asyncns);
+    if (err == ASYNCNS_HANDLE_AGAIN)
     {
-        LOGE("invalid dns query.");
-        close_and_free_query(EV_A_ query_ctx);
+        // try again
         return;
     }
-
-    if (asyncns_wait(asyncns, 0) == -1)
+    else if (err == ASYNCNS_HANDLE_ERROR)
     {
         // asyncns error
         FATAL("asyncns exit unexpectedly.");
     }
 
+    asyncns_query_t *query = asyncns_getnext(asyncns);
+    struct query_ctx *query_ctx= (struct query_ctx*) asyncns_getuserdata(asyncns, query);
+
     if (!asyncns_isdone(asyncns, query))
     {
         // wait reolver
@@ -394,7 +393,7 @@ static void query_resolve_cb(EV_P_ ev_timer *watcher, int revents)
         LOGD("[udp] asyncns resolved.");
     }
 
-    ev_timer_stop(EV_A_ watcher);
+    query_ctx->query = NULL;
 
     err = asyncns_getaddrinfo_done(asyncns, query, &result);
 
@@ -821,9 +820,7 @@ static void server_recv_cb (EV_P_ ev_io *w, int revents)
         query_ctx->addr_header_len = addr_header_len;
         query_ctx->src_addr = src_addr;
         memcpy(query_ctx->addr_header, addr_header, addr_header_len);
-
-        ev_timer_start(EV_A_ &query_ctx->watcher);
-
+        asyncns_setuserdata(server_ctx->asyncns, query, query_ctx);
     }
     else
     {
@@ -862,7 +859,7 @@ int udprelay_init(const char *server_host, const char *server_port,
 #endif
 #endif
 #ifdef UDPRELAY_REMOTE
-                  asyncns_t *asyncns,
+                  int dns_thread_num,
 #endif
                   int method, int timeout, const char *iface)
 {
@@ -877,6 +874,21 @@ int udprelay_init(const char *server_host, const char *server_port,
     //////////////////////////////////////////////////
     // Setup server context
 
+#ifdef UDPRELAY_REMOTE
+    // setup asyncns
+    asyncns_t *asyncns;
+    if (!(asyncns = asyncns_new(dns_thread_num)))
+    {
+        FATAL("asyncns failed");
+    }
+    struct resolve_ctx *resolve_ctx = malloc(sizeof(struct resolve_ctx));
+    resolve_ctx->asyncns = asyncns;
+
+    int asyncnsfd = asyncns_fd(asyncns);
+    ev_io_init (&resolve_ctx->io, query_resolve_cb, asyncnsfd, EV_READ);
+    ev_io_start (loop, &resolve_ctx->io);
+#endif
+
     // Bind to port
     int serverfd = create_server_socket(server_host, server_port);
     if (serverfd < 0)
diff --git a/src/udprelay.h b/src/udprelay.h
index 40915f85..095a1ab9 100644
--- a/src/udprelay.h
+++ b/src/udprelay.h
@@ -39,9 +39,14 @@ struct server_ctx
 };
 
 #ifdef UDPRELAY_REMOTE
+struct resolve_ctx
+{
+    ev_io io;
+    asyncns_t *asyncns;
+};
+
 struct query_ctx
 {
-    ev_timer watcher;
     asyncns_query_t *query;
     struct sockaddr src_addr;
     int buf_len;
@@ -69,7 +74,7 @@ static void remote_recv_cb (EV_P_ ev_io *w, int revents);
 static void remote_timeout_cb(EV_P_ ev_timer *watcher, int revents);
 static char *hash_key(const char *header, const int header_len, const struct sockaddr *addr);
 #ifdef UDPRELAY_REMOTE
-static void query_resolve_cb(EV_P_ ev_timer *watcher, int revents);
+static void query_resolve_cb(EV_P_ ev_io *w, int revents);
 #endif
 static void close_and_free_remote(EV_P_ struct remote_ctx *ctx);