zenith-ebpf 0.1.0

//! eBPF 构建脚本：预编译产物管理与按需编译
//!
//! 设计原则（极端标准）：
//! 1. 预编译产物纳入版本控制（bpf/prebuilt/），零 C 编译开销
//! 2. Fail-Closed：clang 不可用时硬错误，禁止静默降级到空文件
//! 3. SHA-256 哈希校验：构建时计算哈希并生成 Rust 源文件供运行时验证
//! 4. 源码变更触发重编译：cargo:rerun-if-changed 精确依赖追踪
//!
//! 安全模型：
//! - 预编译 .o 文件在构建时计算 SHA-256 哈希
//! - 哈希值写入 generated/hashes.rs，在 xdp_prog.rs 中作为 EXPECTED_HASHES 常量
//! - 运行时 verify_integrity() 使用恒定时间比较防止时序侧信道

use sha2::Digest;
use std::path::{Path, PathBuf};
use std::process::Command;

/// eBPF 程序定义：源文件名 → 目标文件名
const PROGRAMS: &[(&str, &str, &str)] = &[
    ("xdp_main", "xdp_main.c", "xdp_main.o"),
    ("xdp_redirect", "xdp_redirect.c", "xdp_redirect.o"),
    ("xdp_stats", "xdp_stats.c", "xdp_stats.o"),
];

fn main() {
    // 非 Linux 目标：crate 内容整体 cfg 门控为空，跳过 BPF 编译与哈希生成
    if std::env::var("CARGO_CFG_TARGET_OS").map(|v| v != "linux").unwrap_or(false) {
        return;
    }
    let manifest_dir = match std::env::var("CARGO_MANIFEST_DIR") {
        Ok(dir) => dir,
        Err(_) => {
            eprintln!("cargo:error=CARGO_MANIFEST_DIR not set, cannot proceed");
            return;
        }
    };

    let bpf_dir = PathBuf::from(&manifest_dir).join("..").join("..").join("bpf");
    let prebuilt_dir = bpf_dir.join("prebuilt");
    let src_dir = bpf_dir.join("src");
    let include_dir = bpf_dir.join("include");

    // 已发布的 crate 不包含 workspace 级 bpf/ 目录：跳过 BPF 编译与校验，
    // 保留 src/generated/bpf_hashes.rs 中已有的哈希值（打包时随源码一起发布）。
    if !bpf_dir.exists() {
        eprintln!("cargo:warning=bpf/ directory not found (published crate mode), skipping BPF build");
        return;
    }

    // 确保预编译目录存在
    if let Err(e) = std::fs::create_dir_all(&prebuilt_dir) {
        eprintln!(
            "cargo:error=Failed to create prebuilt dir {}: {}",
            prebuilt_dir.display(),
            e
        );
        std::process::exit(1);
    }

    // 检查所有预编译产物是否存在且非空
    let mut need_compile = false;
    for (_name, _src, obj) in PROGRAMS {
        let obj_path = prebuilt_dir.join(obj);
        match std::fs::metadata(&obj_path) {
            Ok(meta) if meta.len() > 0 => {
                // 预编译产物存在且非空
            }
            _ => {
                need_compile = true;
                break;
            }
        }
    }

    if need_compile {
        // 编译 eBPF 程序
        eprintln!("cargo:warning=eBPF prebuilt artifacts missing or empty, compiling...");
        if let Err(e) = compile_all(&src_dir, &prebuilt_dir, &include_dir) {
            eprintln!("cargo:error={}", e);
            // 检查是否允许空字节码（测试场景）
            let allow_missing = std::env::var("ZENITH_ALLOW_MISSING_BPF")
                .map(|v| v == "1" || v == "true")
                .unwrap_or(false);
            if !allow_missing {
                eprintln!(
                    "cargo:error=eBPF compilation failed. Set ZENITH_ALLOW_MISSING_BPF=1 \
                     only for development/testing without kernel."
                );
                std::process::exit(1);
            } else {
                eprintln!("cargo:warning=Creating empty placeholders (ZENITH_ALLOW_MISSING_BPF=1 set)");
                for (_name, _src, obj) in PROGRAMS {
                    let obj_path = prebuilt_dir.join(obj);
                    if !obj_path.exists() {
                        let _ = std::fs::write(&obj_path, &[] as &[u8]);
                    }
                }
            }
        }
    }

    // ELF 完整性校验（fail-closed：校验失败立即硬错误，禁止静默降级）
    // 防止编译产物损坏或被篡改导致运行时加载失败被掩盖
    if let Err(e) = verify_elf_integrity(&prebuilt_dir) {
        eprintln!("cargo:error=eBPF ELF 完整性校验失败: {}", e);
        // 仅在非 allow_missing 模式下硬失败
        let allow_missing = std::env::var("ZENITH_ALLOW_MISSING_BPF")
            .map(|v| v == "1" || v == "true")
            .unwrap_or(false);
        if !allow_missing {
            std::process::exit(1);
        }
    }

    // 计算 SHA-256 哈希并生成 Rust 源文件
    if let Err(e) = generate_hashes(&prebuilt_dir, &manifest_dir) {
        eprintln!("cargo:error=Failed to generate hash file: {}", e);
        std::process::exit(1);
    }

    // 重新编译触发器
    println!("cargo:rerun-if-changed={}", src_dir.display());
    println!("cargo:rerun-if-changed={}", include_dir.display());
    println!("cargo:rerun-if-changed={}", prebuilt_dir.display());
    println!("cargo:rerun-if-env-changed=ZENITH_ALLOW_MISSING_BPF");
    println!("cargo:rerun-if-env-changed=ZENITH_BPF_STRICT_CHECK");
}

