fix: properly handle errors, close connection

This commit is contained in:
DarkCat09 2024-10-18 22:04:46 +04:00
parent 468e13fa6e
commit 7394f91444
Signed by: DarkCat09
GPG key ID: BD3CE9B65916CD82
4 changed files with 18 additions and 9 deletions

View file

@ -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 {

View file

@ -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()))?;
}
}

View file

@ -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))
}
}

View file

@ -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;
}
}
});