TCP proxy PoC

main.rs

// testing
// 1. open server on 127.0.0.1:7979
// $ nc -l -k 127.0.0.1 7979
// 2. connects to 127.0.0.1:7878
// $ nc 127.0.0.1 7878

// proxy is close when either side closed the connection

use std::io::prelude::*;
use std::net::{TcpListener, TcpStream, Shutdown, SocketAddr};
use std::io::{copy, ErrorKind, Error};

use anyhow::{Result};

fn main() {
    let listener = TcpListener::bind("127.0.0.1:7878").unwrap();

    let remote: SocketAddr = "127.0.0.1:7979".parse().unwrap();
    for stream in listener.incoming() {
        let stream = stream.unwrap();
        println!("connected");

        handle_connection(stream, remote);
        println!("proxy completed");
        break;
    }
}

fn handle_connection(mut source: TcpStream, destination: SocketAddr) -> Result<()> {
    // `source` and `destination` are both `TcpStream`
    let mut destination = TcpStream::connect(destination)?;

    let mut source_copy = source.try_clone()?;
    let mut destination_copy = destination.try_clone()?;
    let src2dst = std::thread::spawn(move || {
        std::io::copy(&mut source_copy, &mut destination_copy).unwrap();
        println!("src -> dst completed");
        // client has closed connection, we need to make server aware of it
        //destination_copy.shutdown(Shutdown::Both).unwrap();
        // shutdown twice might result in error on macOS
        match destination_copy.shutdown(Shutdown::Both) {
            Err(err) => match err.kind() {
                ErrorKind::NotConnected => (),
                e => println!("should not happen: {:?}", e),
            },
            Ok(_) => ()
       };
    });

    let dst2src = std::thread::spawn(move || {
        std::io::copy(&mut destination, &mut source).unwrap();
        println!("dst -> src completed");
        // server has closed connection, we need to make client aware of it
        //source.shutdown(Shutdown::Both).unwrap();
        match source.shutdown(Shutdown::Both) {
            Err(err) => match err.kind() {
                ErrorKind::NotConnected => (),
                e => println!("should not happen: {:?}", e),
            },
            Ok(_) => ()
       };
    });

    src2dst.join().unwrap();
    dst2src.join().unwrap();
    Ok(())
}