/// eBPF ELF 完整性校验
///
/// 极端严格标准：构建期即校验编译产物结构，防止：
/// - ELF magic 损坏（非 0x7f 'E' 'L' 'F'）
/// - 文件过小（< 64 字节，ELF 头都不完整）
/// - 缺少关键 section（.maps、xdp）
///
/// 这些校验能在构建期就暴露 clang 编译异常，避免运行时加载失败被单元测试静默掩盖。
fn verify_elf_integrity(prebuilt_dir: &Path) -> Result<(), String> {
    const ELF_MAGIC: [u8; 4] = [0x7f, b'E', b'L', b'F'];
    const MIN_ELF_SIZE: usize = 64; // ELF64 头部最小 64 字节

    for (_name, _src, obj) in PROGRAMS {
        let obj_path = prebuilt_dir.join(obj);
        let data = std::fs::read(&obj_path)
            .map_err(|e| format!("读取 {} 失败: {}", obj_path.display(), e))?;

        if data.is_empty() {
            // 空文件由 ZENITH_ALLOW_MISSING_BPF 创建，跳过校验
            continue;
        }

        // 校验 1：ELF magic
        if data.len() < 4 || data[..4] != ELF_MAGIC {
            return Err(format!(
                "{} ELF magic 错误：期望 {:?}，实际 {:?}",
                obj,
                ELF_MAGIC,
                if data.len() >= 4 { &data[..4] } else { &data }
            ));
        }

        // 校验 2：最小文件大小
        if data.len() < MIN_ELF_SIZE {
            return Err(format!(
                "{} 文件过小：{} 字节 < {} 字节（ELF64 头部最小尺寸）",
                obj,
                data.len(),
                MIN_ELF_SIZE
            ));
        }

        // 校验 3：ELF class 必须为 64 位（e_ident[EI_CLASS] = 2）
        // 偏移 4 是 EI_CLASS，值为 2 表示 ELF64
        if data[4] != 2 {
            return Err(format!(
                "{} ELF class 错误：期望 2 (ELF64)，实际 {}",
                obj, data[4]
            ));
        }

        // 校验 4：eBPF 目标机器类型（EM_BPF = 247）
        // ELF64 header 中 e_machine 位于偏移 18，2 字节小端
        if data.len() >= 20 {
            let e_machine = u16::from_le_bytes([data[18], data[19]]);
            if e_machine != 247 {
                return Err(format!(
                    "{} e_machine 错误：期望 247 (EM_BPF)，实际 {}。\
                     可能 clang -target bpf 未正确设置",
                    obj, e_machine
                ));
            }
        }

        eprintln!("cargo:warning=ELF 校验通过: {} ({} bytes, ELF64, EM_BPF)", obj, data.len());
    }
    Ok(())
}

