zenith-stack 0.1.0

//! # l4_test — 连接协议层 (L3/L4) 真实测试程序
//!
//! 测试 SourceAdmissionEngine IP 黑名单/端口过滤 + TcpSpliceRelay L4 端口转发。
//!
//! 架构：
//! ```text
//! 测试 1: SourceAdmissionEngine 单元级真实网络测试
//!   hping3 发包 → 程序捕获原始 socket → 解析 L3/L4 头 → SourceAdmissionEngine 评估
//!
//! 测试 2: TcpSpliceRelay L4 端口转发
//!   client → listen:19090 → splice(2) → backend:19091
//! ```

#![deny(unsafe_code)]

use std::io::{Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};

// ── SourceAdmissionEngine 测试 ──────────────────────────────────────────────

/// 构造测试用的准入引擎：
/// - 允许 127.0.0.1 的所有流量
/// - 拒绝 10.0.0.0/24 的 TCP 流量
/// - 只允许 TCP 8080 和 443 端口
/// - 默认拒绝
fn build_admission_engine() -> zenith_net::source_admission::SourceAdmissionEngine {
    use zenith_net::source_admission::*;

    let mut engine = SourceAdmissionEngine::deny_all();

    // 规则 1: 允许 127.0.0.1 所有流量
    engine.add_rule(AdmissionRule {
        id: 1,
        src_ip: IpAddr::V4([127, 0, 0, 1]),
        prefix_len: 0,
        src_port: 0,
        dst_port: 0,
        proto: ProtoMatch::Any,
        action: AdmissionAction::Allow,
        enabled: true,
    }).unwrap();

    // 规则 2: 拒绝 10.0.0.1 的所有流量（IP 黑名单）
    engine.add_rule(AdmissionRule {
        id: 2,
        src_ip: IpAddr::V4([10, 0, 0, 1]),
        prefix_len: 0,
        src_port: 0,
        dst_port: 0,
        proto: ProtoMatch::Any,
        action: AdmissionAction::Deny,
        enabled: true,
    }).unwrap();

    // 规则 3: 拒绝 10.0.0.2 的 TCP 流量（IP+协议黑名单）
    engine.add_rule(AdmissionRule {
        id: 3,
        src_ip: IpAddr::V4([10, 0, 0, 2]),
        prefix_len: 0,
        src_port: 0,
        dst_port: 0,
        proto: ProtoMatch::Tcp,
        action: AdmissionAction::Deny,
        enabled: true,
    }).unwrap();

    // 规则 4: 允许 TCP 8080 端口（端口过滤）
    engine.add_rule(AdmissionRule {
        id: 4,
        src_ip: IpAddr::V4_WILDCARD,
        prefix_len: 0,
        src_port: 0,
        dst_port: 8080,
        proto: ProtoMatch::Tcp,
        action: AdmissionAction::Allow,
        enabled: true,
    }).unwrap();

    // 规则 5: 允许 TCP 443 端口
    engine.add_rule(AdmissionRule {
        id: 5,
        src_ip: IpAddr::V4_WILDCARD,
        prefix_len: 0,
        src_port: 0,
        dst_port: 443,
        proto: ProtoMatch::Tcp,
        action: AdmissionAction::Allow,
        enabled: true,
    }).unwrap();

    engine
}

/// 将 std::net::IpAddr 转为 zenith_net::source_admission::IpAddr
fn to_net_ip(ip: std::net::IpAddr) -> zenith_net::source_admission::IpAddr {
    use zenith_net::source_admission::IpAddr;
    match ip {
        std::net::IpAddr::V4(v4) => IpAddr::V4(v4.octets()),
        std::net::IpAddr::V6(v6) => IpAddr::V6(v6.octets()),
    }
}

