tcpproxy.rs

use async_std::{
    net::{TcpListener, TcpStream},
    task,
};
use failure::Error;
use futures::{future, prelude::*, stream::FuturesUnordered};

/// Handle incoming connections.
async fn handle(socket: TcpStream, remote: &str) -> Result<(), Error> {
    let remote = TcpStream::connect(remote).await?;
    let (socket_rx, mut socket_tx) = socket.split();
    let (rx, mut tx) = remote.split();

    let sender = socket_rx.copy_into(&mut tx);
    let receiver = rx.copy_into(&mut socket_tx);

    // NB: at the end of this, either sender or receiver will error or disconnect.
    // You can figure out what happened by looking at the result of the select.
    let _ = future::select(sender, receiver).await;

    Ok(())
}

fn main() -> Result<(), Error> {
    task::block_on(async {
        let listener = TcpListener::bind("127.0.0.1:8080").await?;
        let mut connections = FuturesUnordered::new();

        let mut incoming = listener.incoming().fuse();

        loop {
            // NB: special case required in case of no active connections.
            // See: https://github.com/rust-lang-nursery/futures-rs/issues/1765
            if connections.is_empty() {
                let stream = incoming.select_next_some().await?;
                connections.push(handle(stream, "127.0.0.1:8081").boxed());
            }

            futures::select! {
                stream = incoming.select_next_some() => {
                    let stream = stream?;
                    connections.push(handle(stream, "127.0.0.1:8082").boxed());
                }
                _ = connections.select_next_some() => {
                    println!("lost connection");
                }
            }
        }

        // NB: https://github.com/rust-lang/async-book/issues/22#issuecomment-522537735
        #[allow(unreachable_code)]
        Ok::<(), Error>(())
    })?;

    Ok(())
}