ctf-ad-mirror-tool-rust/tests/forwarding.rs

use mirror_tool::{
    config::{NetworkConfig, ReplicaPolicy, ServiceConfig, Team},
    engine::Engine,
};
use std::{
    net::{IpAddr, Ipv4Addr, SocketAddr},
    time::Duration,
};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpListener, TcpStream},
    task::JoinHandle,
    time::{sleep, timeout},
};

fn settings() -> NetworkConfig {
    NetworkConfig {
        max_buffered_bytes: 65536,
        connect_timeout_ms: 200,
        buffer_size: 4096,
        replica_queue_chunks: 4,
        max_connections: 32,
        drain_timeout_ms: 100,
    }
}
async fn free_address() -> SocketAddr {
    TcpListener::bind("127.0.0.1:0")
        .await
        .unwrap()
        .local_addr()
        .unwrap()
}
fn service(listen: SocketAddr, upstream: SocketAddr) -> ServiceConfig {
    ServiceConfig {
        custom: false,
        custom_ip: None,
        id: "service".into(),
        name: "Test".into(),
        listen,
        upstream_port: upstream.port(),
        primary: Some("primary".into()),
        replicas: vec![],
        enabled: true,
        auto_switch: false,
        replica_policy: ReplicaPolicy::Isolate,
    }
}
fn teams() -> Vec<Team> {
    vec![Team {
        hidden: false,
        id: "primary".into(),
        name: "Primary".into(),
        address: IpAddr::V4(Ipv4Addr::LOCALHOST),
    }]
}
async fn echo() -> (SocketAddr, JoinHandle<()>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let address = listener.local_addr().unwrap();
    let task = tokio::spawn(async move {
        loop {
            let (mut stream, _) = listener.accept().await.unwrap();
            tokio::spawn(async move {
                let mut buffer = [0; 4096];
                loop {
                    let n = stream.read(&mut buffer).await.unwrap_or(0);
                    if n == 0 {
                        let _ = stream.shutdown().await;
                        break;
                    }
                    if stream.write_all(&buffer[..n]).await.is_err() {
                        break;
                    }
                }
            });
        }
    });
    (address, task)
}
async fn exchange(address: SocketAddr, payload: &[u8]) -> Vec<u8> {
    let stream = TcpStream::connect(address).await.unwrap();
    let (mut reader, mut writer) = stream.into_split();
    let upload = async {
        writer.write_all(payload).await.unwrap();
        writer.shutdown().await.unwrap();
    };
    let download = async {
        let mut response = vec![];
        reader.read_to_end(&mut response).await.unwrap();
        response
    };
    let (_, response) = timeout(Duration::from_secs(5), async {
        tokio::join!(upload, download)
    })
    .await
    .unwrap();
    response
}
async fn settled(engine: &Engine) {
    timeout(Duration::from_secs(2), async {
        loop {
            if engine.snapshots()["service"].stats.active == 0 {
                break;
            }
            sleep(Duration::from_millis(5)).await;
        }
    })
    .await
    .unwrap();
}

#[tokio::test]
async fn bulk_bidirectional_bytes_and_half_close() {
    let (upstream, server) = echo().await;
    let listen = free_address().await;
    let engine = Engine::new(settings());
    engine
        .upsert(&service(listen, upstream), &teams())
        .await
        .unwrap();
    let payload: Vec<u8> = (0..2_000_000).map(|n| (n % 251) as u8).collect();
    assert_eq!(exchange(listen, &payload).await, payload);
    settled(&engine).await;
    let stats = engine.snapshots()["service"].stats.clone();
    assert_eq!(stats.bytes_up, payload.len() as u64);
    assert_eq!(stats.bytes_down, payload.len() as u64);
    assert_eq!(stats.completed, 1);
    engine.shutdown().await;
    server.abort();
}