fn test_source_admission_engine() -> (u32, u32) {
    use zenith_net::source_admission::*;

    println!("\n============================================================");
    println!("测试 1: SourceAdmissionEngine IP 黑名单/端口过滤");
    println!("============================================================");

    let engine = build_admission_engine();
    let mut pass = 0u32;
    let mut fail = 0u32;

    // 测试 1a: 127.0.0.1 应被允许
    let action = engine.evaluate(
        IpAddr::V4([127, 0, 0, 1]),
        12345, 8080, 6, // TCP
    );
    if action == AdmissionAction::Allow {
        println!("  [PASS] 127.0.0.1:12345→8080/TCP 允许");
        pass += 1;
    } else {
        println!("  [FAIL] 127.0.0.1:12345→8080/TCP 应允许，实际: {:?}", action);
        fail += 1;
    }

    // 测试 1b: 10.0.0.1 应被拒绝（IP 黑名单）
    let action = engine.evaluate(
        IpAddr::V4([10, 0, 0, 1]),
        12345, 8080, 6,
    );
    if action == AdmissionAction::Deny {
        println!("  [PASS] 10.0.0.1:12345→8080/TCP 被拒绝 (IP 黑名单)");
        pass += 1;
    } else {
        println!("  [FAIL] 10.0.0.1:12345→8080/TCP 应拒绝，实际: {:?}", action);
        fail += 1;
    }

    // 测试 1c: 10.0.0.2 TCP 应被拒绝（IP+协议黑名单）
    let action = engine.evaluate(
        IpAddr::V4([10, 0, 0, 2]),
        12345, 8080, 6,
    );
    if action == AdmissionAction::Deny {
        println!("  [PASS] 10.0.0.2:12345→8080/TCP 被拒绝 (IP+协议黑名单)");
        pass += 1;
    } else {
        println!("  [FAIL] 10.0.0.2:12345→8080/TCP 应拒绝，实际: {:?}", action);
        fail += 1;
    }

    // 测试 1d: 192.168.1.100→8080/TCP 应被允许（端口规则 4）
    let action = engine.evaluate(
        IpAddr::V4([192, 168, 1, 100]),
        54321, 8080, 6,
    );
    if action == AdmissionAction::Allow {
        println!("  [PASS] 192.168.1.100:54321→8080/TCP 允许 (端口规则)");
        pass += 1;
    } else {
        println!("  [FAIL] 192.168.1.100:54321→8080/TCP 应允许，实际: {:?}", action);
        fail += 1;
    }

    // 测试 1e: 192.168.1.100→22/TCP 应被拒绝（端口不在白名单）
    let action = engine.evaluate(
        IpAddr::V4([192, 168, 1, 100]),
        54321, 22, 6,
    );
    if action == AdmissionAction::Deny {
        println!("  [PASS] 192.168.1.100:54321→22/TCP 被拒绝 (端口不在白名单)");
        pass += 1;
    } else {
        println!("  [FAIL] 192.168.1.100:54321→22/TCP 应拒绝，实际: {:?}", action);
        fail += 1;
    }

    // 测试 1f: 192.168.1.100→8080/UDP 应被拒绝（协议不匹配）
    let action = engine.evaluate(
        IpAddr::V4([192, 168, 1, 100]),
        54321, 8080, 17, // UDP
    );
    if action == AdmissionAction::Deny {
        println!("  [PASS] 192.168.1.100:54321→8080/UDP 被拒绝 (协议不匹配)");
        pass += 1;
    } else {
        println!("  [FAIL] 192.168.1.100:54321→8080/UDP 应拒绝，实际: {:?}", action);
        fail += 1;
    }

    // 测试 1g: IPv6 ::1 应被拒绝（默认拒绝 + 无 IPv6 规则）
    let action = engine.evaluate(
        IpAddr::V6([0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1]),
        12345, 8080, 6,
    );
    if action == AdmissionAction::Deny {
        println!("  [PASS] [::1]:12345→8080/TCP 被拒绝 (无 IPv6 白名单规则)");
        pass += 1;
    } else {
        println!("  [FAIL] [::1]:12345→8080/TCP 应拒绝，实际: {:?}", action);
        fail += 1;
    }

    // 测试 1h: 10.0.0.3→8080/TCP 应被允许（规则 4 通配 IP + 端口 8080）
    let action = engine.evaluate(
        IpAddr::V4([10, 0, 0, 3]),
        12345, 8080, 6,
    );
    if action == AdmissionAction::Allow {
        println!("  [PASS] 10.0.0.3:12345→8080/TCP 允许 (通配 IP + 端口 8080)");
        pass += 1;
    } else {
        println!("  [FAIL] 10.0.0.3:12345→8080/TCP 应允许，实际: {:?}", action);
        fail += 1;
    }

    println!("\n  SourceAdmissionEngine: {} 通过, {} 失败", pass, fail);
    (pass, fail)
}

