zenith-stack 0.1.0

//! # full_feature_server — 全功能安全验证服务器
//!
//! 启用 WAF + 代理转发 + IP 黑名单 + 端口转发 + 协议转换等全量安全功能。
//!
//! 配置：
//! - WAF 默认检测器（SecurityPipeline 已内置）
//! - 反向代理路由 `/api/*` → 后端 127.0.0.1:19090（HTTP→HTTP 转发）
//! - 虚拟主机白名单 `localhost` + `127.0.0.1`
//! - IP 黑名单：通过 zenith-forward AddressWhitelist 拒绝特定 IP
//! - 端口转发：L4 TCP 转发到后端
//! - 端口协议转换：HTTP 8080 → HTTP/1.1 9090（通过 ProxyEngine forward）

#![deny(unsafe_code)]

use std::env;
use std::net::SocketAddr;
use std::process::ExitCode;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;

use zenith_stack::api::CanonicalResponse;
use zenith_stack::forward::AddressWhitelist;
use zenith_stack::proxy::LoadBalanceStrategy;
use zenith_stack::web::app::{success_response, App};
use zenith_stack::web::server::{ProtocolServer, ServerConfig};

fn parse_args() -> Result<(SocketAddr,), String> {
    let mut port: u16 = 18080;
    let mut host: String = "127.0.0.1".to_string();
    let mut it = env::args().skip(1);
    while let Some(arg) = it.next() {
        match arg.as_str() {
            "--port" => {
                let val = it.next().ok_or("--port needs <PORT>")?;
                port = val.parse().map_err(|e| format!("port parse: {e}"))?;
            }
            "--host" => {
                host = it.next().ok_or("--host needs <HOST>")?;
            }
            "-h" | "--help" => {
                println!("full_feature_server [--port PORT] [--host HOST]");
                std::process::exit(0);
            }
            other => return Err(format!("unknown arg {other:?}")),
        }
    }
    // IPv6 地址需要方括号包裹（如 [::]:8080），IPv4 不需要
    let bind: SocketAddr = if host.contains(':') {
        format!("[{host}]:{port}")
    } else {
        format!("{host}:{port}")
    }
    .parse()
    .map_err(|e| format!("addr: {e}"))?;
    Ok((bind,))
}

/// 最大并发连接数（超过时拒绝新连接）
const MAX_CONNECTIONS: u32 = 1024;

