From b8fbc275d1e4613181f7b153fe25a4691b1b1b23 Mon Sep 17 00:00:00 2001
From: b23r0 <b23r0@foxmail.com>
Date: Sun, 10 Apr 2022 11:19:20 +0800
Subject: [PATCH] async-std to tokio.

---
 Cargo.toml   |  5 ++---
 README.md    |  2 +-
 src/main.rs  | 35 ++++++++++-------------------------
 src/socks.rs | 24 +++++-------------------
 4 files changed, 18 insertions(+), 48 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index dd3f997..29f51e2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "rsocx"
-version = "0.1.3"
+version = "0.2.0"
 authors = ["b23r0 <b23r0@foxmail.com>"]
 edition = "2021"
 description = "A high performence Socks5 proxy server with bind/reverse support implementation by Rust."
@@ -13,6 +13,5 @@ license = "MIT"
 [dependencies]
 simple_logger = "2.1.0"
 log = "0.4.14"
-futures = "0.3.0"
-async-std = {version = "1" , features = ["attributes"]}
+tokio = { version = "1", features = ["full"] }
 getopts = "0.2"
\ No newline at end of file
diff --git a/README.md b/README.md
index aae6027..10b45bc 100644
--- a/README.md
+++ b/README.md
@@ -3,7 +3,7 @@ A high performence Socks5 proxy server with bind/reverse support
 
 # Features
 
-* Async-std
+* Async
 * No unsafe code
 * Single executable
 * Linux/Windows/Mac/BSD support
diff --git a/src/main.rs b/src/main.rs
index 2e7a018..80f59d2 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,9 +4,7 @@ mod socks;
 use log::LevelFilter;
 use simple_logger::SimpleLogger;
 use getopts::Options;
-use futures::{AsyncReadExt, AsyncWriteExt, FutureExt, StreamExt};
-use async_std::{io, net::{TcpListener, TcpStream}, task};
-use futures::select;
+use tokio::{io::{self, AsyncWriteExt, AsyncReadExt}, task, net::{TcpListener, TcpStream}};
 use utils::MAGIC_FLAG;
 
 fn usage(program: &str, opts: &Options) {
@@ -18,7 +16,7 @@ fn usage(program: &str, opts: &Options) {
 }
 
 
-#[async_std::main]
+#[tokio::main]
 async fn main() -> io::Result<()>  {
 	SimpleLogger::new().with_utc_timestamps().with_utc_timestamps().with_colors(true).init().unwrap();
 	::log::set_max_level(LevelFilter::Info);
@@ -77,11 +75,9 @@ async fn main() -> io::Result<()>  {
 			Ok(p) => p
 		};
 
-		let mut incoming = listener.incoming();
-
-		while let Some(stream) = incoming.next().await {
-			let stream = stream?;
-			task::spawn(async {
+		loop{
+			let (stream , _) = listener.accept().await.unwrap();
+			tokio::spawn(async {
 				socks::socksv5_handle(stream).await;
 			});
 		}
@@ -143,10 +139,8 @@ async fn main() -> io::Result<()>  {
 			Ok(p) => p
 		};
 
-		let mut incoming = listener.incoming();
-
-		while let Some(stream) = incoming.next().await {
-			let mut stream = stream?;
+		loop {
+			let (mut stream , _) = listener.accept().await.unwrap();
 
 			match slave_stream.write_all(&[MAGIC_FLAG[0]]).await{
 				Err(e) => {
@@ -171,8 +165,8 @@ async fn main() -> io::Result<()>  {
 				let mut buf2 = [0u8 ; 1024];
 
 				loop{
-					select! {
-						a = proxy_stream.read(&mut buf1).fuse() => {
+					tokio::select! {
+						a = proxy_stream.read(&mut buf1) => {
 		
 							let len = match a {
 								Err(_) => {
@@ -191,7 +185,7 @@ async fn main() -> io::Result<()>  {
 								break;
 							}
 						},
-						b = stream.read(&mut buf2).fuse() =>  { 
+						b = stream.read(&mut buf2) =>  { 
 							let len = match b{
 								Err(_) => {
 									break;
@@ -208,17 +202,8 @@ async fn main() -> io::Result<()>  {
 								break;
 							}
 						},
-						complete => break,
 					}
 				}
-				match stream.shutdown(std::net::Shutdown::Both){
-					Err(_) => {},
-					_ => {}
-				};
-				match proxy_stream.shutdown(std::net::Shutdown::Both){
-					Err(_) => {},
-					_ => {}
-				};
 				log::info!("transfer [{}:{}] finished" , slave_addr.ip() , slave_addr.port());
 			});
 		}
diff --git a/src/socks.rs b/src/socks.rs
index ca3a3dc..124cb0d 100644
--- a/src/socks.rs
+++ b/src/socks.rs
@@ -1,7 +1,6 @@
-use futures::{AsyncReadExt, AsyncWriteExt, FutureExt};
-use async_std::{net::{TcpStream}};
 use std::{ net::{Ipv6Addr, SocketAddrV6}};
-use futures::select;
+
+use tokio::{net::TcpStream, io::{AsyncWriteExt, AsyncReadExt}};
 
 use crate::utils::makeword;
 
@@ -100,8 +99,8 @@ async fn tcp_transfer(stream : &mut TcpStream , addr : &Addr, address : &String
 	let mut buf1 = [0u8 ; 1024];
 	let mut buf2 = [0u8 ; 1024];
 	loop{
-		select! {
-			a = client.read(&mut buf1).fuse() => {
+		tokio::select! {
+			a = client.read(&mut buf1) => {
 
 				let len = match a {
 					Err(_) => {
@@ -120,7 +119,7 @@ async fn tcp_transfer(stream : &mut TcpStream , addr : &Addr, address : &String
 					break;
 				}
 			},
-			b = stream.read(&mut buf2).fuse() =>  { 
+			b = stream.read(&mut buf2) =>  { 
 				let len = match b{
 					Err(_) => {
 						break;
@@ -137,15 +136,8 @@ async fn tcp_transfer(stream : &mut TcpStream , addr : &Addr, address : &String
 					break;
 				}
 			},
-			complete => break,
 		}
 	}
-	match client.shutdown(std::net::Shutdown::Both){
-		Err(e) => {
-			log::info!("error : {}" , e);
-		},
-		_ => {}
-	};
 }
 
 pub async fn socksv5_handle(mut stream: TcpStream) {
@@ -276,10 +268,4 @@ pub async fn socksv5_handle(mut stream: TcpStream) {
 		log::info!("connection [{}] finished" , address);
 		break;
 	}
-	match stream.shutdown(std::net::Shutdown::Both){
-		Err(e) => {
-			log::error!("error : {}" , e);
-		},
-		_ => {}
-	};
 }
\ No newline at end of file