HTTPS proxy server in Rust using the CONNECT protocol

Question

I'm a Rust newbie. As a pet project, I decided to implement a simple multi-threaded HTTPS proxy server in Rust which uses the CONNECT protocol. I have tested the code below using my browser.

use std::io::prelude::*; // Contains the read/write traits
use std::net::{TcpListener, TcpStream};
use std::str;
use std::io;
use std::thread;

const PORT: i32 = 5000;

// Function to handle connection
fn handle_client(mut stream: TcpStream) {
    println!("[*] Received connection request from {:?}", stream);

    // Read the CONNECT request's bytes into the buffer
    let mut buf = [0; 4096];

    // Read the bytes from the stream
    let nbytes = match stream.read(&mut buf) {
        Ok(n) => n,
        Err(_) => return () // early return if an error
    }; 

    println!("[*] Received {} bytes of data", nbytes);

    // Convert the request to a string
    let req : &str = match str::from_utf8(&buf) {
        Ok(s) => s,
        Err(_) => return () // early return if an error
    }; 

    println!("[*] Received request {}", req);

    // Split the text on whitespace and get the hostname
    let website : &str = req.split_whitespace().collect::<Vec<_>>()[1];

    println!("[*] Connecting to {}", website);

    // Open a TCP connection to the website
    let mut tunnel = match TcpStream::connect(website) {
        Ok(t) => t,
        Err(_) => return ()
    };

    // Send an ack to the client
    match stream.write_all(b"HTTP/1.1 200 Connection established\r\n\r\n") {
        Ok(_) => (),
        Err(_) => return ()
    };

    // Set both sockets to nonblocking mode
    match stream.set_nonblocking(true) {
        Ok(()) => (),
        Err(_) => return ()
    };

    match tunnel.set_nonblocking(true) {
        Ok(()) => (),
        Err(_) => return ()
    }

    let mut stream_buf = [0u8; 4096]; // Buffer containing data received from stream
    let mut tunnel_buf = [0u8; 4096]; // Buffer containing data received from tunnel

    let mut stream_nbytes = 0usize; // The number of bytes pending in stream_buf to be written to tunnel
    let mut tunnel_nbytes = 0usize; // The number of bytes pending in tunnel_buf to be written to stream

    // Keep copying data back and forth
    loop {
        // Read data from stream to be sent to tunnel -- only read if stream_buf is empty
        if stream_nbytes == 0 {
            stream_nbytes = match stream.read(&mut stream_buf) {
                Ok(0) => return (), // Socket closed 
                Ok(n) => n,
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => 0, // If there is no data, return 0 bytes written
                Err(_) => return ()
            };
        }

        // Read data from tunnel to be sent to stream -- only read if tunnel_buf is empty
        if tunnel_nbytes == 0 {
            tunnel_nbytes = match tunnel.read(&mut tunnel_buf) {
                Ok(0) => return (), // Socket closed 
                Ok(n) => n,
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => 0, // If there is no data, return 0 bytes written
                Err(_) => return ()
            };
        }

        // Write data from stream to tunnel
        if stream_nbytes > 0 {
            // Pass the slice corresponding to first `stream_nbytes`
            match tunnel.write(&mut stream_buf[0..stream_nbytes]) {
                Ok(0) => return (), // Socket closed 
                Ok(n) if n == stream_nbytes => stream_nbytes = 0, // If we get equal 
                Ok(_) => { println!("Cannot write partially :("); return () }, // No support for partial nbytes
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => (), // Write the bytes in later
                Err(_) => return ()
            }
        }

        // Write data from tunnel to stream
        if tunnel_nbytes > 0 {
            // Pass the slice corresponding to first `stream_nbytes`
            match stream.write(&mut tunnel_buf[0..tunnel_nbytes]) {
                Ok(0) => return (), // Socket closed
                Ok(n) if n == tunnel_nbytes => tunnel_nbytes = 0, // If we get equal 
                Ok(_) => { println!("Cannot write partially :("); return () }, // No support for partial nbytes
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => (), // Write the bytes in later
                Err(_) => return ()
            }
        }
    }
 }

fn main() -> std::io::Result<()> {
    // Create a server
    let local_addr = format!("localhost:{}", PORT);
    let server = TcpListener::bind(local_addr)?;

    println!("[*] Listening on port {}", PORT);

    // Keep spinning and spawn threads for any incoming connections
    for stream_result in server.incoming() {
        match stream_result {
            Ok(stream) => thread::spawn(move || handle_client(stream)), // Spawn a new thread, ignore the return value because we don't need to join threads
            _          => continue
        };
    }

    Ok(())
}

I am looking for critique on the following facets of my code -

