fix: properly handle errors, close connection

src/error.rs

@@ -1,6 +1,8 @@
 use socks5_server::proto::Error as SocksProtoError;
 use tokio::net::TcpStream;
 
+use crate::serde_addr::TargetAddr;
+
 pub type AppResult<T> = core::result::Result<T, AppError>;
 
 #[derive(Debug)]
@@ -9,7 +11,7 @@ pub enum AppError {
     InvalidConfig(toml::de::Error),
     SocksError(SocksProtoError),
     DataNotUtf8(std::string::FromUtf8Error),
-    ThirdPartyHost,
+    ThirdPartyHost(TargetAddr),
 }
 
 impl From<std::io::Error> for AppError {
@@ -44,8 +46,8 @@ pub type HandlerResult<T> = core::result::Result<T, HandlerError>;
 
 #[derive(Debug)]
 pub struct HandlerError {
-    inner: AppError,
-    stream: TcpStream,
+    pub inner: AppError,
+    pub stream: TcpStream,
 }
 
 impl From<(AppError, TcpStream)> for HandlerError {

src/handler.rs

@@ -82,8 +82,11 @@ pub async fn handler(
             let mut conn = cmd.reply(Reply::Succeeded, Address::unspecified()).await?;
 
             let res = tokio::io::copy_bidirectional(&mut target, &mut conn).await;
+            let _ = conn.shutdown().await;
             let _ = target.shutdown().await;
 
+            // println!("2nd sd: {:?}", conn.shutdown().await);
+
             res.map_err(|e| (e, conn.into_inner()))?;
         }
     }

src/host.rs

@@ -18,7 +18,7 @@ pub fn resolve_domain(host: Vec<u8>, port: u16, config: Arc<Config>) -> AppResul
     } else if let Some(remote) = config.remote.get(&target) {
         Ok(remote.to_socket_addrs()?)
     } else {
-        Err(AppError::ThirdPartyHost)
+        Err(AppError::ThirdPartyHost(target))
     }
 }
 
@@ -30,6 +30,6 @@ pub fn resolve_ip(addr: SocketAddr, config: Arc<Config>) -> AppResult<SocketAddr
     } else if let Some(remote) = config.remote.get(&target) {
         Ok(remote.to_socket_addrs()?.next().unwrap())
     } else {
-        Err(AppError::ThirdPartyHost)
+        Err(AppError::ThirdPartyHost(target))
     }
 }

src/main.rs

@@ -6,11 +6,11 @@ mod serde_addr;
 
 use std::sync::Arc;
 
-use error::AppResult;
+use error::{AppResult, HandlerError};
 use handler::handler;
 
 use socks5_server as socks;
-use tokio::net::TcpListener;
+use tokio::{io::AsyncWriteExt, net::TcpListener};
 
 #[tokio::main]
 async fn main() -> AppResult<()> {
@@ -25,8 +25,12 @@ async fn main() -> AppResult<()> {
         tokio::spawn(async move {
             match handler(conn, config).await {
                 Ok(()) => {}
-                Err(e) => {
-                    //
+                Err(HandlerError {
+                    inner: e,
+                    stream: mut conn,
+                }) => {
+                    eprintln!("{:?}", e);
+                    let _ = conn.shutdown().await;
                 }
             }
         });