Writing a TCP proxy in Rust

by Anirudh on 31 Jul '21


I've been tinkering around with Rust for a little bit now. My experience is mainly with web-dev, but I wanted to spend some time working with systems programming languages and Rust seemed to be a language gaining in popularity, so I decided to dive in and start building stuff with it.


Interested in more Rust or Distributed Systems-related blogs? Check out my Gossip Glomers series!

I spend a chunk of my time building tools and modifying/improving systems for an old MMORPG known as Tales of Pirates. So, I decided to take this opportunity to work on something that I've been meaning to do for a while, and learn Rust along the way.


Skip to the code


Problem


The intention of this proxy was to create a "middleman" that could parse data coming in from game clients, sanitize/modify it for the sake of security, and build additional game features without having to modify too much of the server-side source code of the game, since the server-side code is really old, and making significant changes often lead to unexpected bugs, data inconsistencies or straight up crashes.


The architecture would be something like this:


Base Architecture

This would also play the role of a sort of "VPN", letting players from all over the world connect to proxies that are close to them, leading to better network connectivity and lower effective response times.


Solution


Creating a proxy that just relays data back and forth between two tcp connections is fairly simple.


I'm using the tokio framework to act as a runtime in order to make async programming in Rust easier.

use tokio::{net::{TcpListener, TcpStream}, try_join};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>>{
    let proxy_server = TcpListener::bind("127.0.0.1:1973").await?;

    while let Ok((client, _)) = proxy_server.accept().await {
        tokio::spawn(async move {
            handle_client_conn(client).await;
        });
    }

    Ok(())
}

async fn handle_client_conn(mut client_conn: TcpStream) -> Result<(), Box<dyn std::error::Error>> {
    let mut main_server_conn = TcpStream::connect("127.0.0.1:1974").await?;
    let (mut client_recv, mut client_send) = client_conn.split();
    let (mut server_recv, mut server_send) = main_server_conn.split();

    let handle_one = async {
        tokio::io::copy(&mut server_recv, &mut client_send).await
    };

    let handle_two = async {
        tokio::io::copy(&mut client_recv, &mut server_send).await
    };

    try_join!(handle_one, handle_two)?;

    Ok(())
}

Let's quickly go through what's happening here. I won't dive too much into the tokio-specific code.


We basically go through the following steps:


  1. Start listening for TCP connections on port 1973
  2. Whenever we encounter a connection from a client, we create a new tokio task, and execute our handle_client_conn function in it
  3. In our connection handler, we establish a new connection with the "server".
  4. We split both our client and server connections into "Read" streams and "Write" streams
  5. We use the io::copy function to send data back and forth between the client and the server.

io::copy

The io::copy function copies bytes from a given "read" stream into a given "write" stream. This process continues until the read stream reaches EOF (no more data left) and the write stream has flushed all of the data that it was given.


Fairly straightforward, but it doesn't allow us a lot of control on the data that is being transmitted between the client and the server.


Buffers, Threads and Channels


As data flows in from either one of the connections (server or client), we need a way to be able to store, parse, sanitize and modify that data.


We can stream data from a TCP socket directly into a buffer. For this, I chose to use the bytes crate in Rust. It provides a performant byte buffer, and maintains internal cursors which are advanced whenever data is read or written to the buffer.


Channels allow us to communicate data between threads. Rust provides us with a multiple producer, single-consumer channel. Tokio has their own implementations of different types of channels, built to be used with the await syntax in Rust.


Here's a small snippet of the code I wrote which uses channels and buffers to send data to a ProxyServer instance which handles parsing this data, sanitizing it and then forwarding it to the other party (client or server, depending on where the data came from).


tokio::spawn(async move {
    let mut buffer = BytesMut::with_capacity(65536);

    loop {
        match reader.read_buf(&mut buffer).await {
            Ok(0) => {
                println!("Gateserver connection closed");
                break;
            }

            Ok(n) => loop {
                let mut cursor = Cursor::new(&buffer[..]);
                match Packet::check_frame(&mut cursor) {
                    Ok(len) => {
                        let size_to_parse = len;

                        let packet =
                            Packet::parse_frame(&mut buffer, size_to_parse).unwrap();
                        channel.send(packet).await;
                        buffer.advance(size_to_parse);
                    }

                    Err(err) => {
                        break;
                    }
                }
            },

            Err(ref e) => {
                if e.kind() == io::ErrorKind::WouldBlock {
                    continue;
                }
            }
        }
    }
});

There's a lot going on here, but let's break it down into generic terms.


  1. We create a buffer, and loop over data coming in from the reader, which is the "read" half of the TCP connection that we established with the server.
  2. This gets all data coming in from the main game server, and puts it into the buffer, returning the number of bytes that were read.
  3. We check if the data that is present in the buffer so far can be used to create a whole "Packet", which is a unit of data being used to represent an operation within the game (moving your character, attacking a monster etc.).
  4. If we have enough data for a Packet, we parse that packet, and send it across the channel to the ProxyServer, which handles the sanitizing and then forwards the data to the game client.
  5. Then, we advance the buffer, which advances the internal cursor of the buffer, so that data is written to/read from it from the correct position.
  6. If the read_buf function ever returns a 0, that means the connection has been closed from the server's side and there will be no more data coming in through the stream.

If you'd like to read more about how the data is flowing through this system, I've written some notes about it here


You can find the entire code to the work I've done here


In the next post, I'll go into more detail about the exact packet structure that the game uses, how I'm parsing it and what I intend to do with it.


Thanks for reading!