diff --git a/tokio/examples/proxy.rs b/tokio/examples/proxy.rs new file mode 100644 index 000000000..795dda82a --- /dev/null +++ b/tokio/examples/proxy.rs @@ -0,0 +1,71 @@ +//! A proxy that forwards data to another server and forwards that server's +//! responses back to clients. +//! +//! Because the Tokio runtime uses a thread pool, each TCP connection is +//! processed concurrently with all other TCP connections across multiple +//! threads. +//! +//! You can showcase this by running this in one terminal: +//! +//! cargo run --example proxy +//! +//! This in another terminal +//! +//! cargo run --example echo +//! +//! And finally this in another terminal +//! +//! cargo run --example connect 127.0.0.1:8081 +//! +//! This final terminal will connect to our proxy, which will in turn connect to +//! the echo server, and you'll be able to see data flowing between them. + +#![deny(warnings, rust_2018_idioms)] +#![feature(async_await)] + +use futures::{future::try_join, FutureExt, StreamExt}; +use std::{env, error::Error, net::SocketAddr}; +use tokio::{ + io::AsyncReadExt, + net::{TcpListener, TcpStream}, +}; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let listen_addr = env::args().nth(1).unwrap_or("127.0.0.1:8081".to_string()); + let listen_addr = listen_addr.parse::()?; + + let server_addr = env::args().nth(2).unwrap_or("127.0.0.1:8080".to_string()); + let server_addr = server_addr.parse::()?; + + println!("Listening on: {}", listen_addr); + println!("Proxying to: {}", server_addr); + + let mut incoming = TcpListener::bind(&listen_addr)?.incoming(); + + while let Some(Ok(inbound)) = incoming.next().await { + let transfer = transfer(inbound, server_addr).map(|r| { + if let Err(e) = r { + println!("Failed to transfer; error={}", e); + } + }); + + tokio::spawn(transfer); + } + + Ok(()) +} + +async fn transfer(inbound: TcpStream, proxy_addr: SocketAddr) -> Result<(), Box> { + let outbound = TcpStream::connect(&proxy_addr).await?; + + let (mut ri, mut wi) = inbound.split(); + let (mut ro, mut wo) = outbound.split(); + + let client_to_server = ri.copy(&mut wo); + let server_to_client = ro.copy(&mut wi); + + try_join(client_to_server, server_to_client).await?; + + Ok(()) +} diff --git a/tokio/examples_old/proxy.rs b/tokio/examples_old/proxy.rs deleted file mode 100644 index 3f4e6754e..000000000 --- a/tokio/examples_old/proxy.rs +++ /dev/null @@ -1,128 +0,0 @@ -//! A proxy that forwards data to another server and forwards that server's -//! responses back to clients. -//! -//! Because the Tokio runtime uses a thread pool, each TCP connection is -//! processed concurrently with all other TCP connections across multiple -//! threads. -//! -//! You can showcase this by running this in one terminal: -//! -//! cargo run --example proxy -//! -//! This in another terminal -//! -//! cargo run --example echo -//! -//! And finally this in another terminal -//! -//! cargo run --example connect 127.0.0.1:8081 -//! -//! This final terminal will connect to our proxy, which will in turn connect to -//! the echo server, and you'll be able to see data flowing between them. - -#![deny(warnings, rust_2018_idioms)] - -use std::env; -use std::io::{self, Read, Write}; -use std::net::{Shutdown, SocketAddr}; -use std::sync::{Arc, Mutex}; -use tokio; -use tokio::io::{copy, shutdown}; -use tokio::net::{TcpListener, TcpStream}; -use tokio::prelude::*; - -fn main() -> Result<(), Box> { - let listen_addr = env::args().nth(1).unwrap_or("127.0.0.1:8081".to_string()); - let listen_addr = listen_addr.parse::()?; - - let server_addr = env::args().nth(2).unwrap_or("127.0.0.1:8080".to_string()); - let server_addr = server_addr.parse::()?; - - // Create a TCP listener which will listen for incoming connections. - let socket = TcpListener::bind(&listen_addr)?; - println!("Listening on: {}", listen_addr); - println!("Proxying to: {}", server_addr); - - let done = socket - .incoming() - .map_err(|e| println!("error accepting socket; error = {:?}", e)) - .for_each(move |client| { - let server = TcpStream::connect(&server_addr); - let amounts = server.and_then(move |server| { - // Create separate read/write handles for the TCP clients that we're - // proxying data between. Note that typically you'd use - // `AsyncRead::split` for this operation, but we want our writer - // handles to have a custom implementation of `shutdown` which - // actually calls `TcpStream::shutdown` to ensure that EOF is - // transmitted properly across the proxied connection. - // - // As a result, we wrap up our client/server manually in arcs and - // use the impls below on our custom `MyTcpStream` type. - let client_reader = MyTcpStream(Arc::new(Mutex::new(client))); - let client_writer = client_reader.clone(); - let server_reader = MyTcpStream(Arc::new(Mutex::new(server))); - let server_writer = server_reader.clone(); - - // Copy the data (in parallel) between the client and the server. - // After the copy is done we indicate to the remote side that we've - // finished by shutting down the connection. - let client_to_server = copy(client_reader, server_writer) - .and_then(|(n, _, server_writer)| shutdown(server_writer).map(move |_| n)); - - let server_to_client = copy(server_reader, client_writer) - .and_then(|(n, _, client_writer)| shutdown(client_writer).map(move |_| n)); - - client_to_server.join(server_to_client) - }); - - let msg = amounts - .map(move |(from_client, from_server)| { - println!( - "client wrote {} bytes and received {} bytes", - from_client, from_server - ); - }) - .map_err(|e| { - // Don't panic. Maybe the client just disconnected too soon. - println!("error: {}", e); - }); - - tokio::spawn(msg); - - Ok(()) - }); - - tokio::run(done); - Ok(()) -} - -// This is a custom type used to have a custom implementation of the -// `AsyncWrite::shutdown` method which actually calls `TcpStream::shutdown` to -// notify the remote end that we're done writing. -#[derive(Clone)] -struct MyTcpStream(Arc>); - -impl Read for MyTcpStream { - fn read(&mut self, buf: &mut [u8]) -> io::Result { - self.0.lock().unwrap().read(buf) - } -} - -impl Write for MyTcpStream { - fn write(&mut self, buf: &[u8]) -> io::Result { - self.0.lock().unwrap().write(buf) - } - - fn flush(&mut self) -> io::Result<()> { - Ok(()) - } -} - -impl AsyncRead for MyTcpStream {} - -impl AsyncWrite for MyTcpStream { - fn shutdown(&mut self) -> Poll<(), io::Error> { - self.0.lock().unwrap().shutdown(Shutdown::Write)?; - Ok(().into()) - } -}