// ── TcpSpliceRelay L4 端口转发测试 ──────────────────────────────────────────

fn test_l4_port_forward() -> (u32, u32) {
    println!("\n============================================================");
    println!("测试 2: L4 TCP 端口转发 (TcpRelay 统一 API)");
    println!("============================================================");

    let mut pass = 0u32;
    let mut fail = 0u32;

    // 启动后端服务器（端口 0 = OS 自动分配，避免固定端口冲突）
    let backend = TcpListener::bind("127.0.0.1:0").unwrap();
    let backend_addr = backend.local_addr().unwrap();
    println!("  后端服务器: {}", backend_addr);

    // 启动 L4 转发监听（端口 0 = OS 自动分配）
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let listen_addr = listener.local_addr().unwrap();
    println!("  转发监听: {} → {}", listen_addr, backend_addr);

    // 后端线程：接受一个连接，回显收到的数据
    let backend_thread = std::thread::spawn(move || {
        let (mut conn, peer) = backend.accept().unwrap();
        println!("  后端接受连接 from {}", peer);
        let mut buf = [0u8; 4096];
        let n = conn.read(&mut buf).unwrap();
        let received = &buf[..n];
        println!("  后端收到 {} 字节: {}", n, String::from_utf8_lossy(received));
        // 回显
        conn.write_all(received).unwrap();
        println!("  后端回显 {} 字节", n);
    });

    // 转发线程：接受客户端连接，splice 到后端
    let relay_thread = std::thread::spawn(move || {
        let (client_conn, peer) = listener.accept().unwrap();
        println!("  转发接受连接 from {}", peer);

        // 连接后端（使用 OS 分配的动态端口）
        let upstream = TcpStream::connect(backend_addr).unwrap();
        println!("  转发连接后端 {}", backend_addr);

        // TcpRelay 统一 API：自动选择 splice(2) 或 std::io::copy
        let relay = zenith_forward::TcpRelay::default();
        match relay.relay(&client_conn, &upstream) {
            Ok((c2u, u2c)) => {
                println!("  TcpRelay 完成: client→upstream={}B upstream→client={}B", c2u, u2c);
            }
            Err(e) => {
                println!("  TcpRelay 错误: {}", e);
            }
        }
    });

    // 客户端：连接转发端口（使用 OS 分配的动态端口），发送数据，验证回显
    std::thread::sleep(std::time::Duration::from_millis(100));
    let mut client = TcpStream::connect(listen_addr).unwrap();
    let test_data = b"hello-l4-forward";
    client.write_all(test_data).unwrap();
    println!("  客户端发送 {} 字节: {}", test_data.len(), String::from_utf8_lossy(test_data));

    // 读取回显
    let mut buf = [0u8; 4096];
    let n = client.read(&mut buf).unwrap();
    let echoed = &buf[..n];

    if echoed == test_data {
        println!("  [PASS] L4 端口转发回显正确 ({} 字节)", n);
        pass += 1;
    } else {
        println!("  [FAIL] L4 端口转发回显不匹配: 期望 {:?} 实际 {:?}", test_data, echoed);
        fail += 1;
    }

    // 关闭客户端连接
    drop(client);

    // 等待线程完成
    let _ = relay_thread.join();
    let _ = backend_thread.join();

    println!("\n  TcpSpliceRelay: {} 通过, {} 失败", pass, fail);
    (pass, fail)
}

// ── ForwardEngine + AddressWhitelist L4 IP 黑名单测试 ──────────────────────