1. Is this idiomatic Rust? It feels quite verbose and convoluted when compared to implementations in other languages. I'm wondering if that is just a feature of Rust or I'm missing something. Is my use of Rust primitives (like slices, arrays, etc.) optimal? Can I refactor the code?
2. Is there a better way to handle errors?
3. Can I improve the performance of the code? Are there better alternatives than threading which are not too complicated to implement?
4. Can I write better comments and improve readability?
5. Can I improve logging?

Feel free to point out anything else that could have been done better.

Answer 1

Idiomatic Rust

Standard utilities

Much of your code is contained in the std function std::io::copy, which you can use to avoid handling each side of the pipe:

match io::copy(&mut reader, &mut writer) {
    Ok(0) => return (),
    Err(e) if e.kind() != io::ErrorKind::WouldBlock => return (),
    _ => ()
}

And you can also use a for loop to avoid repeating yourself here:

let pipes = [
    (&stream, &tunnel),
    (&tunnel, &stream)
];

loop {
    for (mut reader, mut writer) in pipes.iter() {
        // ...
    }
}

Error Handling

You have a pattern repeating in your code for handling errors:

let value = match expression {
    Ok(value) => value,
    Err(_) => return ()
};

This pattern is actually so common that Rust provides the ? operator which behaves the same way:

let values = expression?;

Okay, it's actually a little different: When using the ? operator, Rust needs to know how to convert the value to your return type. It does this using an implementation of E: Into<ReturnType::Error>, and so your function will need to return a compatible type. For small applications, using a crate like anyhow makes this easy:

use anyhow::*;

fn handle_client(...) -> Result<()> {
    // Read the CONNECT request's bytes into the buffer
    let mut buf = [0; 4096];
    let nbytes = stream.read(&mut buf)?;
}

Now that we are receiving errors from the handler, we can let the user know about them

thread::spawn(move || {
    if let Err(error) = handle_client(stream) {
        error!("error while handling stream: {}", error);
    }
});

Also, returning an error from main uses a default implementation that isn't very user friendly:

Error: Os { code: 10048, kind: AddrInUse, message: "Only one usage of each socket address (protocol/network address/port) is normally permitted." }

It might be better to provide a more descriptive error message:

let server = TcpListener::bind(("localhost", PORT)).unwrap_or_else(|e| {
    if e.kind() == io::ErrorKind::AddrInUse {
        error!("Port {} is already being used by another program", PORT);
        std::process::exit(1);
    } else {
        panic!("{:?}", e);
    }
});

Logging

The [log] crate is a pretty conventional way of writing log messages in Rust - many libraries support it directly, letting you see what's happening inside their code too. Just set up a logger using a crate like [env_logger] and output messages using its' macros.

use log::*;

fn main() {
    env_logger::init();
    // ...
    info!("Listening on port {}", PORT);
}

And in no particular order...

• TcpListener::bind takes a ToSocketAddrs. You can avoid allocating with bind(("localhost", PORT)).
• You can write a match with one arm as an if let statement
• You can get the 2nd item in an iterator with nth:

  let website = req
       .split_whitespace()
       .nth(1)
  

• Personally, I don't use top level imports for items I only use a single time, preferring to use let req = std::str::from_utf8(&buf)?;
• Speaking of that line, you might want to convert &buf[..nbytes] instead, to avoid checking the empty part of the buffer.

Comments

If you ask me, the code is easier to read with less comments - you should be describing the conceptual purpose of a block of code. "Read the bytes from the stream" doesn't tell me anything more about why you are calling stream.read.

Rather, I think this is more helpful

// Read the CONNECT request's bytes into the buffer
let mut buf = [0; 4096];
let nbytes = stream.read(&mut buf)?;

There is also no problem treating log messages as documentation - they can teach the reader about the behaviour of the code just as well.

Another example would be the request parsing:

// Split the text on whitespace and get the hostname
let website : &str = req.split_whitespace().collect::<Vec<_>>()[1];

I can already see that you are splitting the string on whitespace - it says it right there in the method name. It's more helpful to know what this line is doing. It could be "Parse the hostname from the request", but I'd be happy just using the variable name to document this line.

On the other hand, abbreviations can get in the way of a reader's understanding - try to avoid them when they describe an important part of the code:

// Send an acknowledgement to the client
stream.write_all(b"HTTP/1.1 200 Connection established\r\n\r\n")?;

Performance

Each thread is currently using a busy loop to wait for data on the sockets, and isn't gaining much in performance from it. You can make it much easier on your computer by adding a small std::thread::sleep between each iteration of the loop.

Even better though: this kind of application is particularly suited to async: The proxy will be able to handle many more clients without choking on the heavyweight threads, and wake up each task the moment data is ready on the socket. With smol and futures, you can write code that'll look nearly identical to the sync code (and actually fix a problem where packets can get dropped if the write socket is full). Altogether, these changes look like this: proxy.rs