Mantis/mantis/build.rs

use std::env;
use std::fs;
use std::io::{BufRead as _, BufReader};
use std::path::PathBuf;
use std::process::{Child, Command, Stdio};
use std::time::SystemTime;

use cargo_metadata::{Artifact, CompilerMessage, Message, Metadata, MetadataCommand, Package, Target};

fn main() {
    let suricata_eve_socket = PathBuf::from("/").join("tmp").join("suricata-alerts.sock");

    let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
    let artifact_dir = manifest_dir.join("static").join("artifacts");
    let rule_dir = manifest_dir.join("static").join("rules");
    let db_dir = manifest_dir.join("static").join("db");
    let csv_dir = manifest_dir.parent().unwrap().join("records");

    let lib_dir = manifest_dir.parent().unwrap().join("lib");
    let onnxruntime_dir = lib_dir.join("onnxruntime").join("libonnxruntime.so");
    let ingress_edpf_dir = lib_dir.join("ebpf").join("mantis-ingress");
    let egress_edpf_dir = lib_dir.join("ebpf").join("mantis-egress");

    let static_web = manifest_dir.join("static").join("web");
    let project_name = manifest_dir.file_name().unwrap().to_string_lossy().into_owned();
    let frontend_dir = manifest_dir
        .parent()
        .unwrap()
        .join(format!("{}-frontend", project_name));

    println!("cargo:rustc-env=ARTIFACTCS_PATH={}", artifact_dir.display());
    println!("cargo:rustc-env=CSV_RECORD_PATH={}", csv_dir.display());
    println!("cargo:rustc-env=ONNXRUNTIME_PATH={}", onnxruntime_dir.display());
    println!("cargo:rustc-env=INGRESS_PATH={}", ingress_edpf_dir.display());
    println!("cargo:rustc-env=EGRESS_PATH={}", egress_edpf_dir.display());
    println!("cargo:rustc-env=RULE_PATH={}", rule_dir.display());
    println!("cargo:rustc-env=DB_PATH={}", db_dir.display());
    println!("cargo:rustc-env=RULE_EVE_PATH={}", suricata_eve_socket.display());

    for item in &[
        "src",
        "public",
        "package.json",
        "package-lock.json",
        "next.config.js",
        "tailwind.config.js",
        "postcss.config.js",
        "tsconfig.json",
    ] {
        println!("cargo:rerun-if-changed={}", frontend_dir.join(item).display());
    }

    if env::var_os("SKIP_EBPF_BUILD").is_some() {
        for path in &[&ingress_edpf_dir, &egress_edpf_dir] {
            if !path.exists() {
                fs::write(path, []).unwrap_or_else(|e| panic!("cannot write stub {path:?}: {e}"));
            }
        }
        return;
    }
    build_ingress_ebpf(&ingress_edpf_dir);
    build_egress_ebpf(&egress_edpf_dir);
    build_frontend(&frontend_dir, &static_web);
}