#[tokio::test]
async fn response_after_request_eof_and_idle_over_three_seconds() {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let upstream = listener.local_addr().unwrap();
    let server = tokio::spawn(async move {
        let (mut stream, _) = listener.accept().await.unwrap();
        let mut request = vec![];
        stream.read_to_end(&mut request).await.unwrap();
        assert_eq!(request, b"after idle");
        stream.write_all(b"after eof").await.unwrap();
    });
    let listen = free_address().await;
    let engine = Engine::new(settings());
    engine
        .upsert(&service(listen, upstream), &teams())
        .await
        .unwrap();
    let mut stream = TcpStream::connect(listen).await.unwrap();
    sleep(Duration::from_millis(3200)).await;
    stream.write_all(b"after idle").await.unwrap();
    stream.shutdown().await.unwrap();
    let mut reply = vec![];
    timeout(Duration::from_secs(2), stream.read_to_end(&mut reply))
        .await
        .unwrap()
        .unwrap();
    assert_eq!(reply, b"after eof");
    server.await.unwrap();
    engine.shutdown().await;
}

#[tokio::test]
async fn retarget_only_new_sessions_and_reenable_without_port_leak() {
    let (a, server_a) = echo().await;
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let b = listener.local_addr().unwrap();
    let server_b = tokio::spawn(async move {
        let (mut stream, _) = listener.accept().await.unwrap();
        let mut request = vec![];
        stream.read_to_end(&mut request).await.unwrap();
        stream.write_all(b"new target").await.unwrap();
    });
    let listen = free_address().await;
    let engine = Engine::new(settings());
    let mut config = service(listen, a);
    engine.upsert(&config, &teams()).await.unwrap();
    let mut old = TcpStream::connect(listen).await.unwrap();
    old.write_all(b"before").await.unwrap();
    let mut before = [0; 6];
    old.read_exact(&mut before).await.unwrap();
    config.upstream_port = b.port();
    engine.upsert(&config, &teams()).await.unwrap();
    old.write_all(b"old target").await.unwrap();
    old.shutdown().await.unwrap();
    let mut reply = vec![];
    old.read_to_end(&mut reply).await.unwrap();
    assert_eq!(reply, b"old target");
    assert_eq!(exchange(listen, b"request").await, b"new target");
    config.enabled = false;
    engine.upsert(&config, &teams()).await.unwrap();
    assert!(!engine.snapshots()["service"].running);
    assert!(TcpStream::connect(listen).await.is_err());
    config.enabled = true;
    config.upstream_port = a.port();
    engine.upsert(&config, &teams()).await.unwrap();
    assert_eq!(exchange(listen, b"enabled").await, b"enabled");
    engine.shutdown().await;
    server_a.abort();
    server_b.await.unwrap();
}

#[tokio::test]
async fn bind_error_is_reported_and_working_listener_preserved() {
    let (upstream, server) = echo().await;
    let listen = free_address().await;
    let occupied = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let engine = Engine::new(settings());
    let mut config = service(listen, upstream);
    engine.upsert(&config, &teams()).await.unwrap();
    config.listen = occupied.local_addr().unwrap();
    assert!(engine.upsert(&config, &teams()).await.is_err());
    assert_eq!(exchange(listen, b"still running").await, b"still running");
    engine.shutdown().await;
    server.abort();
}

#[tokio::test]
async fn process_connection_limit_sheds_and_shutdown_cancels_active_sessions() {
    let (upstream, server) = echo().await;
    let listen = free_address().await;
    let engine = Engine::new(NetworkConfig {
        max_connections: 1,
        ..settings()
    });
    engine
        .upsert(&service(listen, upstream), &teams())
        .await
        .unwrap();
    let mut first = TcpStream::connect(listen).await.unwrap();
    first.write_all(b"a").await.unwrap();
    let mut byte = [0];
    first.read_exact(&mut byte).await.unwrap();
    let second_listen = free_address().await;
    let mut second_config = service(second_listen, upstream);
    second_config.id = "second".into();
    engine.upsert(&second_config, &teams()).await.unwrap();
    let mut rejected = TcpStream::connect(second_listen).await.unwrap();
    assert_eq!(
        timeout(Duration::from_secs(1), rejected.read(&mut byte))
            .await
            .unwrap()
            .unwrap(),
        0
    );
    assert_eq!(engine.snapshots()["second"].stats.rejected, 1);
    timeout(Duration::from_secs(1), engine.shutdown())
        .await
        .unwrap();
    assert_eq!(engine.snapshots()["service"].stats.active, 0);
    assert!(!engine.snapshots()["service"].running);
    assert_eq!(
        timeout(Duration::from_secs(1), first.read(&mut byte))
            .await
            .unwrap()
            .unwrap(),
        0
    );
    server.abort();
}

