RustでTCPトンネルを実装

RustでTCPトンネルを実装してみたので備忘録的にメモ

まずサーバー側のプロセスは以下のように送られてきた内容をそのまま返すだけのエコーサーバーを用意しました。

fn tcp_echo(mut stream: TcpStream) -> Result<(), Error> {
    log::info!("connect from[{}] to[{}]", stream.peer_addr()?, stream.local_addr()?);
    let mut buffer = [0u8; 1024];
    loop {
        let nbytes = stream.read(&mut buffer)?;
        if nbytes == 0 {
            return Ok(());
        }

        // exitの文字列が送られてきたら終了
        let str = match str::from_utf8(&buffer[..nbytes]) {
            Ok(str) => str,
            Err(e) => ""
        }.trim();
        log::info!("from[{}] to[{}]: {}", stream.peer_addr()?, stream.local_addr()?, str);
        if str.eq("exit") {
            stream.shutdown(Shutdown::Both).expect("shutdown call failed");
            return Ok(());
        }

        // 来た内容をそのまま返す
        stream.write(&buffer[..nbytes])?;
        stream.flush()?;
    }
}

fn main() {
    init_logger();
    let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
    let port = 6543;
    let listener = TcpListener::bind(SocketAddr::new(ip, port)).expect("Error. failed to bind to the address");
    for streams in listener.incoming() {
        match streams {
            Err(e) => { eprintln!("error: {}", e)},
            Ok(stream) => {
                // 接続毎にスレッド生成
                thread::spawn(move || {
                    println!("{:?}", stream);
                    tcp_echo(stream).unwrap_or_else(|error| eprintln!("receive: {:?}", error));
                });
            }
        }
    }
}

試してみると以下のようにエコーサーバーとして機能していることが確認できます。

$ telnet localhost 6543
Trying 127.0.0.1...
Connected to localhost.
Escape character is '^]'.
abc
abc
def
def
ghi
ghi
exit
Connection closed by foreign host.
$

それから、トンネル用のプログラムとして以下を実装しました。

fn stream_tunnel(mut from_stream: TcpStream, mut to_stream: TcpStream) -> Result<(), Error>{
    let mut buffer = [0u8; 1024];
    loop {
        let nbytes = from_stream.read(&mut buffer)?;
        if nbytes == 0 {
            to_stream.shutdown(Shutdown::Write).expect("shutdown call failed");
            return Ok(());
        }

        let str = match str::from_utf8(&buffer[..nbytes]) {
            Ok(str) => str,
            Err(e) => ""
        }.trim();

        // 来た内容をそのまま送る
        to_stream.write(&buffer[..nbytes])?;
        to_stream.flush()?;
    }
}


fn tcp_tunnel(client_stream: TcpStream) -> Result<(), Error> {
    let clinet_stream2 = client_stream.try_clone()?;

    let server_socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 6543);
    let server_stream = TcpStream::connect_timeout(&server_socket, Duration::from_secs(1)).expect("Could not connect.");
    let server_strea2 = server_stream.try_clone()?;

    let handle1 = thread::spawn(move || {
        stream_tunnel(client_stream, server_stream);
    });
    let handle2 = thread::spawn(move || {
        stream_tunnel(server_strea2, clinet_stream2);
    });


    handle1.join().unwrap();
    handle2.join().unwrap();
    Ok(())
}

fn main() {
    let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
    let port = 3456;
    let listener = TcpListener::bind(SocketAddr::new(ip, port)).expect("Error. failed to bind to the address");
    for streams in listener.incoming() {
        match streams {
            Err(e) => { eprintln!("error: {}", e)},
            Ok(stream) => {
                // 接続毎にスレッド生成
                thread::spawn(move || {
                    println!("{:?}", stream);
                    tcp_tunnel(stream).unwrap_or_else(|error| eprintln!("receive: {:?}", error));
                });
            }
        }
    }
    
}

ここではローカルの3456ポートへの書き込みをそのまま6543ポートへ書き込むのとその逆をやっているだけです。TcpStreamのスレッド間での共有が必要だったので最初Arc<Mutex>にする必要があったのかと思ったのですが、TcpStream側にtry_cloneメソッドが用意されていたので、そちらを呼び出すだけで簡単にストリームの共有を行うことができました。