How to write correct proxy?

I tried creating it using this link and this example.
here's my code.

use clap::{Arg, Command};
use tokio::{io::{self, AsyncWriteExt}, net::{TcpListener, TcpStream}};

async fn proxy(client: &str, server: &str) -> io::Result<()> {
    let listener = TcpListener::bind(client).await?;

    while let Ok((inbound, _)) = listener.accept().await {
        

        transfer(inbound, server).await?;
    }

    Ok(())
}

async fn transfer(mut inbound: TcpStream, proxy_addr: &str) -> io::Result<()> {
    let mut outbound = TcpStream::connect(proxy_addr).await?;

    let (mut ri, mut wi) = inbound.split();
    let (mut ro, mut wo) = outbound.split();

    let client_to_server =  async {
        io::copy(&mut ri, &mut wo).await?;
        wo.shutdown().await
    };

    let server_to_client = async {
        io::copy(&mut ro, &mut wi).await?;
        wi.shutdown().await
    };

    let (a, b) = tokio::join!(client_to_server, server_to_client);
    a.unwrap();
    b.unwrap();

    Ok(())
}

#[tokio::main]
async fn main() -> io::Result<()> {
    let matches = Command::new("proxy")
        .arg(
            Arg::new("client")
                .short('c')
                .long("client")
                .value_name("ADDRESS")
                .help("The address of the eyeball that we will be proxying traffic for")
                .takes_value(true)
                .required(true),
        )
        .arg(
            Arg::new("server")
                .short('s')
                .long("server")
                .value_name("ADDRESS")
                .help("The address of the origin that we will be proxying traffic for")
                .takes_value(true)
                .required(true),
        ).get_matches();

    let client = matches.get_one::<String>("client").unwrap();
    let server = matches.get_one::<String>("server").unwrap();

    proxy(client, server).await
}

This code correctly if and only if the client has to send and recieve data and drop the connection. For example it works perfectly fine for HTTP. But, for other protocol such as websockets where the connection has to be kept alive it doesn't work.
My aim is to make this tcp proxy work for all the protocols that use tcp under the hood. How to write this correctly?

I would assume you don't call shutdown in that case.

On a separate note, using io::copy is probably going to be super inefficient. You should explore io_uring and see if it suits your purpose.

1 Like

I think the issue is you also have to forward the TCP settings like nodelay, ttl, and, most relevantly here, keepalive, among others. These are available on TcpStream as getters and setters.

3 Likes