fn test_forward_whitelist() -> (u32, u32) {
    use zenith_forward::{AddressWhitelist, ForwardEngine, ForwardProtocol};

    println!("\n============================================================");
    println!("测试 3: ForwardEngine + AddressWhitelist L4 IP 黑名单");
    println!("============================================================");

    let mut pass = 0u32;
    let mut fail = 0u32;

    // 白名单只允许 127.0.0.1，拒绝所有其他 IP
    let whitelist = AddressWhitelist::new()
        .with_allowed_ip(std::net::IpAddr::from([127u8, 0, 0, 1]));
    let mut engine = ForwardEngine::new().with_whitelist(whitelist);

    // 测试 3a: 127.0.0.1 应被允许
    let allowed = engine.is_address_allowed(&"127.0.0.1:8080".parse().unwrap());
    if allowed {
        println!("  [PASS] 127.0.0.1:8080 允许通过白名单");
        pass += 1;
    } else {
        println!("  [FAIL] 127.0.0.1:8080 应被允许");
        fail += 1;
    }

    // 测试 3b: 10.0.0.1 应被拒绝
    let allowed = engine.is_address_allowed(&"10.0.0.1:8080".parse().unwrap());
    if !allowed {
        println!("  [PASS] 10.0.0.1:8080 被 IP 白名单拒绝");
        pass += 1;
    } else {
        println!("  [FAIL] 10.0.0.1:8080 应被拒绝");
        fail += 1;
    }

    // 测试 3c: 192.168.1.1 应被拒绝
    let allowed = engine.is_address_allowed(&"192.168.1.1:8080".parse().unwrap());
    if !allowed {
        println!("  [PASS] 192.168.1.1:8080 被 IP 白名单拒绝");
        pass += 1;
    } else {
        println!("  [FAIL] 192.168.1.1:8080 应被拒绝");
        fail += 1;
    }

    // 测试 3d: create_session 拒绝非白名单 IP
    let session = engine.create_session(
        ForwardProtocol::Tcp,
        "10.0.0.1:12345".parse().unwrap(),
        "127.0.0.1:9090".parse().unwrap(),
    );
    if session.is_none() {
        println!("  [PASS] create_session 拒绝非白名单客户端 IP (L4 SSRF 防护)");
        pass += 1;
    } else {
        println!("  [FAIL] create_session 应拒绝非白名单 IP");
        fail += 1;
    }

    // 测试 3e: create_session 拒绝非白名单上游 IP
    let session = engine.create_session(
        ForwardProtocol::Tcp,
        "127.0.0.1:12345".parse().unwrap(),
        "10.0.0.1:9090".parse().unwrap(), // 上游也不在白名单
    );
    if session.is_none() {
        println!("  [PASS] create_session 拒绝非白名单上游 IP (L4 SSRF 防护)");
        pass += 1;
    } else {
        println!("  [FAIL] create_session 应拒绝非白名单上游");
        fail += 1;
    }

    // 测试 3f: create_session 允许白名单内的 IP
    let session = engine.create_session(
        ForwardProtocol::Tcp,
        "127.0.0.1:12345".parse().unwrap(),
        "127.0.0.1:9090".parse().unwrap(),
    );
    if session.is_some() {
        println!("  [PASS] create_session 允许白名单内 IP 双向通过");
        pass += 1;
    } else {
        println!("  [FAIL] create_session 应允许白名单内 IP");
        fail += 1;
    }

    println!("\n  ForwardEngine 白名单: {} 通过, {} 失败", pass, fail);
    (pass, fail)
}

// ── CIDR 子网 IP 黑名单测试 ─────────────────────────────────────────────────

