From 91d25d3cc0d00c4339908b7b24a5fcbc58d78b7e Mon Sep 17 00:00:00 2001
From: zonyitoo <zonyitoo@gmail.com>
Date: Mon, 28 Dec 2020 13:23:00 +0800
Subject: [PATCH] ignore mode in local-dns local resolver

ref shadowsocks/shadowsocks-android#2635
---
 Cargo.lock                                    |  6 +--
 .../src/local/dns/server.rs                   | 42 ++++++++-----------
 2 files changed, 20 insertions(+), 28 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 21881600b76e..3047670b665e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2061,7 +2061,7 @@ checksum = "efd1f82c56340fdf16f2a953d7bda4f8fdffba13d93b00844c25572110b26079"
 [[package]]
 name = "trust-dns-proto"
 version = "0.20.0-alpha.3"
-source = "git+https://github.com/bluejekyll/trust-dns.git?branch=main#540906be77b03dc59494aef453b2dbcc875a2a0a"
+source = "git+https://github.com/bluejekyll/trust-dns.git?branch=main#57245de7d73008c9f960ba8c952e2dbcee54bb84"
 dependencies = [
  "async-trait",
  "cfg-if 1.0.0",
@@ -2085,7 +2085,7 @@ dependencies = [
 [[package]]
 name = "trust-dns-resolver"
 version = "0.20.0-alpha.3"
-source = "git+https://github.com/bluejekyll/trust-dns.git?branch=main#540906be77b03dc59494aef453b2dbcc875a2a0a"
+source = "git+https://github.com/bluejekyll/trust-dns.git?branch=main#57245de7d73008c9f960ba8c952e2dbcee54bb84"
 dependencies = [
  "cfg-if 1.0.0",
  "futures-util",
@@ -2109,7 +2109,7 @@ dependencies = [
 [[package]]
 name = "trust-dns-rustls"
 version = "0.20.0-alpha.3"
-source = "git+https://github.com/bluejekyll/trust-dns.git?branch=main#540906be77b03dc59494aef453b2dbcc875a2a0a"
+source = "git+https://github.com/bluejekyll/trust-dns.git?branch=main#57245de7d73008c9f960ba8c952e2dbcee54bb84"
 dependencies = [
  "futures-channel",
  "futures-io",
diff --git a/crates/shadowsocks-service/src/local/dns/server.rs b/crates/shadowsocks-service/src/local/dns/server.rs
index e1f460b2aa50..9786b2a97399 100644
--- a/crates/shadowsocks-service/src/local/dns/server.rs
+++ b/crates/shadowsocks-service/src/local/dns/server.rs
@@ -678,36 +678,28 @@ impl DnsClient {
         message.set_recursion_desired(true);
         message.add_query(query.clone());
 
-        let mut last_err = io::Error::new(ErrorKind::InvalidData, "resolve empty");
-
         // Query UDP then TCP
 
-        if self.mode.enable_udp() {
-            match self
-                .client_cache
-                .lookup_udp_local(local_addr, message.clone(), self.context.connect_opts_ref())
-                .await
-            {
-                Ok(msg) => return Ok(msg),
-                Err(err) => {
-                    last_err = err.into();
-                }
-            }
-        }
+        let udp_query =
+            self.client_cache
+                .lookup_udp_local(local_addr, message.clone(), self.context.connect_opts_ref());
+        let tcp_query = async move {
+            // Send TCP query after 500ms, because UDP will always return faster than TCP, there is no need to send queries simutaneously
+            time::sleep(Duration::from_millis(500)).await;
 
-        if self.mode.enable_tcp() {
-            match self
-                .client_cache
+            self.client_cache
                 .lookup_tcp_local(local_addr, message, self.context.connect_opts_ref())
                 .await
-            {
-                Ok(msg) => return Ok(msg),
-                Err(err) => {
-                    last_err = err.into();
-                }
-            }
-        }
+        };
 
-        Err(last_err)
+        tokio::pin!(udp_query);
+        tokio::pin!(tcp_query);
+
+        match future::select(udp_query, tcp_query).await {
+            Either::Left((Ok(m), ..)) => Ok(m),
+            Either::Left((Err(..), next)) => next.await.map_err(From::from),
+            Either::Right((Ok(m), ..)) => Ok(m),
+            Either::Right((Err(..), next)) => next.await.map_err(From::from),
+        }
     }
 }