tokioのcopy_bidirectionalによる非同期TCPバイト列転送とhalf-closeの伝播

TCPコネクションを別のTCPコネクションへそのまま中継する処理をRust+tokioで実装します。HTTPなどL7のプロトコルは一切パースせず、TCPのバイト列をそのまま転送するだけの処理です。非同期ランタイムにはtokioを採用しています。

tokioの依存設定

非同期ランタイムにはtokioを使い、Cargo.tomlでは次の機能を有効にしています。

[dependencies]
tokio = { version = "1", features = ["rt-multi-thread", "net", "io-util", "macros"] }
  • rt-multi-thread: マルチスレッドの非同期ランタイム本体
  • net: TcpListener / TcpStream
  • io-util: copy_bidirectionalなどのI/Oユーティリティ
  • macros: #[tokio::main]

tokioはfullfeatureを指定すればファイルI/O・シグナルハンドリング・同期プリミティブなどをまとめて有効化できますが、ここでは実際に使う機能だけを個別に指定しています。tokioの各featureは#[cfg(feature = "...")]でコード自体をコンパイル対象から外し入れする仕組みになっており、有効化しなければ関連するモジュールも依存も一切ビルドに含まれません。

リスナーとバックエンドへの転送

固定のリッスンアドレスから固定の単一バックエンドへ転送するだけの、最小構成のコードです。

use std::net::SocketAddr;

use tokio::io::copy_bidirectional;
use tokio::net::{TcpListener, TcpStream};

const LISTEN_ADDR: &str = "127.0.0.1:8000";
const BACKEND_ADDR: &str = "127.0.0.1:9000";

#[tokio::main]
async fn main() -> std::io::Result<()> {
    let listener = TcpListener::bind(LISTEN_ADDR).await?;
    println!("listening on {LISTEN_ADDR}, forwarding to {BACKEND_ADDR}");

    loop {
        let (inbound, peer_addr) = listener.accept().await?;
        tokio::spawn(async move {
            if let Err(e) = handle_connection(inbound, peer_addr).await {
                eprintln!("connection error ({peer_addr}): {e}");
            }
        });
    }
}

async fn handle_connection(mut inbound: TcpStream, peer_addr: SocketAddr) -> std::io::Result<()> {
    let mut outbound = TcpStream::connect(BACKEND_ADDR).await?;
    let (from_client, from_backend) = copy_bidirectional(&mut inbound, &mut outbound).await?;
    println!("{peer_addr}: {from_client} bytes client->backend, {from_backend} bytes backend->client");
    Ok(())
}

accept()のエラーとコネクション個別のエラーとで、扱いを意図的に分けています。listener.accept()が失敗するのはリスナー自体に問題がある状況なので?でそのままプロセスを終了させ、一方で個々のコネクション処理(バックエンドへの接続失敗や転送中のI/Oエラー)はhandle_connectionの中で完結させ、失敗してもeprintln!でログを出すだけで全体は動き続けます。障害の影響範囲をコネクション単位に閉じ込める設計です。

copy_bidirectionalはhalf-closeを個別に伝播する

tokio::io::copy_bidirectionalは、2つのAsyncRead + AsyncWriteストリーム間でバイト列を双方向にコピーする関数です。単純にクライアント→バックエンド、バックエンド→クライアントの2つのcopyjoin!しているだけに見えますが、実際にはTCPの片方向クローズ(half-close)を正しく扱う設計になっています。

一方向の読み取りがEOFに達すると、copy_bidirectionalはその時点でもう一方の書き込み先に対してshutdown()を呼び、FINを送出します。しかし逆方向のコピーはそこで打ち切られず、そちら側が独自にEOFに達するまで動き続けます。つまり「送信を終えた(shutdown(write)した)が、相手からの応答はまだ受け取りたい」というhalf-closeを伴う通信でも、両方向が正しく完了するまで待ち続けます。仮に片方のEOFで両方向を即座に打ち切る実装だった場合、相手が応答を送り切る前にコネクションが切断されてしまいます。TCPのコネクション状態だけを見てバイト列を中継する立場では、この片方向クローズの伝播はプロトコルを問わず必要になる基本的な性質です。

動作確認

バックエンド役にpython3 -m http.server、クライアント役にcurlを使って透過性を確認します。

# バックエンド
python3 -m http.server 9000

# 転送プログラム
cargo run

# クライアント(別ターミナル)
curl -v http://127.0.0.1:8000/

ログには次のように転送バイト数が出力されます。

listening on 127.0.0.1:8000, forwarding to 127.0.0.1:9000
127.0.0.1:54608: 824 bytes client->backend, 824 bytes backend->client

curlhttp.serverが返すディレクトリ一覧のHTMLをそのまま受け取っており、HTTPのリクエスト行もヘッダも一切パースせず、TCPのバイト列を右から左に流しているだけであることが、この動作確認から確認できます。