/// 编译所有 eBPF C 源文件
fn compile_all(src_dir: &Path, out_dir: &Path, include_dir: &Path) -> Result<(), String> {
    let clang = find_clang()?;

    for (_name, src, obj) in PROGRAMS {
        let src_path = src_dir.join(src);
        let obj_path = out_dir.join(obj);

        if !src_path.exists() {
            return Err(format!("Source file not found: {}", src_path.display()));
        }

        let status = Command::new(&clang)
            .args([
                "-O2",
                "-target",
                "bpf",
                "-g",
                "-c",
                src_path.to_str().ok_or("Invalid source path")?,
                "-o",
                obj_path.to_str().ok_or("Invalid output path")?,
                "-I",
                include_dir.to_str().ok_or("Invalid include path")?,
                "-I",
                "/usr/include",
                "-Wall",
                "-Werror",
                "-fno-stack-protector",
                "-Wno-macro-redefined",
                "-Wno-unused-value",
                "-Wno-pointer-to-int-cast",
            ])
            .status()
            .map_err(|e| format!("Failed to execute clang: {}", e))?;

        if !status.success() {
            return Err(format!(
                "clang compilation failed for {} (exit code: {:?})",
                src,
                status.code()
            ));
        }

        eprintln!("cargo:warning=Compiled {} successfully", obj);
    }

    Ok(())
}

/// 查找 clang 编译器
fn find_clang() -> Result<PathBuf, String> {
    // 检查环境变量
    if let Ok(clang_path) = std::env::var("CLANG_PATH") {
        let p = PathBuf::from(&clang_path);
        if p.exists() {
            return Ok(p);
        }
    }

    // 尝试 clang 命令
    match Command::new("clang").arg("--version").output() {
        Ok(output) if output.status.success() => Ok(PathBuf::from("clang")),
        _ => {
            // 尝试 clang-18, clang-17 等
            for ver in ["clang-18", "clang-17", "clang-16", "clang-15", "clang-14"] {
                if Command::new(ver).arg("--version").output().map(|o| o.status.success()).unwrap_or(false) {
                    return Ok(PathBuf::from(ver));
                }
            }
            Err("clang not found. Please install clang (>=14) or set CLANG_PATH.".to_string())
        }
    }
}

/// 计算预编译产物的 SHA-256 哈希并生成 Rust 源文件
fn generate_hashes(prebuilt_dir: &Path, manifest_dir: &str) -> Result<(), String> {
    let out_dir = PathBuf::from(manifest_dir).join("src").join("generated");
    std::fs::create_dir_all(&out_dir)
        .map_err(|e| format!("Failed to create generated dir: {}", e))?;

    let hashes_path = out_dir.join("bpf_hashes.rs");

    let mut hashes = Vec::new();
    for (name, _src, obj) in PROGRAMS {
        let obj_path = prebuilt_dir.join(obj);
        let hash = if obj_path.exists() {
            compute_sha256(&obj_path)?
        } else {
            // 文件不存在时使用空字符串哈希（开发模式）
            compute_sha256_bytes(&[])?
        };
        hashes.push((name, hash));
    }

    // 生成 Rust 源文件
    let mut content = String::with_capacity(512);
    content.push_str("// Auto-generated SHA-256 hashes for eBPF prebuilt bytecode\n");
    content.push_str("// Generated by build.rs - DO NOT EDIT\n\n");
    content.push_str("/// 预编译 eBPF 字节码的期望 SHA-256 哈希值表（程序名 → 哈希）。\n");
    content.push_str("pub const EXPECTED_HASHES: &[(&str, &str)] = &[\n");
    for (name, hash) in &hashes {
        content.push_str(&format!("    (\"{}\", \"{}\"),\n", name, hash));
    }
    content.push_str("];\n");

    // 仅当生成内容与现有文件不同步时才落盘重写，避免每次增量构建都
    // 修改 src/generated/bpf_hashes.rs 的 mtime/内容，弄脏 git 工作区
    match std::fs::read_to_string(&hashes_path) {
        Ok(existing) if existing == content => {
            eprintln!("cargo:warning=SHA-256 hashes unchanged, skip rewriting {}", hashes_path.display());
        }
        _ => {
            std::fs::write(&hashes_path, &content)
                .map_err(|e| format!("Failed to write hash file: {}", e))?;
            eprintln!("cargo:warning=Generated SHA-256 hashes in {}", hashes_path.display());
        }
    }
    Ok(())
}

/// 计算文件的 SHA-256 哈希
fn compute_sha256(path: &Path) -> Result<String, String> {
    let data = std::fs::read(path)
        .map_err(|e| format!("Failed to read {}: {}", path.display(), e))?;
    compute_sha256_bytes(&data)
}

/// 计算字节数组的 SHA-256 哈希（纯 Rust sha2 实现，不依赖外部命令）
fn compute_sha256_bytes(data: &[u8]) -> Result<String, String> {
    let digest = sha2::Sha256::digest(data);
    let mut hex = String::with_capacity(digest.len() * 2);
    for byte in digest {
        hex.push_str(&format!("{:02x}", byte));
    }
    Ok(hex)
}