fn build_app() -> Result<App, String> {
    let mut app = App::new();
    app.get("/", |_req, _rm| {
        Ok(success_response(
            "full_feature_server OK\n",
            "text/plain; charset=utf-8",
        ))
    });
    app.get("/health", |_req, _rm| {
        Ok(success_response(r#"{"status":"healthy"}"#, "application/json"))
    });
    app.get("/hello/:name", |_req, rm| {
        let name = rm.get("name").unwrap_or("stranger");
        Ok(success_response(
            format!("Hello, {name}!\n"),
            "text/plain; charset=utf-8",
        ))
    });
    app.post("/echo", |req, _rm| {
        let mut resp = CanonicalResponse::new(200);
        let _ = resp.add_header(b"content-type", b"application/octet-stream");
        resp.set_body(req.body().to_vec());
        Ok(resp)
    });
    app.validate().map_err(|e| format!("route validation failed: {e}"))?;
    Ok(app)
}

fn main() -> ExitCode {
    let _ = tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .try_init();

    let _global_rt = zenith_stack::rt::init_global(zenith_stack::rt::RuntimeConfig::auto());

    let (bind,) = match parse_args() {
        Ok(a) => a,
        Err(e) => {
            eprintln!("args error: {e}");
            return ExitCode::from(2);
        }
    };

    let app = match build_app() {
        Ok(a) => a,
        Err(e) => {
            eprintln!("{e}");
            return ExitCode::from(4);
        }
    };

    // ── 配置安全管道 ──────────────────────────────────────────────────────
    // 1. 虚拟主机白名单（IdentityMiddleware）
    // 2. WAF 严格按需：本示例显式开启，演示内置 5 检测器拦截
    let cfg = ServerConfig::new()
        .with_allowed_hosts(["localhost", "127.0.0.1"])
        .with_waf(true);

    let server = match ProtocolServer::try_with_config(app, cfg) {
        Ok(s) => s,
        Err(e) => {
            eprintln!("server init failed: {e}");
            return ExitCode::from(3);
        }
    };

    // 2. 反向代理路由：/api/* → 后端 127.0.0.1:19090（HTTP→HTTP 端口转发）
    server.add_proxy_route(
        "/api",
        vec![("127.0.0.1:19090".to_string(), 1)],
        LoadBalanceStrategy::RoundRobin,
    );

    // 3. 反向代理路由：/proxy/* → 后端 127.0.0.1:19091（端口协议转换转发）
    //    不同端口、不同后端，模拟协议转换场景
    server.add_proxy_route(
        "/proxy",
        vec![("127.0.0.1:19091".to_string(), 1)],
        LoadBalanceStrategy::RoundRobin,
    );

    // 4. IP 黑名单：通过 zenith-forward 白名单实现 fail-closed
    //    白名单只允许 127.0.0.1，拒绝其他所有 IP（L4 转发层 SSRF 防护）
    let whitelist = AddressWhitelist::new()
        .with_allowed_ip(std::net::IpAddr::from([127u8, 0, 0, 1]));
    // ForwardEngine 仅演示白名单配置，未接入实际转发链路
    let _forward_engine = zenith_stack::forward::ForwardEngine::new()
        .with_whitelist(whitelist);

    println!("[full_feature] bind={bind}");
    println!("[full_feature] WAF=enabled (5 detectors: SQLi/XSS/SSRF/cmd/path)");
    println!("[full_feature] proxy /api/* → 127.0.0.1:19090 (RoundRobin)");
    println!("[full_feature] proxy /proxy/* → 127.0.0.1:19091 (RoundRobin, port-convert)");
    println!("[full_feature] allowed_hosts=localhost,127.0.0.1 (Host whitelist → 421)");
    println!("[full_feature] IP whitelist=127.0.0.1 (ForwardEngine L4 SSRF block)");
    println!("[full_feature] Ctrl+C to exit");

    let server_arc = Arc::new(server);

    zenith_stack::rt::block_on(async move {
        let listener = match tokio::net::TcpListener::bind(bind).await {
            Ok(l) => l,
            Err(e) => {
                eprintln!("[full_feature] bind failed: {e}");
                return;
            }
        };
        println!("[full_feature] listening on {}", listener.local_addr().unwrap_or(bind));

        let conn_count = Arc::new(AtomicU32::new(0));

        loop {
            let (stream, peer) = match listener.accept().await {
                Ok((s, p)) => (s, p),
                Err(e) => {
                    eprintln!("[full_feature] accept error: {e}");
                    continue;
                }
            };

            // 连接数限制：超过上限时拒绝新连接
            let prev = conn_count.fetch_add(1, Ordering::AcqRel);
            if prev >= MAX_CONNECTIONS {
                conn_count.fetch_sub(1, Ordering::AcqRel);
                eprintln!(
                    "[full_feature] connection limit ({MAX_CONNECTIONS}) reached, dropping {peer}"
                );
                continue;
            }

            let server_clone = Arc::clone(&server_arc);
            let conn_count_clone = Arc::clone(&conn_count);
            zenith_stack::rt::spawn(async move {
                let std_stream = match stream.into_std() {
                    Ok(s) => {
                        let _ = s.set_nonblocking(false);
                        s
                    }
                    Err(e) => {
                        eprintln!("[full_feature] stream convert: {e}");
                        conn_count_clone.fetch_sub(1, Ordering::AcqRel);
                        return;
                    }
                };
                let _ = tokio::task::spawn_blocking(move || {
                    let _ = server_clone.serve_std_tcp_conn(std_stream, peer, None);
                }).await;
                conn_count_clone.fetch_sub(1, Ordering::AcqRel);
            });
        }
    });

    ExitCode::SUCCESS
}