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 socks5_server::proto::Error as SocksProtoError;
|
||||||
use tokio::net::TcpStream;
|
use tokio::net::TcpStream;
|
||||||
|
|
||||||
|
use crate::serde_addr::TargetAddr;
|
||||||
|
|
||||||
pub type AppResult<T> = core::result::Result<T, AppError>;
|
pub type AppResult<T> = core::result::Result<T, AppError>;
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
|
@ -9,7 +11,7 @@ pub enum AppError {
|
||||||
InvalidConfig(toml::de::Error),
|
InvalidConfig(toml::de::Error),
|
||||||
SocksError(SocksProtoError),
|
SocksError(SocksProtoError),
|
||||||
DataNotUtf8(std::string::FromUtf8Error),
|
DataNotUtf8(std::string::FromUtf8Error),
|
||||||
ThirdPartyHost,
|
ThirdPartyHost(TargetAddr),
|
||||||
}
|
}
|
||||||
|
|
||||||
impl From<std::io::Error> for AppError {
|
impl From<std::io::Error> for AppError {
|
||||||
|
@ -44,8 +46,8 @@ pub type HandlerResult<T> = core::result::Result<T, HandlerError>;
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
pub struct HandlerError {
|
pub struct HandlerError {
|
||||||
inner: AppError,
|
pub inner: AppError,
|
||||||
stream: TcpStream,
|
pub stream: TcpStream,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl From<(AppError, TcpStream)> for HandlerError {
|
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 mut conn = cmd.reply(Reply::Succeeded, Address::unspecified()).await?;
|
||||||
|
|
||||||
let res = tokio::io::copy_bidirectional(&mut target, &mut conn).await;
|
let res = tokio::io::copy_bidirectional(&mut target, &mut conn).await;
|
||||||
|
let _ = conn.shutdown().await;
|
||||||
let _ = target.shutdown().await;
|
let _ = target.shutdown().await;
|
||||||
|
|
||||||
|
// println!("2nd sd: {:?}", conn.shutdown().await);
|
||||||
|
|
||||||
res.map_err(|e| (e, conn.into_inner()))?;
|
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) {
|
} else if let Some(remote) = config.remote.get(&target) {
|
||||||
Ok(remote.to_socket_addrs()?)
|
Ok(remote.to_socket_addrs()?)
|
||||||
} else {
|
} 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) {
|
} else if let Some(remote) = config.remote.get(&target) {
|
||||||
Ok(remote.to_socket_addrs()?.next().unwrap())
|
Ok(remote.to_socket_addrs()?.next().unwrap())
|
||||||
} else {
|
} 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 std::sync::Arc;
|
||||||
|
|
||||||
use error::AppResult;
|
use error::{AppResult, HandlerError};
|
||||||
use handler::handler;
|
use handler::handler;
|
||||||
|
|
||||||
use socks5_server as socks;
|
use socks5_server as socks;
|
||||||
use tokio::net::TcpListener;
|
use tokio::{io::AsyncWriteExt, net::TcpListener};
|
||||||
|
|
||||||
#[tokio::main]
|
#[tokio::main]
|
||||||
async fn main() -> AppResult<()> {
|
async fn main() -> AppResult<()> {
|
||||||
|
@ -25,8 +25,12 @@ async fn main() -> AppResult<()> {
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
match handler(conn, config).await {
|
match handler(conn, config).await {
|
||||||
Ok(()) => {}
|
Ok(()) => {}
|
||||||
Err(e) => {
|
Err(HandlerError {
|
||||||
//
|
inner: e,
|
||||||
|
stream: mut conn,
|
||||||
|
}) => {
|
||||||
|
eprintln!("{:?}", e);
|
||||||
|
let _ = conn.shutdown().await;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
Loading…
Reference in a new issue