fix: properly handle errors, close connection
This commit is contained in:
parent
468e13fa6e
commit
7394f91444
4 changed files with 18 additions and 9 deletions
|
@ -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 {
|
||||
|
|
|
@ -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()))?;
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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))
|
||||
}
|
||||
}
|
||||
|
|
12
src/main.rs
12
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;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
|
Loading…
Reference in a new issue