Simple Rust Proxy - Async Version

proxy_async.rs

use async_std::prelude::*;
use async_std::net::{TcpListener, TcpStream};
use async_std::task::spawn;

use futures::join;
use futures::stream::StreamExt;

#[derive(Debug)]
enum Direction {
    Back,
    Forth
}

#[async_std::main]
async fn main() {
    let server = TcpListener::bind("127.0.0.1:15432").await.unwrap();
    server.incoming().for_each_concurrent(None, |client| async move {
        println!("New conn...");
        let client = client.unwrap();
        client.set_nodelay(true).unwrap();

        spawn(handle_connection(client));
    }).await;
}

async fn handle_connection(mut client: TcpStream) {
    let mut upstream = TcpStream::connect("127.0.0.1:5432").await.unwrap();
    upstream.set_nodelay(true).unwrap();

    let mut client_clone = client.clone();
    let mut upstream_clone = upstream.clone();

    let t_forth = pipe(&mut client, &mut upstream, Direction::Forth);
    let t_back = pipe(&mut upstream_clone, &mut client_clone, Direction::Back);

    join!(t_back, t_forth);
}

async fn pipe(input: &mut TcpStream, output: &mut TcpStream, direction: Direction) {
    // println!("Direction: {:?}", direction);
    let mut buffer = [0u8; 32_768];
    loop {
        match input.read(&mut buffer).await {
            Ok(size) => {
                if size == 0 {
                    output.shutdown(std::net::Shutdown::Both).unwrap();
                    break;
                }

                /*
                let mut size = size;
                if is_forth {
                    size = cache::processor_forth(size, &mut buffer);
                }
                */

                if output.write(&buffer[..size]).await.is_ok() {
                    output.flush().await.unwrap();
                }
            },
            Err(error) => { println!("Error: {:?}", error); }
        }
    }
}