From 94835c41a4f45583a631b1e21f376ca89794838f Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Thu, 10 Aug 2023 12:57:35 +0800
Subject: [PATCH] UDP Associate

---
 src/http.rs      |  2 +-
 src/socks.rs     | 26 ++++++++++++++++++++------
 src/tun2proxy.rs | 12 +++++++++---
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/src/http.rs b/src/http.rs
index b7ddfce..d3ca650 100644
--- a/src/http.rs
+++ b/src/http.rs
@@ -391,7 +391,7 @@ impl ConnectionManager for HttpManager {
         info.protocol == IpProtocol::Tcp
     }
 
-    fn new_tcp_proxy(&self, info: &ConnectionInfo) -> Result<Box<dyn TcpProxy>, Error> {
+    fn new_tcp_proxy(&self, info: &ConnectionInfo, _: bool) -> Result<Box<dyn TcpProxy>, Error> {
         if info.protocol != IpProtocol::Tcp {
             return Err("Invalid protocol".into());
         }
diff --git a/src/socks.rs b/src/socks.rs
index ec4bf39..4d01125 100644
--- a/src/socks.rs
+++ b/src/socks.rs
@@ -1,5 +1,5 @@
 use crate::{
-    error::Error,
+    error::{Error, Result},
     tun2proxy::{
         ConnectionInfo, ConnectionManager, Direction, IncomingDataEvent, IncomingDirection, OutgoingDataEvent,
         OutgoingDirection, TcpProxy,
@@ -7,7 +7,7 @@ use crate::{
 };
 use smoltcp::wire::IpProtocol;
 use socks5_impl::protocol::{self, handshake, password_method, Address, AuthMethod, StreamOperation, UserKey, Version};
-use std::{collections::VecDeque, net::SocketAddr};
+use std::{collections::VecDeque, convert::TryFrom, net::SocketAddr};
 
 #[derive(Eq, PartialEq, Debug)]
 #[allow(dead_code)]
@@ -31,10 +31,17 @@ struct SocksProxyImpl {
     data_buf: VecDeque<u8>,
     version: Version,
     credentials: Option<UserKey>,
+    command: protocol::Command,
+    udp_associate: Option<SocketAddr>,
 }
 
 impl SocksProxyImpl {
-    fn new(info: &ConnectionInfo, credentials: Option<UserKey>, version: Version) -> Result<Self, Error> {
+    fn new(
+        info: &ConnectionInfo,
+        credentials: Option<UserKey>,
+        version: Version,
+        command: protocol::Command,
+    ) -> Result<Self> {
         let mut result = Self {
             info: info.clone(),
             state: SocksState::ServerHello,
@@ -45,6 +52,8 @@ impl SocksProxyImpl {
             data_buf: VecDeque::default(),
             version,
             credentials,
+            command,
+            udp_associate: None,
         };
         result.send_client_hello()?;
         Ok(result)
@@ -195,8 +204,7 @@ impl SocksProxyImpl {
     }
 
     fn send_request_socks5(&mut self) -> Result<(), Error> {
-        protocol::Request::new(protocol::Command::Connect, self.info.dst.clone())
-            .write_to_stream(&mut self.server_outbuf)?;
+        protocol::Request::new(self.command, self.info.dst.clone()).write_to_stream(&mut self.server_outbuf)?;
         self.state = SocksState::ReceiveResponse;
         self.state_change()
     }
@@ -216,6 +224,9 @@ impl SocksProxyImpl {
         if response.reply != protocol::Reply::Succeeded {
             return Err(format!("SOCKS connection failed: {}", response.reply).into());
         }
+        if self.command == protocol::Command::UdpAssociate {
+            self.udp_associate = Some(SocketAddr::try_from(response.address)?);
+        }
 
         self.server_outbuf.append(&mut self.data_buf);
         self.data_buf.clear();
@@ -325,14 +336,17 @@ impl ConnectionManager for SocksProxyManager {
         info.protocol == IpProtocol::Tcp
     }
 
-    fn new_tcp_proxy(&self, info: &ConnectionInfo) -> Result<Box<dyn TcpProxy>, Error> {
+    fn new_tcp_proxy(&self, info: &ConnectionInfo, udp_associate: bool) -> Result<Box<dyn TcpProxy>> {
         if info.protocol != IpProtocol::Tcp {
             return Err("Invalid protocol".into());
         }
+        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,
         )?))
     }
 
diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs
index e33f43b..0cabacf 100644
--- a/src/tun2proxy.rs
+++ b/src/tun2proxy.rs
@@ -193,7 +193,7 @@ pub(crate) trait UdpProxy {
 
 pub(crate) trait ConnectionManager {
     fn handles_connection(&self, info: &ConnectionInfo) -> bool;
-    fn new_tcp_proxy(&self, info: &ConnectionInfo) -> Result<Box<dyn TcpProxy>, Error>;
+    fn new_tcp_proxy(&self, info: &ConnectionInfo, udp_associate: bool) -> Result<Box<dyn TcpProxy>, Error>;
     fn close_connection(&self, info: &ConnectionInfo);
     fn get_server_addr(&self) -> SocketAddr;
     fn get_credentials(&self) -> &Option<UserKey>;
@@ -447,7 +447,7 @@ impl<'a> TunToProxy<'a> {
                 if first_packet {
                     let mut done = false;
                     for manager in self.connection_managers.iter_mut() {
-                        let tcp_proxy_handler = manager.new_tcp_proxy(&connection_info);
+                        let tcp_proxy_handler = manager.new_tcp_proxy(&connection_info, false);
                         if tcp_proxy_handler.is_err() {
                             continue;
                         }
@@ -504,8 +504,14 @@ impl<'a> TunToProxy<'a> {
                         self.expect_smoltcp_send()?;
                         self.sockets.remove(handle);
                     }
+                } else {
+                    // Another UDP packet
+                    let cm = self.get_connection_manager(&connection_info);
+                    if cm.is_none() {
+                        return Ok(());
+                    }
+                    // TODO: Handle UDP packets
                 }
-                // Otherwise, UDP is not yet supported.
             } else {
                 log::warn!("Unsupported protocol: {} ({})", connection_info, dst);
             }