#[tokio::test]
async fn disable_keeps_existing_connection_and_remove_joins_it() {
    let (upstream, server) = echo().await;
    let listen = free_address().await;
    let engine = Engine::new(settings());
    let mut config = service(listen, upstream);
    engine.upsert(&config, &teams()).await.unwrap();
    let mut stream = TcpStream::connect(listen).await.unwrap();
    stream.write_all(b"a").await.unwrap();
    let mut byte = [0];
    stream.read_exact(&mut byte).await.unwrap();
    config.enabled = false;
    engine.upsert(&config, &teams()).await.unwrap();
    stream.write_all(b"b").await.unwrap();
    stream.read_exact(&mut byte).await.unwrap();
    assert_eq!(&byte, b"b");
    timeout(Duration::from_secs(1), engine.remove("service"))
        .await
        .unwrap()
        .unwrap();
    assert!(engine.snapshots().is_empty());
    assert_eq!(stream.read(&mut byte).await.unwrap(), 0);
    engine.shutdown().await;
    server.abort();
}

// Distinct loopback addresses let primary and replica share the configured upstream port.
async fn replica_fixture(
    policy: ReplicaPolicy,
) -> (Engine, SocketAddr, TcpListener, JoinHandle<()>) {
    let (upstream, server) = echo().await;
    let replica = TcpListener::bind(SocketAddr::new("::1".parse().unwrap(), upstream.port()))
        .await
        .unwrap();
    let listen = free_address().await;
    let mut config = service(listen, upstream);
    config.replica_policy = policy;
    config.replicas = vec!["replica".into()];
    let mut teams = teams();
    teams.push(Team {
        hidden: false,
        id: "replica".into(),
        name: "Replica".into(),
        address: "::1".parse().unwrap(),
    });
    let engine = Engine::new(settings());
    engine.upsert(&config, &teams).await.unwrap();
    (engine, listen, replica, server)
}

#[tokio::test]
async fn strict_mirror_exact_bytes_and_drains_large_replica_replies() {
    let (engine, listen, replica, server) = replica_fixture(ReplicaPolicy::Strict).await;
    let payload: Vec<u8> = (0..1_000_000).map(|n| (n % 251) as u8).collect();
    let expected = payload.clone();
    let worker = tokio::spawn(async move {
        let (stream, _) = replica.accept().await.unwrap();
        let (mut reader, mut writer) = stream.into_split();
        let read = async {
            let mut bytes = vec![];
            reader.read_to_end(&mut bytes).await.unwrap();
            assert_eq!(bytes, expected);
        };
        let write = async {
            writer.write_all(&vec![17; 2_000_000]).await.unwrap();
            writer.shutdown().await.unwrap();
        };
        tokio::join!(read, write);
    });
    assert_eq!(exchange(listen, &payload).await, payload);
    timeout(Duration::from_secs(3), worker)
        .await
        .unwrap()
        .unwrap();
    settled(&engine).await;
    assert_eq!(
        engine.snapshots()["service"].stats.replica_bytes,
        payload.len() as u64
    );
    assert_eq!(engine.snapshots()["service"].stats.replica_disconnects, 0);
    engine.shutdown().await;
    server.abort();
}

#[tokio::test]
async fn isolate_unavailable_replica_preserves_primary_and_counts_failure() {
    let (engine, listen, replica, server) = replica_fixture(ReplicaPolicy::Isolate).await;
    drop(replica);
    // Give the worker time to fail before EOF rather than overflowing its bounded queue.
    let mut client = TcpStream::connect(listen).await.unwrap();
    client.write_all(b"one").await.unwrap();
    let mut bytes = [0; 3];
    client.read_exact(&mut bytes).await.unwrap();
    sleep(Duration::from_millis(50)).await;
    client.write_all(b"two").await.unwrap();
    client.shutdown().await.unwrap();
    let mut reply = vec![];
    client.read_to_end(&mut reply).await.unwrap();
    assert_eq!(reply, b"two");
    settled(&engine).await;
    let stats = engine.snapshots()["service"].stats.clone();
    assert_eq!(stats.replica_errors, 1);
    assert_eq!(stats.replica_disconnects, 1);
    assert_eq!(stats.io_errors, 0);
    engine.shutdown().await;
    server.abort();
}