fn build_ingress_ebpf(dst: &PathBuf) {
    let Metadata { packages, .. } = MetadataCommand::new().no_deps().exec().unwrap();
    let ebpf_package = packages
        .into_iter()
        .find(|Package { name, .. }| **name == "ingress-ebpf")
        .unwrap();

    let out_dir = env::var_os("OUT_DIR").unwrap();
    let out_dir = PathBuf::from(out_dir);

    let endian = env::var_os("CARGO_CFG_TARGET_ENDIAN").unwrap();
    let target = if endian == "big" {
        "bpfeb"
    } else if endian == "little" {
        "bpfel"
    } else {
        panic!("unsupported endian={:?}", endian)
    };

    let arch = env::var_os("CARGO_CFG_TARGET_ARCH").unwrap();
    let target = format!("{target}-unknown-none");

    let Package { manifest_path, .. } = ebpf_package;
    let ebpf_dir = manifest_path.parent().unwrap();

    println!("cargo:rerun-if-changed={}", ebpf_dir.as_str());

    let mut cmd = Command::new("cargo");
    cmd.args([
        "build",
        "-Z",
        "build-std=core",
        "--bins",
        "--message-format=json",
        "--release",
        "--target",
        &target,
    ]);

    cmd.env("CARGO_CFG_BPF_TARGET_ARCH", arch);

    for key in ["RUSTUP_TOOLCHAIN", "RUSTC", "RUSTC_WORKSPACE_WRAPPER"] {
        cmd.env_remove(key);
    }
    cmd.current_dir(ebpf_dir);

    let ebpf_target_dir = out_dir.join("../ingress-ebpf");
    cmd.arg("--target-dir").arg(&ebpf_target_dir);

    let mut child = cmd
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .unwrap_or_else(|err| panic!("failed to spawn {cmd:?}: {err}"));
    let Child { stdout, stderr, .. } = &mut child;

    let stderr = stderr.take().unwrap();
    let stderr = BufReader::new(stderr);
    let stderr = std::thread::spawn(move || {
        for line in stderr.lines() {
            let line = line.unwrap();
            println!("{line}");
        }
    });

    let stdout = stdout.take().unwrap();
    let stdout = BufReader::new(stdout);
    let mut executables = Vec::new();
    for message in Message::parse_stream(stdout) {
        #[allow(clippy::collapsible_match)]
        match message.expect("valid JSON") {
            Message::CompilerArtifact(Artifact {
                executable,
                target: Target { name, .. },
                ..
            }) => {
                if let Some(executable) = executable {
                    executables.push((name, executable.into_std_path_buf()));
                }
            }
            Message::CompilerMessage(CompilerMessage { message, .. }) => {
                for line in message.rendered.unwrap_or_default().split('\n') {
                    println!("{line}");
                }
            }
            Message::TextLine(line) => {
                println!("{line}");
            }
            _ => {}
        }
    }

    let status = child
        .wait()
        .unwrap_or_else(|err| panic!("failed to wait for {cmd:?}: {err}"));
    assert_eq!(status.code(), Some(0), "{cmd:?} failed: {status:?}");

    stderr.join().map_err(std::panic::resume_unwind).unwrap();

    for (_name, binary) in executables {
        let _: u64 = fs::copy(&binary, dst).unwrap_or_else(|err| panic!("failed to copy {binary:?} to {dst:?}: {err}"));
    }
}

fn build_egress_ebpf(dst: &PathBuf) {
    let Metadata { packages, .. } = MetadataCommand::new().no_deps().exec().unwrap();
    let ebpf_package = packages
        .into_iter()
        .find(|Package { name, .. }| **name == "egress-ebpf")
        .unwrap();

    let out_dir = env::var_os("OUT_DIR").unwrap();
    let out_dir = PathBuf::from(out_dir);

    let endian = env::var_os("CARGO_CFG_TARGET_ENDIAN").unwrap();
    let target = if endian == "big" {
        "bpfeb"
    } else if endian == "little" {
        "bpfel"
    } else {
        panic!("unsupported endian={:?}", endian)
    };

    let arch = env::var_os("CARGO_CFG_TARGET_ARCH").unwrap();
    let target = format!("{target}-unknown-none");

    let Package { manifest_path, .. } = ebpf_package;
    let ebpf_dir = manifest_path.parent().unwrap();

    println!("cargo:rerun-if-changed={}", ebpf_dir.as_str());

    let mut cmd = Command::new("cargo");
    cmd.args([
        "build",
        "-Z",
        "build-std=core",
        "--bins",
        "--message-format=json",
        "--release",
        "--target",
        &target,
    ]);

    cmd.env("CARGO_CFG_BPF_TARGET_ARCH", arch);
    cmd.env("CARGO_TERM_COLOR", "always");

    for key in ["RUSTUP_TOOLCHAIN", "RUSTC", "RUSTC_WORKSPACE_WRAPPER"] {
        cmd.env_remove(key);
    }
    cmd.current_dir(ebpf_dir);

    let ebpf_target_dir = out_dir.join("../egress-ebpf");
    cmd.arg("--target-dir").arg(&ebpf_target_dir);

    let mut child = cmd
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .unwrap_or_else(|err| panic!("failed to spawn {cmd:?}: {err}"));
    let Child { stdout, stderr, .. } = &mut child;

    let stderr = stderr.take().unwrap();
    let stderr = BufReader::new(stderr);
    let stderr = std::thread::spawn(move || {
        for line in stderr.lines() {
            let line = line.unwrap();
            println!("{line}");
        }
    });

    let stdout = stdout.take().unwrap();
    let stdout = BufReader::new(stdout);
    let mut executables = Vec::new();
    for message in Message::parse_stream(stdout) {
        #[allow(clippy::collapsible_match)]
        match message.expect("valid JSON") {
            Message::CompilerArtifact(Artifact {
                executable,
                target: Target { name, .. },
                ..
            }) => {
                if let Some(executable) = executable {
                    executables.push((name, executable.into_std_path_buf()));
                }
            }
            Message::CompilerMessage(CompilerMessage { message, .. }) => {
                for line in message.rendered.unwrap_or_default().split('\n') {
                    println!("{line}");
                }
            }
            Message::TextLine(line) => {
                println!("{line}");
            }
            _ => {}
        }
    }

    let status = child
        .wait()
        .unwrap_or_else(|err| panic!("failed to wait for {cmd:?}: {err}"));
    assert_eq!(status.code(), Some(0), "{cmd:?} failed: {status:?}");

    stderr.join().map_err(std::panic::resume_unwind).unwrap();

    for (_name, binary) in executables {
        let _: u64 = fs::copy(&binary, dst).unwrap_or_else(|err| panic!("failed to copy {binary:?} to {dst:?}: {err}"));
    }
}

