proxy_tokio.rs

use tokio::net::{TcpListener, TcpStream};
use tokio::io::{Error, AsyncReadExt, AsyncWriteExt};
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};

#[tokio::main]
async fn main() {
    let server = TcpListener::bind("127.0.0.1:5433").await.unwrap();
    tokio::task::spawn(async move {
        loop {
            let (socket, _) = server.accept().await.unwrap();
            // socket.set_nodelay(true).unwrap();

            tokio::spawn(async move {
                handle_connection(socket).await.unwrap();
            });
        }
    }).await.unwrap();
}

async fn handle_connection(client: TcpStream) -> Result<(), Error> {
    let upstream = TcpStream::connect("127.0.0.1:5432").await?;
    // upstream.set_nodelay(true)?;

    let (mut c_rx, mut c_tx) = client.into_split();
    let (mut s_rx, mut s_tx) = upstream.into_split();

    let t1 = pipe(&mut c_rx, &mut s_tx);
    let t2 = pipe(&mut s_rx, &mut c_tx);

    let (t1, t2) = tokio::join!(t1, t2);

    t1?;
    t2?;

    Ok(())
}

async fn pipe(input: &mut OwnedReadHalf, output: &mut OwnedWriteHalf) -> Result<(), Error> {
    let mut buf = [0u8; 32_768];

    loop {
        let size = input.read(&mut buf).await?;
        if size == 0 {
            break;
        } else if size > 0 {
            if output.write(&buf[0..size]).await.is_ok() {
                output.flush().await?;
            }
        }
    }

    Ok(())
}