#[tokio::test]
async fn strict_unavailable_replica_ends_session_before_primary_receives_payload() {
    let (engine, listen, replica, server) = replica_fixture(ReplicaPolicy::Strict).await;
    drop(replica);
    let mut client = TcpStream::connect(listen).await.unwrap();
    let _ = client.write_all(b"must not forward").await;
    let mut bytes = vec![];
    let _ = timeout(Duration::from_secs(1), client.read_to_end(&mut bytes))
        .await
        .unwrap();
    settled(&engine).await;
    let stats = engine.snapshots()["service"].stats.clone();
    assert_eq!(stats.bytes_up, 0);
    assert_eq!(stats.io_errors, 1);
    assert_eq!(stats.replica_errors, 1);
    engine.shutdown().await;
    server.abort();
}

#[tokio::test]
async fn isolate_slow_replica_disconnects_instead_of_corrupting_primary_stream() {
    let (engine, listen, replica, server) = replica_fixture(ReplicaPolicy::Isolate).await;
    let slow = tokio::spawn(async move {
        let (_stream, _) = replica.accept().await.unwrap();
        sleep(Duration::from_secs(20)).await;
    });
    let payload = vec![93; 16_000_000];
    assert_eq!(exchange(listen, &payload).await, payload);
    settled(&engine).await;
    assert!(engine.snapshots()["service"].stats.replica_disconnects >= 1);
    engine.shutdown().await;
    server.abort();
    slow.abort();
}

#[tokio::test]
async fn strict_slow_replica_backpressures_and_shutdown_joins_workers() {
    let (engine, listen, replica, server) = replica_fixture(ReplicaPolicy::Strict).await;
    let slow = tokio::spawn(async move {
        let (_stream, _) = replica.accept().await.unwrap();
        sleep(Duration::from_secs(20)).await;
    });
    let mut client = TcpStream::connect(listen).await.unwrap();
    let payload = vec![31; 32_000_000];
    assert!(
        timeout(Duration::from_millis(200), client.write_all(&payload))
            .await
            .is_err()
    );
    let stats = engine.snapshots()["service"].stats.clone();
    assert!(stats.bytes_up < payload.len() as u64);
    assert_eq!(stats.replica_disconnects, 0);
    timeout(Duration::from_secs(1), engine.shutdown())
        .await
        .unwrap();
    assert_eq!(engine.snapshots()["service"].stats.active, 0);
    server.abort();
    slow.abort();
}

#[tokio::test]
async fn quiesce_keeps_existing_stream_but_route_rollback_cannot_reopen_listener() {
    let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let address = upstream.local_addr().unwrap();
    let task = tokio::spawn(async move {
        let (mut stream, _) = upstream.accept().await.unwrap();
        let mut bytes = vec![];
        stream.read_to_end(&mut bytes).await.unwrap();
        stream.write_all(&bytes).await.unwrap();
    });
    let engine = Engine::new(settings());
    let listen = free_address().await;
    let original = service(listen, address);
    engine.upsert(&original, &teams()).await.unwrap();
    let mut client = TcpStream::connect(listen).await.unwrap();
    timeout(Duration::from_secs(1), async {
        while engine.snapshots()["service"].stats.active == 0 {
            tokio::task::yield_now().await;
        }
    })
    .await
    .unwrap();
    engine.quiesce().await;
    assert!(TcpStream::connect(listen).await.is_err());
    engine.upsert(&original, &teams()).await.unwrap();
    assert!(TcpStream::connect(listen).await.is_err());
    assert!(!engine.snapshots()["service"].running);
    client
        .write_all(b"existing connection survives quiesce")
        .await
        .unwrap();
    client.shutdown().await.unwrap();
    let mut bytes = vec![];
    timeout(Duration::from_secs(1), client.read_to_end(&mut bytes))
        .await
        .unwrap()
        .unwrap();
    assert_eq!(bytes, b"existing connection survives quiesce");
    task.await.unwrap();
    engine.shutdown().await;
}