home

A TCP Proxy in 30 lines of Rust

Jul 2021

TCP proxies accept connections from one computer and forward them to another. AWS Global Accelerator and Cloudflare Spectrum are examples of TCP proxies you can pay to use in the wild.

TCP proxies can be quite simple. Here we’ll write one in 30 some lines of Rust. After we write it we’ll demonstrate the proxy’s functionality by proxying traffic from localhost:1212 to localhost:1313.

Get to the code already

We’ll use two libraries to implement our proxy: tokio, and clap. After running cargo new tcp-proxy to create a new Rust project we can add those as dependencies by adding the following two lines to Cargo.toml.

[dependencies]
+ tokio = { version = "1", features = ["full"] }
+ clap = "2.33.3"

We’ll be using clap for parsing command line arguments and tokio for implementing the actual proxy.

Command line arguments

In src/main.rs we’ll start by adding our command line arguments. Note that we’re using #[tokio::main] to make our main function async as our actual proxy will use it.

use clap::{App, Arg};
use tokio::io;

#[tokio::main]
async fn main() -> io::Result<()> {
    let matches = App::new("proxy")
        .version("0.1")
        .author("Zeke M. <[email protected]>")
        .about("A simple tcp proxy")
        .arg(
            Arg::with_name("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::with_name("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.value_of("client").unwrap();
    let server = matches.value_of("server").unwrap();

    println!("(client, server) = ({}, {})", client, server);

    Ok(())
}

Now we can test that our argument parsing is working as expected.

$ cargo run -- -c foo -s bar
   Compiling proxy v0.1.0 (/Users/zekemedley/Desktop/projects/netsys/proxy)
    Finished dev [unoptimized + debuginfo] target(s) in 0.67s
     Running `target/debug/proxy -c foo -s bar`
(client, server) = (foo, bar)

Nice.

The proxy

Our proxy will listen for incoming tcp connections from the client and pass them along to the server. Once the client has initiated a connection it makes a new connection with the server and copies bytes to and from it.

We’ll start by making those connections.

async fn proxy(client: &str, server: &str) -> io::Result<()> {
    // Listen for connections from the eyeball and forward to the
    // origin.

    let listener = TcpListener::bind(client).await?;
    loop {
        let (eyeball, _) = listener.accept().await?;
        let origin = TcpStream::connect(server).await?;
    }
}

Once we have the connections we’ll copy bytes from the read end of each connection to the write end of the other one. We can do this by splitting the connections into their read and write ends and spawning some tokio threads to do the copying.

        let (mut eread, mut ewrite) = eyeball.into_split();
        let (mut oread, mut owrite) = origin.into_split();

        let e2o = tokio::spawn(async move { io::copy(&mut eread, &mut owrite).await });
        let o2e = tokio::spawn(async move { io::copy(&mut oread, &mut ewrite).await });

Finally if either end of the connection closes we want to close the other one. We can do that using tokio’s select macro.

        select! {
                _ = e2o => println!("e2o done"),
                _ = o2e => println!("o2e done"),

        }

The final code for our proxy function looks like this:

async fn proxy(client: &str, server: &str) -> io::Result<()> {
    let listener = TcpListener::bind(client).await?;
    loop {
        let (eyeball, _) = listener.accept().await?;
        let origin = TcpStream::connect(server).await?;

        let (mut eread, mut ewrite) = eyeball.into_split();
        let (mut oread, mut owrite) = origin.into_split();

        let e2o = tokio::spawn(async move { io::copy(&mut eread, &mut owrite).await });
        let o2e = tokio::spawn(async move { io::copy(&mut oread, &mut ewrite).await });

        select! {
                _ = e2o => println!("e2o done"),
                _ = o2e => println!("o2e done"),

        }
    }
}

We can then replace the call to println with a call to proxy in our main function.

-    println!("(client, server) = ({}, {})", client, server);
+    proxy(client, server).await

Kicking the tires

To test that our proxy works we’ll start by running it and proxying connections from 127.0.0.1:1212 to 127.0.0.1:1313.

cargo run -- -c 0.0.0.0:1212 -s 127.0.0.1:1313

Then in another window we’ll use netcat to listen for tcp connections on port 1313. This is the server end of our proxy.

nc -l 1313

And finally we’ll make a connection to 127.0.0.1:1212 and verify that our server sees the traffic.

echo hello | nc 127.0.0.1 1212

If everything was set up properly you’ll see your server print “hello”.

Congrats. You just implemented a tcp proxy in Rust. The complete code can be found here.