From 30420059cc737884a4aa34b9a7bd1fa05ba5a698 Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Fri, 11 Aug 2023 22:54:20 +0800
Subject: [PATCH] send_udp_packet

---
 src/socks.rs     | 16 ++++++----------
 src/tun2proxy.rs | 25 ++++++++++++++-----------
 2 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/src/socks.rs b/src/socks.rs
index a1af57c..3e4a7a9 100644
--- a/src/socks.rs
+++ b/src/socks.rs
@@ -158,11 +158,11 @@ impl SocksProxyImpl {
             return Err("SOCKS5 server requires an unsupported authentication method.".into());
         }
 
-        if auth_method == AuthMethod::UserPass {
-            self.state = SocksState::SendAuthData;
+        self.state = if auth_method == AuthMethod::UserPass {
+            SocksState::SendAuthData
         } else {
-            self.state = SocksState::SendRequest;
-        }
+            SocksState::SendRequest
+        };
         self.state_change()
     }
 
@@ -344,12 +344,8 @@ impl ConnectionManager for SocksProxyManager {
     fn new_tcp_proxy(&self, info: &ConnectionInfo, udp_associate: bool) -> Result<Box<dyn TcpProxy>> {
         use socks5_impl::protocol::Command::{Connect, UdpAssociate};
         let command = if udp_associate { UdpAssociate } else { Connect };
-        Ok(Box::new(SocksProxyImpl::new(
-            info,
-            self.credentials.clone(),
-            self.version,
-            command,
-        )?))
+        let credentials = self.credentials.clone();
+        Ok(Box::new(SocksProxyImpl::new(info, credentials, self.version, command)?))
     }
 
     fn close_connection(&self, _: &ConnectionInfo) {}
diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs
index 687392c..22e538f 100644
--- a/src/tun2proxy.rs
+++ b/src/tun2proxy.rs
@@ -484,17 +484,7 @@ impl<'a> TunToProxy<'a> {
                 let payload = &frame[payload_offset..payload_offset + payload_size];
                 if let (Some(virtual_dns), true) = (&mut self.options.virtual_dns, port == 53) {
                     let response = virtual_dns.receive_query(payload)?;
-                    {
-                        let rx_buffer = udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]);
-                        let tx_buffer = udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]);
-                        let mut socket = udp::Socket::new(rx_buffer, tx_buffer);
-                        socket.bind(dst)?;
-                        let meta = UdpMetadata::from(connection_info.src);
-                        socket.send_slice(response.as_slice(), meta)?;
-                        let handle = self.sockets.add(socket);
-                        self.expect_smoltcp_send()?;
-                        self.sockets.remove(handle);
-                    }
+                    self.send_udp_packet(dst, connection_info.src, response.as_slice())?;
                 } else {
                     // Another UDP packet
                     if !self.connection_map.contains_key(&connection_info) {
@@ -516,6 +506,7 @@ impl<'a> TunToProxy<'a> {
                     if let Some(udp_associate) = state.tcp_proxy_handler.get_udp_associate() {
                         log::debug!("UDP associate address: {}", udp_associate);
                         // Send packets via UDP associate...
+                        // self.send_udp_packet(connection_info.src, udp_associate, &s5_udp_data)?;
                     } else {
                         // UDP associate tunnel not ready yet, we must cache the packet...
                     }
@@ -562,6 +553,18 @@ impl<'a> TunToProxy<'a> {
         Ok(state)
     }
 
+    fn send_udp_packet(&mut self, src: SocketAddr, dst: SocketAddr, data: &[u8]) -> Result<()> {
+        let rx_buffer = udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]);
+        let tx_buffer = udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]);
+        let mut socket = udp::Socket::new(rx_buffer, tx_buffer);
+        socket.bind(src)?;
+        socket.send_slice(data, UdpMetadata::from(dst))?;
+        let handle = self.sockets.add(socket);
+        self.expect_smoltcp_send()?;
+        self.sockets.remove(handle);
+        Ok(())
+    }
+
     fn write_to_server(&mut self, info: &ConnectionInfo) -> Result<(), Error> {
         if let Some(state) = self.connection_map.get_mut(info) {
             let event = state.tcp_proxy_handler.peek_data(OutgoingDirection::ToServer);