fn build_frontend(frontend_dir: &PathBuf, static_dir: &PathBuf) {
    if !frontend_dir.exists() {
        panic!("Frontend directory {:?} does not exist", frontend_dir);
    }

    let out_dir = frontend_dir.join("out");
    if !needs_frontend_rebuild(frontend_dir, &out_dir, static_dir) {
        return;
    }

    let status = Command::new("npm")
        .arg("install")
        .current_dir(frontend_dir)
        .status()
        .unwrap_or_else(|err| panic!("failed to run npm install: {err}"));
    if !status.success() {
        panic!("npm install failed with exit code: {:?}", status.code());
    }

    let status = Command::new("npx")
        .args(["next", "build"])
        .current_dir(frontend_dir)
        .status()
        .unwrap_or_else(|err| panic!("failed to run next build: {err}"));
    if !status.success() {
        panic!("next build failed with exit code: {:?}", status.code());
    }

    if static_dir.exists() {
        fs::remove_dir_all(static_dir).unwrap_or_else(|err| panic!("failed to remove {:?}: {err}", static_dir));
    }
    fs::create_dir_all(static_dir).unwrap_or_else(|err| panic!("failed to create {:?}: {err}", static_dir));
    copy_dir_all(&out_dir, static_dir).unwrap_or_else(|err| panic!("failed to copy frontend build: {err}"));
}

fn needs_frontend_rebuild(frontend_dir: &PathBuf, out_dir: &PathBuf, static_dir: &PathBuf) -> bool {
    if !out_dir.exists() {
        return true;
    }

    if !static_dir.exists() {
        return true;
    }

    let out_modified = match fs::metadata(out_dir).and_then(|m| m.modified()) {
        Ok(time) => time,
        Err(_) => {
            return true;
        }
    };

    let static_modified = match fs::metadata(static_dir).and_then(|m| m.modified()) {
        Ok(time) => time,
        Err(_) => {
            return true;
        }
    };

    let essential_items = [
        "src",
        "public",
        "package.json",
        "next.config.js",
        "tailwind.config.js",
        "postcss.config.js",
        "tsconfig.json",
        "package-lock.json",
    ];

    for item_name in essential_items {
        let item_path = frontend_dir.join(item_name);
        if !item_path.exists() {
            continue;
        }

        let item_modified = match get_dir_last_modified(&item_path) {
            Some(time) => time,
            None => continue,
        };

        if item_modified > out_modified {
            return true;
        }
    }

    if out_modified > static_modified {
        return true;
    }

    false
}

fn get_dir_last_modified(path: &PathBuf) -> Option<SystemTime> {
    if path.is_file() {
        return fs::metadata(path).and_then(|m| m.modified()).ok();
    }

    if path.is_dir() {
        let mut latest = fs::metadata(path).and_then(|m| m.modified()).ok()?;

        if let Ok(entries) = fs::read_dir(path) {
            for entry in entries.flatten() {
                if let Some(modified) = get_dir_last_modified(&entry.path()) {
                    if modified > latest {
                        latest = modified;
                    }
                }
            }
        }

        return Some(latest);
    }

    None
}

fn copy_dir_all(src: &PathBuf, dst: &PathBuf) -> std::io::Result<()> {
    for entry in fs::read_dir(src)? {
        let entry = entry?;
        let file_type = entry.file_type()?;
        let src_path = entry.path();
        let dst_path = dst.join(entry.file_name());

        if file_type.is_dir() {
            fs::create_dir_all(&dst_path)?;
            copy_dir_all(&src_path, &dst_path)?;
        } else {
            fs::copy(&src_path, &dst_path)?;
        }
    }
    Ok(())
}