fn test_cidr_blacklist() -> (u32, u32) {
    use zenith_forward::AddressWhitelist;

    println!("\n============================================================");
    println!("测试 4: CIDR 子网 IP 黑名单 (AddressWhitelist)");
    println!("============================================================");

    let mut pass = 0u32;
    let mut fail = 0u32;

    // 白名单：允许 127.0.0.0/8 子网 + 192.168.0.0/16 子网
    let whitelist = AddressWhitelist::new()
        .with_allowed_cidr_ip(std::net::IpAddr::from([127u8, 0, 0, 0]), 8)
        .with_allowed_cidr_ip(std::net::IpAddr::from([192u8, 168, 0, 0]), 16);

    // 测试 4a: 127.0.0.1 在 /8 子网内
    if whitelist.is_allowed(&"127.0.0.1:8080".parse().unwrap()) {
        println!("  [PASS] 127.0.0.1 在 127.0.0.0/8 子网内 → 允许");
        pass += 1;
    } else {
        println!("  [FAIL] 127.0.0.1 应在 127.0.0.0/8 内");
        fail += 1;
    }

    // 测试 4b: 127.0.0.254 在 /8 子网内
    if whitelist.is_allowed(&"127.0.0.254:8080".parse().unwrap()) {
        println!("  [PASS] 127.0.0.254 在 127.0.0.0/8 子网内 → 允许");
        pass += 1;
    } else {
        println!("  [FAIL] 127.0.0.254 应在 127.0.0.0/8 内");
        fail += 1;
    }

    // 测试 4c: 192.168.1.100 在 /16 子网内
    if whitelist.is_allowed(&"192.168.1.100:8080".parse().unwrap()) {
        println!("  [PASS] 192.168.1.100 在 192.168.0.0/16 子网内 → 允许");
        pass += 1;
    } else {
        println!("  [FAIL] 192.168.1.100 应在 192.168.0.0/16 内");
        fail += 1;
    }

    // 测试 4d: 10.0.0.1 不在任何子网内 → 拒绝
    if !whitelist.is_allowed(&"10.0.0.1:8080".parse().unwrap()) {
        println!("  [PASS] 10.0.0.1 不在白名单子网 → 拒绝");
        pass += 1;
    } else {
        println!("  [FAIL] 10.0.0.1 应被拒绝");
        fail += 1;
    }

    // 测试 4e: 192.169.1.1 不在 192.168.0.0/16 内 → 拒绝
    if !whitelist.is_allowed(&"192.169.1.1:8080".parse().unwrap()) {
        println!("  [PASS] 192.169.1.1 不在 192.168.0.0/16 → 拒绝");
        pass += 1;
    } else {
        println!("  [FAIL] 192.169.1.1 应被拒绝");
        fail += 1;
    }

    // 测试 4f: 空白名单默认拒绝所有（fail-closed）
    let empty_wl = AddressWhitelist::new();
    if !empty_wl.is_allowed(&"127.0.0.1:8080".parse().unwrap()) {
        println!("  [PASS] 空白名单 fail-closed 拒绝所有");
        pass += 1;
    } else {
        println!("  [FAIL] 空白名单应拒绝所有");
        fail += 1;
    }

    // 测试 4g: IP 级别检查（不含端口）
    if whitelist.is_ip_allowed(&std::net::IpAddr::from([127u8, 0, 0, 1])) {
        println!("  [PASS] is_ip_allowed(127.0.0.1) → true");
        pass += 1;
    } else {
        println!("  [FAIL] is_ip_allowed(127.0.0.1) 应为 true");
        fail += 1;
    }

    if !whitelist.is_ip_allowed(&std::net::IpAddr::from([10u8, 0, 0, 1])) {
        println!("  [PASS] is_ip_allowed(10.0.0.1) → false");
        pass += 1;
    } else {
        println!("  [FAIL] is_ip_allowed(10.0.0.1) 应为 false");
        fail += 1;
    }

    println!("\n  CIDR 黑名单: {} 通过, {} 失败", pass, fail);
    (pass, fail)
}

fn main() {
    println!("Zenith 连接协议层 (L3/L4) 真实测试");

    let (p1, f1) = test_source_admission_engine();
    let (p2, f2) = test_l4_port_forward();
    let (p3, f3) = test_forward_whitelist();
    let (p4, f4) = test_cidr_blacklist();

    let total_pass = p1 + p2 + p3 + p4;
    let total_fail = f1 + f2 + f3 + f4;

    println!("\n============================================================");
    println!("  总计: {} 通过, {} 失败", total_pass, total_fail);
    println!("============================================================");

    if total_fail > 0 {
        std::process::exit(1);
    }
}