example: use copy_bidirectional in proxy.rs (#5856)

This commit is contained in:
kiron1
2023-07-28 18:06:44 +08:00
committed by GitHub
parent 5128601898
commit 6aca07bee7

View File

@@ -22,8 +22,7 @@
#![warn(rust_2018_idioms)]
use tokio::io;
use tokio::io::AsyncWriteExt;
use tokio::io::copy_bidirectional;
use tokio::net::{TcpListener, TcpStream};
use futures::FutureExt;
@@ -44,36 +43,19 @@ async fn main() -> Result<(), Box<dyn Error>> {
let listener = TcpListener::bind(listen_addr).await?;
while let Ok((inbound, _)) = listener.accept().await {
let transfer = transfer(inbound, server_addr.clone()).map(|r| {
if let Err(e) = r {
println!("Failed to transfer; error={}", e);
}
});
while let Ok((mut inbound, _)) = listener.accept().await {
let mut outbound = TcpStream::connect(server_addr.clone()).await?;
tokio::spawn(transfer);
tokio::spawn(async move {
copy_bidirectional(&mut inbound, &mut outbound)
.map(|r| {
if let Err(e) = r {
println!("Failed to transfer; error={}", e);
}
})
.await
});
}
Ok(())
}
async fn transfer(mut inbound: TcpStream, proxy_addr: String) -> Result<(), Box<dyn Error>> {
let mut outbound = TcpStream::connect(proxy_addr).await?;
let (mut ri, mut wi) = inbound.split();
let (mut ro, mut wo) = outbound.split();
let client_to_server = async {
io::copy(&mut ri, &mut wo).await?;
wo.shutdown().await
};
let server_to_client = async {
io::copy(&mut ro, &mut wi).await?;
wi.shutdown().await
};
tokio::try_join!(client_to_server, server_to_client)?;
Ok(())
}