fix: UDP upstream remote address

Banyc · Jan 14, 2024 · 8ad1575 · 8ad1575
1 parent 3176a64
commit 8ad1575
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 6 deletions.
diff --git a/access_server/src/socks5/server/udp.rs b/access_server/src/socks5/server/udp.rs
@@ -134,6 +134,7 @@ impl Socks5ServerUdpAccess {
         let upstream =
             UdpProxyClient::establish(proxy_chain.chain.clone(), flow.flow().upstream.0.clone())
                 .await?;
+        let upstream_remote = upstream.remote_addr().clone();
 
         let (upstream_read, upstream_write) = upstream.into_split();
 
@@ -167,7 +168,7 @@ impl Socks5ServerUdpAccess {
                 response_header: Some(response_header),
             };
             let _ = io_copy
-                .serve_as_access_server(session_table, upstream_local, "SOCKS UDP")
+                .serve_as_access_server(session_table, upstream_local, upstream_remote, "SOCKS UDP")
                 .await;
         });
         Ok(())

diff --git a/access_server/src/udp/mod.rs b/access_server/src/udp/mod.rs
@@ -141,6 +141,7 @@ impl UdpAccess {
         let proxy_chain = self.proxy_table.choose_chain();
         let upstream =
             UdpProxyClient::establish(proxy_chain.chain.clone(), self.destination.clone()).await?;
+        let upstream_remote = upstream.remote_addr().clone();
 
         let (upstream_read, upstream_write) = upstream.into_split();
 
@@ -164,7 +165,7 @@ impl UdpAccess {
                 response_header: None,
             };
             let _ = io_copy
-                .serve_as_access_server(session_table, upstream_local, "UDP")
+                .serve_as_access_server(session_table, upstream_local, upstream_remote, "UDP")
                 .await;
         });
         Ok(())

diff --git a/common/src/udp/io_copy.rs b/common/src/udp/io_copy.rs
@@ -17,6 +17,7 @@ use tokio_throughput::{ReadGauge, WriteGauge};
 use tracing::{info, trace, warn};
 
 use crate::{
+    addr::InternetAddr,
     error::AnyError,
     udp::{metrics::FlowRecord, TIMEOUT},
 };
@@ -101,6 +102,7 @@ where
         self,
         session_table: Option<UdpSessionTable>,
         upstream_local: Option<SocketAddr>,
+        upstream_remote: InternetAddr,
         log_prefix: &str,
     ) -> Result<FlowMetrics, CopyBiError> {
         let session = session_table.as_ref().map(|s| {
@@ -114,7 +116,7 @@ where
                 end: None,
                 destination: Some(self.flow.flow().upstream.0.clone()),
                 upstream_local,
-                upstream_remote: self.flow.flow().upstream.0.clone(),
+                upstream_remote,
                 downstream_remote: self.flow.flow().downstream.0,
                 up_gauge: Mutex::new(up_handle),
                 dn_gauge: Mutex::new(dn_handle),

diff --git a/proxy_client/src/udp.rs b/proxy_client/src/udp.rs
@@ -28,6 +28,7 @@ use tracing::{error, instrument, trace, warn};
 pub struct UdpProxyClient {
     write: UdpProxyClientWriteHalf,
     read: UdpProxyClientReadHalf,
+    upstream_addr: InternetAddr,
 }
 
 impl UdpProxyClient {
@@ -60,11 +61,15 @@ impl UdpProxyClient {
             let upstream = Arc::new(upstream);
             let write = UdpProxyClientWriteHalf::new(upstream.clone(), Vec::new().into());
             let read = UdpProxyClientReadHalf::new(upstream, Vec::new().into(), proxies);
-            return Ok(UdpProxyClient { write, read });
+            return Ok(UdpProxyClient {
+                write,
+                read,
+                upstream_addr: destination,
+            });
         }
 
         // Connect to upstream
-        let proxy_addr = &proxies[0].address;
+        let proxy_addr = proxies[0].address.clone();
         let addr =
             proxy_addr
                 .to_socket_addr()
@@ -103,12 +108,20 @@ impl UdpProxyClient {
         let header_bytes: Arc<[_]> = buf.into();
         let write = UdpProxyClientWriteHalf::new(upstream.clone(), header_bytes.clone());
         let read = UdpProxyClientReadHalf::new(upstream, header_bytes, proxies);
-        Ok(UdpProxyClient { write, read })
+        Ok(UdpProxyClient {
+            write,
+            read,
+            upstream_addr: proxy_addr,
+        })
     }
 
     pub fn into_split(self) -> (UdpProxyClientReadHalf, UdpProxyClientWriteHalf) {
         (self.read, self.write)
     }
+
+    pub fn remote_addr(&self) -> &InternetAddr {
+        &self.upstream_addr
+    }
 }
 
 #[derive(Debug, Error)]