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

use mirror_tool::{
    config::{Config, NetworkConfig, ReplicaPolicy, ServiceConfig, Team},
    engine::Engine,
    flags::{FlagConfig, Transformer, Vault},
};
use std::{path::Path, time::Duration};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpListener, TcpStream},
    time::timeout,
};

fn flag(byte: u8) -> Vec<u8> {
    [vec![byte; 31], vec![b'=']].concat()
}
fn settings(root: &Path) -> FlagConfig {
    FlagConfig {
        database: root.join("vault.sqlite3"),
        partial_timeout_ms: 100,
        ..FlagConfig::default()
    }
}
async fn released(config: &FlagConfig) {
    let mut name = config.database.as_os_str().to_os_string();
    name.push(".lock");
    timeout(Duration::from_secs(2), async {
        loop {
            let file = std::fs::OpenOptions::new()
                .read(true)
                .write(true)
                .open(&name)
                .unwrap();
            if fs2::FileExt::try_lock_exclusive(&file).is_ok() {
                break;
            }
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await
    .unwrap();
}
#[test]
fn opt_in_is_false_and_missing_vault_is_rejected() {
    let base = r#"[[services]]
id="sample"
name="Sample"
listen="127.0.0.1:28181"
upstream_port=28182
"#;
    let config: Config = toml::from_str(base).unwrap();
    assert!(!config.services[0].flag_protection);
    assert!(config.flags.is_none());
    config.validate().unwrap();
    let mut bad = config;
    bad.services[0].flag_protection = true;
    assert!(bad.validate().is_err());
    for length in [0, 15, 129] {
        let cfg = FlagConfig {
            payload_length: length,
            ..FlagConfig::default()
        };
        assert!(cfg.validate().is_err());
    }
}
#[tokio::test]
async fn every_split_binary_integrity_unknown_output_and_eof() {
    let tmp = tempfile::tempdir().unwrap();
    let vault = Vault::with_key(settings(tmp.path()), [7; 32]).unwrap();
    let original = flag(b'A');
    let token = vault
        .translate("a", vec![original.clone()], true)
        .await
        .unwrap()
        .remove(0);
    assert_ne!(token, original);
    assert_eq!(token.len(), original.len());
    for split in 0..=original.len() {
        let mut up = Transformer::new(vault.clone(), "a".into(), true);
        let mut result = up.feed(&original[..split], false).await.unwrap();
        result.extend(up.feed(&original[split..], false).await.unwrap());
        result.extend(up.feed(&[], true).await.unwrap());
        assert_eq!(result, token);
        let mut down = Transformer::new(vault.clone(), "a".into(), false);
        let mut result = down.feed(&token[..split], false).await.unwrap();
        result.extend(down.feed(&token[split..], true).await.unwrap());
        assert_eq!(result, original);
    }
    let input = [
        b"\x00\xffPOST ".as_slice(),
        &original,
        b"\r\nbody=lowercase\x00",
        &original,
        b"!ABC",
    ]
    .concat();
    let mut up = Transformer::new(vault.clone(), "a".into(), true);
    let mut output = Vec::new();
    for chunk in input.chunks(3) {
        output.extend(up.feed(chunk, false).await.unwrap());
    }
    output.extend(up.feed(&[], true).await.unwrap());
    let expected = [
        b"\x00\xffPOST ".as_slice(),
        &token,
        b"\r\nbody=lowercase\x00",
        &token,
        b"!ABC",
    ]
    .concat();
    assert_eq!(output, expected);
    let unknown = flag(b'Z');
    assert_eq!(
        vault
            .translate("a", vec![unknown.clone()], false)
            .await
            .unwrap(),
        vec![unknown]
    );
    assert_eq!(
        vault
            .translate("another", vec![token.clone()], false)
            .await
            .unwrap(),
        vec![token.clone()]
    );
    assert!(vault.translate("a", vec![token], true).await.is_err());
}
#[tokio::test]
async fn partial_candidate_times_out_without_plaintext_fallback() {
    let tmp = tempfile::tempdir().unwrap();
    let vault = Vault::with_key(settings(tmp.path()), [7; 32]).unwrap();
    let mut codec = Transformer::new(vault.clone(), "a".into(), true);
    assert!(codec.feed(b"AAAA", false).await.unwrap().is_empty());
    let (_writer, mut reader) = tokio::io::duplex(64);
    let mut buffer = [0; 32];
    let error = timeout(Duration::from_secs(1), codec.read(&mut reader, &mut buffer))
        .await
        .unwrap()
        .unwrap_err();
    assert_eq!(error.kind(), std::io::ErrorKind::TimedOut);
    assert_eq!(vault.status().errors, 1);
}
#[tokio::test]
async fn concurrent_repeats_share_one_durable_mapping_and_restart_restores() {
    let tmp = tempfile::tempdir().unwrap();
    let cfg = settings(tmp.path());
    let vault = Vault::with_key(cfg.clone(), [7; 32]).unwrap();
    assert!(Vault::with_key(cfg.clone(), [7; 32]).is_err()); // Single-owner lock.
    let mut jobs = tokio::task::JoinSet::new();
    for _ in 0..24 {
        let v = vault.clone();
        jobs.spawn(async move {
            v.translate("a", vec![flag(b'A')], true)
                .await
                .unwrap()
                .remove(0)
        });
    }
    let first = jobs.join_next().await.unwrap().unwrap();
    while let Some(result) = jobs.join_next().await {
        assert_eq!(result.unwrap(), first);
    }
    assert_eq!(vault.status().entries, 1);
    let other = vault
        .translate("b", vec![flag(b'A')], true)
        .await
        .unwrap()
        .remove(0);
    assert_ne!(other, first);
    let bytes = std::fs::read(&cfg.database).unwrap();
    assert!(
        !bytes
            .windows(32)
            .any(|s| s == flag(b'A') || s == first.as_slice())
    );
    drop(vault);
    released(&cfg).await;
    assert!(Vault::with_key(cfg.clone(), [8; 32]).is_err());
    let mut changed = cfg.clone();
    changed.suffix = "!".into();
    assert!(Vault::with_key(changed, [7; 32]).is_err());
    let vault = Vault::with_key(cfg.clone(), [7; 32]).unwrap();
    assert_eq!(
        vault
            .translate("a", vec![first.clone()], false)
            .await
            .unwrap(),
        vec![flag(b'A')]
    );
    assert_eq!(
        vault.translate("a", vec![flag(b'A')], true).await.unwrap(),
        vec![first]
    );
    drop(vault);
    released(&cfg).await;
    let db = rusqlite::Connection::open(&cfg.database).unwrap();
    db.execute(
        "UPDATE flag_entries SET sealed=zeroblob(length(sealed))",
        [],
    )
    .unwrap();
    drop(db);
    assert!(Vault::with_key(cfg, [7; 32]).is_err());
}
#[tokio::test]
async fn capacity_and_disk_failure_do_not_publish_partial_mappings() {
    let tmp = tempfile::tempdir().unwrap();
    let mut cfg = settings(tmp.path());
    cfg.max_entries = 1;
    let vault = Vault::with_key(cfg.clone(), [7; 32]).unwrap();
    assert!(
        vault
            .translate("a", vec![flag(b'A'), flag(b'B')], true)
            .await
            .is_err()
    );
    assert_eq!(vault.status().entries, 0);
    let db = rusqlite::Connection::open(&cfg.database).unwrap();
    db.execute_batch("CREATE TRIGGER fail_insert BEFORE INSERT ON flag_entries BEGIN SELECT RAISE(ABORT,'synthetic write failure'); END;").unwrap();
    assert!(vault.translate("a", vec![flag(b'A')], true).await.is_err());
    assert_eq!(vault.status().entries, 0);
    assert_eq!(
        db.query_row("SELECT COUNT(*) FROM flag_entries", [], |r| r
            .get::<_, u64>(0))
            .unwrap(),
        0
    );
    db.execute_batch("DROP TRIGGER fail_insert;").unwrap();
    let token = vault.translate("a", vec![flag(b'A')], true).await.unwrap();
    assert_eq!(vault.status().entries, 1);
    assert!(vault.translate("a", vec![flag(b'B')], true).await.is_err());
    assert_eq!(
        vault.translate("a", token, false).await.unwrap(),
        vec![flag(b'A')]
    );
}
async fn capture_echo(listener: TcpListener) -> Vec<u8> {
    let (mut stream, _) = listener.accept().await.unwrap();
    let mut bytes = Vec::new();
    stream.read_to_end(&mut bytes).await.unwrap();
    stream.write_all(&bytes).await.unwrap();
    stream.shutdown().await.unwrap();
    bytes
}
fn service(listen: std::net::SocketAddr, port: u16) -> ServiceConfig {
    ServiceConfig {
        id: "echo".into(),
        name: "Echo".into(),
        listen,
        upstream_port: port,
        primary: Some("a".into()),
        replicas: vec!["b".into()],
        enabled: true,
        auto_switch: false,
        replica_policy: ReplicaPolicy::Strict,
        custom: false,
        custom_ip: None,
        flag_protection: true,
    }
}
#[tokio::test]
async fn real_tcp_encodes_once_for_replicas_decodes_reply_and_toggle_is_per_connection() {
    let tmp = tempfile::tempdir().unwrap();
    let vault = Vault::with_key(settings(tmp.path()), [7; 32]).unwrap();
    let primary = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = primary.local_addr().unwrap().port();
    let replica = TcpListener::bind((std::net::Ipv6Addr::LOCALHOST, port))
        .await
        .unwrap();
    let reservation = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let listen = reservation.local_addr().unwrap();
    drop(reservation);
    let engine = Engine::with_flag_vault(NetworkConfig::default(), Some(vault.clone()));
    let teams = vec![
        Team {
            id: "a".into(),
            name: "A".into(),
            address: "127.0.0.1".parse().unwrap(),
            hidden: false,
        },
        Team {
            id: "b".into(),
            name: "B".into(),
            address: "::1".parse().unwrap(),
            hidden: false,
        },
    ];
    let mut cfg = service(listen, port);
    engine.upsert(&cfg, &teams).await.unwrap();
    let p = tokio::spawn(capture_echo(primary));
    let r = tokio::spawn(capture_echo(replica));
    let mut client = TcpStream::connect(listen).await.unwrap();
    // Capture the route by sending a complete non-candidate prefix first.
    client.write_all(b"hello\n").await.unwrap();
    timeout(Duration::from_secs(2), async {
        while engine.snapshots()["echo"].stats.bytes_up < 6 {
            tokio::task::yield_now().await;
        }
    })
    .await
    .unwrap();
    cfg.flag_protection = false;
    engine.upsert(&cfg, &teams).await.unwrap();
    let original = flag(b'A');
    for chunk in original.chunks(3) {
        client.write_all(chunk).await.unwrap();
    }
    client.shutdown().await.unwrap();
    let mut reply = Vec::new();
    timeout(Duration::from_secs(3), client.read_to_end(&mut reply))
        .await
        .unwrap()
        .unwrap();
    assert_eq!(reply, [b"hello\n".as_slice(), &original].concat());
    let a = p.await.unwrap();
    let b = timeout(Duration::from_secs(3), r).await.unwrap().unwrap();
    assert_eq!(a, b);
    assert_ne!(&a[6..], &original);
    assert_eq!(a.len(), reply.len());
    let primary = TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, port))
        .await
        .unwrap();
    cfg.replicas.clear();
    engine.upsert(&cfg, &teams).await.unwrap();
    let p = tokio::spawn(capture_echo(primary));
    let mut client = TcpStream::connect(listen).await.unwrap();
    client.write_all(&original).await.unwrap();
    client.shutdown().await.unwrap();
    let mut reply = Vec::new();
    client.read_to_end(&mut reply).await.unwrap();
    assert_eq!(reply, original);
    assert_eq!(p.await.unwrap(), original);
    engine.shutdown().await;
}

#[tokio::test]
async fn custom_literal_format_preserves_http_lengths_and_repeated_matches() {
    let tmp = tempfile::tempdir().unwrap();
    let cfg = FlagConfig {
        prefix: "FLAG{".into(),
        suffix: "}".into(),
        alphabet: "0123456789ABCDEF".into(),
        payload_length: 32,
        ..settings(tmp.path())
    };
    let vault = Vault::with_key(cfg, [7; 32]).unwrap();
    let flag = b"FLAG{0123456789ABCDEF0123456789ABCDEF}";
    let body = [flag.as_slice(), b" ", flag].concat();
    let header = format!("POST / HTTP/1.1\r\nContent-Length: {}\r\n\r\n", body.len());
    let input = [header.as_bytes(), &body].concat();
    let mut up = Transformer::new(vault.clone(), "a".into(), true);
    let mut encrypted = Vec::new();
    for c in input.chunks(1) {
        encrypted.extend(up.feed(c, false).await.unwrap());
    }
    encrypted.extend(up.feed(&[], true).await.unwrap());
    assert_eq!(encrypted.len(), input.len());
    assert!(encrypted.starts_with(header.as_bytes()));
    assert_ne!(encrypted, input);
    assert_eq!(
        &encrypted[header.len()..header.len() + flag.len()],
        &encrypted[header.len() + flag.len() + 1..]
    );
    let mut down = Transformer::new(vault, "a".into(), false);
    let output = down.feed(&encrypted, true).await.unwrap();
    assert_eq!(output, input);
}

#[tokio::test]
async fn full_vault_fails_closed_before_primary_receives_flag() {
    let tmp = tempfile::tempdir().unwrap();
    let cfg = FlagConfig {
        max_entries: 1,
        ..settings(tmp.path())
    };
    let vault = Vault::with_key(cfg, [7; 32]).unwrap();
    vault
        .translate("echo", vec![flag(b'A')], true)
        .await
        .unwrap();
    let primary = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = primary.local_addr().unwrap().port();
    let reservation = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let listen = reservation.local_addr().unwrap();
    drop(reservation);
    let engine = Engine::with_flag_vault(NetworkConfig::default(), Some(vault));
    let mut cfg = service(listen, port);
    cfg.replicas.clear();
    let teams = vec![Team {
        id: "a".into(),
        name: "A".into(),
        address: "127.0.0.1".parse().unwrap(),
        hidden: false,
    }];
    engine.upsert(&cfg, &teams).await.unwrap();
    let p = tokio::spawn(capture_echo(primary));
    let mut client = TcpStream::connect(listen).await.unwrap();
    client.write_all(&flag(b'B')).await.unwrap();
    client.shutdown().await.unwrap();
    let mut reply = Vec::new();
    let _ = timeout(Duration::from_secs(2), client.read_to_end(&mut reply))
        .await
        .unwrap();
    assert!(p.await.unwrap().is_empty());
    assert!(reply.is_empty());
    assert_eq!(engine.snapshots()["echo"].stats.io_errors, 1);
    engine.shutdown().await;
}