diff --git a/src/error.rs b/src/error.rs index a175f45..a80e263 100644 --- a/src/error.rs +++ b/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 = core::result::Result; #[derive(Debug)] @@ -9,7 +11,7 @@ pub enum AppError { InvalidConfig(toml::de::Error), SocksError(SocksProtoError), DataNotUtf8(std::string::FromUtf8Error), - ThirdPartyHost, + ThirdPartyHost(TargetAddr), } impl From for AppError { @@ -44,8 +46,8 @@ pub type HandlerResult = core::result::Result; #[derive(Debug)] pub struct HandlerError { - inner: AppError, - stream: TcpStream, + pub inner: AppError, + pub stream: TcpStream, } impl From<(AppError, TcpStream)> for HandlerError { diff --git a/src/handler.rs b/src/handler.rs index 88fda29..9dd8fd2 100644 --- a/src/handler.rs +++ b/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()))?; } } diff --git a/src/host.rs b/src/host.rs index 57ecafc..be73c89 100644 --- a/src/host.rs +++ b/src/host.rs @@ -18,7 +18,7 @@ pub fn resolve_domain(host: Vec, port: u16, config: Arc) -> 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) -> AppResult 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; } } });