A simple proxy with async/await syntax!

async_await_proxy.rs

#![feature(async_await)]
extern crate structopt;

use structopt::StructOpt;

use std::{
    future::Future as StdFuture,
    net::{IpAddr, SocketAddr},
};

use futures::{
    executor,
    future::{self, TryFutureExt},
    io, AsyncRead, AsyncReadExt, AsyncWrite, StreamExt,
};

use romio::{TcpListener, TcpStream};

#[derive(Debug, StructOpt)]
#[structopt(name = "SimpleProxy", about = "A simple, async proxy")]
struct Cli {
    /// Localhost, can be set to internal only.
    #[structopt(long = "localhost", parse(try_from_str), default_value = "0.0.0.0")]
    localhost: IpAddr,

    /// The port the proxy should run on.
    #[structopt(short = "l", parse(try_from_str), default_value = "8080")]
    localport: u16,

    /// The host this server is a proxy to.
    #[structopt(short = "r", long = "remotehost", parse(try_from_str))]
    remotehost: IpAddr,

    /// The port of the host this server should proxy to.
    #[structopt(short = "p", long = "remoteport", parse(try_from_str))]
    remoteport: u16,
}

fn main() -> std::io::Result<()> {
    let app = Cli::from_args();

    let server_addr = SocketAddr::new(app.localhost, app.localport);
    let host_addr = SocketAddr::new(app.remotehost, app.remoteport);

    let mut tcp = TcpListener::bind(&server_addr).unwrap();
    let mut incoming = tcp.incoming();

    executor::block_on(async {
        while let Some(stream) = incoming.next().await {
            let stream = stream?;
            juliex::spawn(async move {
                let _ = proxy(stream, &host_addr).await;
            });
        }

        Ok(())
    })
}

fn copy<'a, R: 'a, W>(reader: R, writer: &'a mut W) -> impl StdFuture<Output = io::Result<()>> + 'a
where
    R: AsyncReadExt,
    W: AsyncWrite + Unpin,
{
    reader
        .copy_into(writer)
        .map_ok(|n| println!("wrote {} bytes", n))
}

async fn proxy<R>(server: R, host: &SocketAddr) -> io::Result<()>
where
    R: AsyncRead + AsyncWrite + Send + Sync,
{
    let stream = TcpStream::connect(host).await?;

    let (server_reader, mut server_writer) = server.split();
    let (host_reader, mut host_writer) = stream.split();

    let sending = copy(server_reader, &mut host_writer);
    let receiving = copy(host_reader, &mut server_writer);

    let proxy = future::select(sending, receiving);
    let _ = proxy.await.factor_first().0?;

    Ok(())
}