Skip to content

Commit

Permalink
example: use copy_bidirectional in proxy.rs (#5856)
Browse files Browse the repository at this point in the history
  • Loading branch information
kiron1 authored Jul 28, 2023
1 parent 5128601 commit 6aca07b
Showing 1 changed file with 12 additions and 30 deletions.
42 changes: 12 additions & 30 deletions examples/proxy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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(async move {
copy_bidirectional(&mut inbound, &mut outbound)
.map(|r| {
if let Err(e) = r {
println!("Failed to transfer; error={}", e);
}
})
.await
});

tokio::spawn(transfer);
}

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

0 comments on commit 6aca07b

Please sign in to comment.