From 3493ac3e77e29c97a910a41372d924b3eacad4f6 Mon Sep 17 00:00:00 2001
From: dev0 <dev@watfaq.com>
Date: Mon, 25 Dec 2023 22:47:40 +1100
Subject: [PATCH] add wg dns timeout

---
 clash/tests/data/config/wg.yaml     |  2 +-
 clash_lib/src/proxy/wg/device.rs    | 14 +++++++++++---
 clash_lib/src/proxy/wg/wireguard.rs | 16 +++++++++-------
 3 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/clash/tests/data/config/wg.yaml b/clash/tests/data/config/wg.yaml
index d238645df..dae22bb5c 100644
--- a/clash/tests/data/config/wg.yaml
+++ b/clash/tests/data/config/wg.yaml
@@ -52,7 +52,7 @@ proxies:
     allowed-ips: ['0.0.0.0/0']
     remote-dns-resolve: true
     dns:
-      - 1.1.1.1
+      - 8.8.8.8
     udp: true
 
 
diff --git a/clash_lib/src/proxy/wg/device.rs b/clash_lib/src/proxy/wg/device.rs
index 6624758bd..23bc0bed5 100644
--- a/clash_lib/src/proxy/wg/device.rs
+++ b/clash_lib/src/proxy/wg/device.rs
@@ -148,14 +148,22 @@ impl DeviceManager {
 
         let pkt = UdpPacket::new(
             msg.to_vec().unwrap().into(),
-            (self.addr, self.get_ephemeral_udp_port().await).into(),
+            SocksAddr::any_ipv4(),
             server.into(),
         );
 
         socket.feed(pkt).await.ok()?;
-        socket.close().await.ok()?;
+        socket.flush().await.ok()?;
+        trace!("sent dns query: {:?}", msg);
+
+        let pkt = match tokio::time::timeout(Duration::from_secs(5), socket.next()).await {
+            Ok(Some(pkt)) => pkt,
+            _ => {
+                warn!("wg dns query timed out with server {server}");
+                return None;
+            }
+        };
 
-        let pkt = socket.next().await?;
         let msg = hickory_proto::op::Message::from_vec(&pkt.data).ok()?;
         trace!("got dns response: {:?}", msg);
         msg.answers()
diff --git a/clash_lib/src/proxy/wg/wireguard.rs b/clash_lib/src/proxy/wg/wireguard.rs
index 0f0915bcb..1816169aa 100644
--- a/clash_lib/src/proxy/wg/wireguard.rs
+++ b/clash_lib/src/proxy/wg/wireguard.rs
@@ -109,7 +109,6 @@ impl WireguardTunnel {
             }
             boringtun::noise::TunnResult::WriteToNetwork(packet) => {
                 self.udp.send_to(&packet, self.endpoint).await?;
-                trace!("sent packet to {}", self.endpoint);
             }
             _ => {
                 error!("unexpected result from encapsulate");
@@ -133,8 +132,6 @@ impl WireguardTunnel {
     }
 
     pub async fn start_forwarding(&self) {
-        trace!("wg stack writing data");
-
         let mut packet_reader = self.packet_reader.lock().await;
         loop {
             match packet_reader.recv().await {
@@ -235,6 +232,8 @@ impl WireguardTunnel {
                 }
 
                 TunnResult::WriteToTunnelV4(packet, addr) => {
+                    trace_ip_packet("Received IP packet", packet);
+
                     if !self.is_ip_allowed(addr.into()) {
                         trace!(
                             "received packet from {} which is not in allowed_ips",
@@ -246,7 +245,7 @@ impl WireguardTunnel {
                     let _ =
                         trace_span!("wg_write_stack", endpoint = %self.endpoint, size = packet.len())
                             .entered();
-                    trace_ip_packet("Received IP packet", packet);
+
                     if let Some(proto) = self.route_protocol(packet) {
                         if let Err(e) = self
                             .packet_writer
@@ -256,10 +255,12 @@ impl WireguardTunnel {
                             error!("failed to send packet to virtual device: {}", e);
                         }
                     } else {
-                        trace!("wg stack recevied unkown data");
+                        warn!("wg stack recevied unkown data");
                     }
                 }
                 TunnResult::WriteToTunnelV6(packet, addr) => {
+                    trace_ip_packet("Received IP packet", packet);
+
                     if !self.is_ip_allowed(addr.into()) {
                         trace!(
                             "received packet from {} which is not in allowed_ips",
@@ -271,7 +272,6 @@ impl WireguardTunnel {
                     let _ =
                         trace_span!("wg_write_stack", endpoint = %self.endpoint, size = packet.len())
                             .entered();
-                    trace_ip_packet("Received IP packet", packet);
                     if let Some(proto) = self.route_protocol(packet) {
                         if let Err(e) = self
                             .packet_writer
@@ -281,7 +281,7 @@ impl WireguardTunnel {
                             error!("failed to send packet to virtual device: {}", e);
                         }
                     } else {
-                        trace!("wg stack recevied unkown data");
+                        warn!("wg stack recevied unkown data");
                     }
                 }
             }
@@ -299,6 +299,8 @@ impl WireguardTunnel {
                 let mut buf = vec![0u8; 65535];
                 let mut peer = self.peer.lock().await;
                 let tun_result = peer.format_handshake_initiation(&mut buf[..], false);
+                drop(peer);
+
                 self.handle_routine_result(tun_result).await;
             }
             TunnResult::Err(e) => {