sbpf-linker 0.2.0

#![expect(unused_crate_dependencies, reason = "used in test harness")]

use std::{
    collections::HashMap,
    env,
    ffi::OsString,
    fs, io,
    path::{Path, PathBuf},
    process::Command,
};

use either::Either;
use object::{File, Object as _, ObjectSection as _};
use sbpf_assembler::{
    OptimizationConfig, SbpfArch,
    astnode::{ASTNode, ROData},
    header::ProgramHeader,
    parser::Token,
};
use sbpf_common::{
    inst_param::Number,
    instruction::{AsmFormat, Instruction},
    opcode::Opcode,
};
use sbpf_linker::{ProgramOptions, byteparser::parse_bytecode};

const NO_TESTS_FILTER: &str = "__no_tests_match_this_sbpf_arch__";

const DEFAULT_STACK_FRAME_SIZE: i32 = 4096;

trait TestArch {
    const ARCH: SbpfArch;

    fn decode_instruction(
        data: &[u8],
    ) -> Result<Instruction, sbpf_common::errors::SBPFError>;

    fn arch_arg() -> String {
        format!("v{}", Self::ARCH.e_flags())
    }

    fn dump(src: &Path, dst: &Path)
    where
        Self: Sized,
    {
        sbpf_dump::<Self>(src, dst);
    }
}

struct SbpfV0;

impl TestArch for SbpfV0 {
    const ARCH: SbpfArch = SbpfArch::V0;

    fn decode_instruction(
        data: &[u8],
    ) -> Result<Instruction, sbpf_common::errors::SBPFError> {
        Instruction::from_bytes(data)
    }
}

struct SbpfV3;

impl TestArch for SbpfV3 {
    const ARCH: SbpfArch = SbpfArch::V3;

    fn decode_instruction(
        data: &[u8],
    ) -> Result<Instruction, sbpf_common::errors::SBPFError> {
        Instruction::from_bytes_sbpf_v3(data)
    }
}

fn rustc_cmd() -> Command {
    Command::new(
        env::var_os("RUSTC").unwrap_or_else(|| OsString::from("rustc")),
    )
}

fn find_binary(binary_re_str: &str) -> PathBuf {
    let binary_re = regex::Regex::new(binary_re_str).unwrap();
    let mut binary = which::which_re(binary_re).expect(binary_re_str);
    binary.next().unwrap_or_else(|| panic!("could not find {binary_re_str}"))
}

fn run_mode<A, F>(target: &str, mode: &str, sysroot: &Path, cfg: Option<F>)
where
    A: TestArch,
    F: Fn(&mut compiletest_rs::Config),
{
    let arch_arg = A::arch_arg();
    let target_rustcflags = format!(
        "-C linker={} -C link-arg=--arch={} --sysroot {}",
        env!("CARGO_BIN_EXE_sbpf-linker"),
        arch_arg,
        sysroot.display()
    );

    let llvm_filecheck = Some(find_binary(r"^FileCheck(-\d+)?$"));

    let mode = mode.parse().expect("invalid compiletest mode");
    let mut config = compiletest_rs::Config {
        target: target.to_owned(),
        target_rustcflags: Some(target_rustcflags),
        llvm_filecheck,
        mode,
        src_base: PathBuf::from(format!("tests/{mode}")),
        ..Default::default()
    };
    config.link_deps();

    if let Some(cfg) = cfg {
        cfg(&mut config);
    }

    config.filters = test_filters_for_arch::<A>(&config.src_base)
        .expect("failed to filter tests by sBPF arch");

    compiletest_rs::run_tests(&config);
}

fn sbpf_dump<A: TestArch>(src: &Path, dst: &Path) {
    let dump = render_emitted_program::<A>(src).unwrap_or_else(|err| {
        panic!("failed to render {}: {err}", src.display())
    });
    fs::write(dst, dump).unwrap_or_else(|err| {
        panic!("failed to write {}: {err}", dst.display())
    });
}

fn test_filters_for_arch<A: TestArch>(
    src_base: &Path,
) -> io::Result<Vec<String>> {
    let suite_name =
        src_base.file_name().unwrap_or_default().to_string_lossy();
    let arch_arg = A::arch_arg();
    let mut filters = Vec::new();
    collect_test_filters_for_arch(
        src_base,
        src_base,
        &suite_name,
        &mut filters,
        &arch_arg,
    )?;
    filters.sort();
    if filters.is_empty() {
        filters.push(NO_TESTS_FILTER.to_owned());
    }
    Ok(filters)
}

fn collect_test_filters_for_arch(
    src_base: &Path,
    dir: &Path,
    suite_name: &str,
    filters: &mut Vec<String>,
    arch_arg: &str,
) -> io::Result<()> {
    for entry in fs::read_dir(dir)? {
        let entry = entry?;
        let path = entry.path();
        if path.is_dir() {
            // Compiletest builds auxiliary crates only when a fixture requests them.
            if entry.file_name() != "auxiliary" {
                collect_test_filters_for_arch(
                    src_base, &path, suite_name, filters, arch_arg,
                )?;
            }
        } else if path.extension().is_some_and(|extension| extension == "rs")
            && !ignored_for_arch(&path, arch_arg)?
        {
            let relative_path = path.strip_prefix(src_base).unwrap_or(&path);
            filters.push(format!("{suite_name}/{}", relative_path.display()));
        }
    }
    Ok(())
}

fn ignored_for_arch(path: &Path, arch_arg: &str) -> io::Result<bool> {
    let contents = fs::read_to_string(path)?;
    Ok(contents.lines().any(|line| {
        line.trim_start()
            .strip_prefix("//")
            .map(str::trim_start)
            .and_then(|line| line.strip_prefix("ignore-sbpf-arch:"))
            .is_some_and(|ignored_arches| {
                ignored_arches
                    .split([',', ' ', '\t'])
                    .any(|ignored_arch| ignored_arch.trim() == arch_arg)
            })
    }))
}

#[test]
fn compile_test() {
    // Assembly fixtures live in `tests/assembly`. Each file is a tiny Rust
    // crate with compiletest directives at the top and inline `CHECK:` lines
    // at the bottom. Use `// ignore-sbpf-arch: v0` or `v3` to skip a fixture
    // for one linker arch. Run just this harness with:
    //
    // `cargo test --test tests compile_test -- --nocapture`
    //
    // or run the whole suite with `cargo test`.
    let target = "bpfel-unknown-none";
    let root_dir = env::var_os("CARGO_MANIFEST_DIR")
        .expect("could not determine the root directory of the project");
    let root_dir = Path::new(&root_dir);
    let bpf_sysroot = if let Some(bpf_sysroot) =
        env::var_os("BPFEL_SYSROOT_DIR")
    {
        PathBuf::from(bpf_sysroot)
    } else {
        let rustc_src = rustc_build_sysroot::rustc_sysroot_src(rustc_cmd())
            .expect("could not determine sysroot source directory");
        let directory = root_dir.join("target/sysroot");
        let mut cargo = Command::new(
            env::var_os("CARGO").unwrap_or_else(|| OsString::from("cargo")),
        );
        cargo.env("RUSTC_BOOTSTRAP", "1");
        match rustc_build_sysroot::SysrootBuilder::new(&directory, target)
            .cargo(cargo)
            .build_mode(rustc_build_sysroot::BuildMode::Build)
            .sysroot_config(rustc_build_sysroot::SysrootConfig::NoStd)
            .build_from_source(&rustc_src)
            .expect("failed to build sysroot")
        {
            rustc_build_sysroot::SysrootStatus::AlreadyCached => {}
            rustc_build_sysroot::SysrootStatus::SysrootBuilt => {}
        }
        directory
    };

    run_mode::<SbpfV0, _>(
        target,
        "assembly",
        &bpf_sysroot,
        Some(|cfg: &mut compiletest_rs::Config| {
            cfg.llvm_filecheck_preprocess = Some(SbpfV0::dump);
        }),
    );
    run_mode::<SbpfV3, _>(
        target,
        "assembly",
        &bpf_sysroot,
        Some(|cfg: &mut compiletest_rs::Config| {
            cfg.llvm_filecheck_preprocess = Some(SbpfV3::dump);
        }),
    );
}

// TODO: add below query methods to sbpf and update below to use them
fn render_emitted_program<A: TestArch>(path: &Path) -> anyhow::Result<String> {
    let bytes = fs::read(path)?;
    let syscall_labels = collect_syscall_labels::<A>(&bytes)?;
    let parse_result = parse_bytecode(
        &bytes,
        ProgramOptions::new(
            OptimizationConfig::enabled(),
            A::ARCH,
            DEFAULT_STACK_FRAME_SIZE,
        ),
    )?;
    let ph_count = if parse_result.prog_is_static { 1u64 } else { 3u64 };
    let rodata_base =
        parse_result.code_section.get_size() + 64 + ph_count * 56;
    let rodata_len = parse_result.data_section.get_size();

    let mut out = Vec::new();
    let rodata_nodes = parse_result.data_section.get_nodes();
    let mut rodata_labels = HashMap::new();
    let mut code_labels = HashMap::new();
    out.push(format!("rodata-count: {}", rodata_nodes.len()));

    for node in rodata_nodes {
        if let ASTNode::ROData { rodata, offset } = node {
            let label = format!("data_{offset:04x}");
            rodata_labels.insert(*offset, label.clone());
            out.push(format!("rodata-label[{offset}]: {label}"));
            out.push(format!("rodata[{offset}]: {}", render_rodata(rodata)?));
        }
    }

    let code_nodes = parse_result.code_section.get_nodes();
    for node in code_nodes {
        if let ASTNode::Label { label, offset } = node {
            code_labels.insert(*offset as i64, label.name.clone());
        }
    }

    for node in code_nodes {
        match node {
            ASTNode::Label { label, offset } => {
                out.push(format!("{offset:04x}: label {}", label.name));
            }
            ASTNode::Instruction { instruction, offset } => {
                for asm in render_instruction::<A>(
                    instruction,
                    *offset,
                    rodata_base,
                    rodata_len,
                    &rodata_labels,
                    &code_labels,
                    &syscall_labels,
                )? {
                    out.push(format!("{offset:04x}: {asm}"));
                }
            }
            _ => {}
        }
    }

    Ok(out.join("\n"))
}

fn render_instruction<A: TestArch>(
    instruction: &Instruction,
    offset: u64,
    rodata_base: u64,
    rodata_len: u64,
    rodata_labels: &HashMap<u64, String>,
    code_labels: &HashMap<i64, String>,
    syscall_labels: &HashMap<u64, String>,
) -> anyhow::Result<Vec<String>> {
    if instruction.opcode == Opcode::Call
        && let Some(label) = syscall_labels.get(&offset)
    {
        return Ok(vec![format!("call {label}")]);
    }

    if instruction.opcode == Opcode::Call
        && let Some(Either::Right(Number::Int(value) | Number::Addr(value))) =
            &instruction.imm
    {
        let target = offset as i64 + 8 + value * 8;
        if let Some(label) = code_labels.get(&target) {
            return Ok(vec![
                instruction.to_asm(AsmFormat::Default)?,
                format!("call {label}"),
            ]);
        }
    }

    if instruction.opcode == Opcode::Lddw
        && let Some(Either::Right(number)) = &instruction.imm
        && let Number::Int(value) | Number::Addr(value) = number
        && let Some(offset) =
            rodata_offset_for_lddw::<A>(*value as u64, rodata_base, rodata_len)
    {
        let dst = instruction.dst.as_ref().ok_or_else(|| {
            anyhow::anyhow!("lddw is missing a destination register")
        })?;
        let mut rendered = vec![format!("lddw r{}, rodata[{offset}]", dst.n)];
        if let Some(label) = rodata_labels.get(&offset) {
            rendered.push(format!("lddw r{}, {}", dst.n, label));
        }
        return Ok(rendered);
    }

    Ok(vec![instruction.to_asm(AsmFormat::Default)?])
}

fn rodata_offset_for_lddw<A: TestArch>(
    value: u64,
    rodata_base: u64,
    rodata_len: u64,
) -> Option<u64> {
    let rodata_vaddr =
        ProgramHeader::new_load(rodata_base, rodata_len, false, A::ARCH)
            .p_vaddr;
    (value >= rodata_vaddr && value < rodata_vaddr + rodata_len)
        .then_some(value - rodata_vaddr)
}

fn collect_syscall_labels<A: TestArch>(
    bytes: &[u8],
) -> anyhow::Result<HashMap<u64, String>> {
    let obj = File::parse(bytes)?;
    let Some(text) = obj.section_by_name(".text") else {
        return Ok(HashMap::new());
    };
    let data = text.data()?;

    let mut labels = HashMap::new();
    let mut offset = 0usize;
    while offset < data.len() {
        let instruction =
            A::decode_instruction(&data[offset..]).map_err(|err| {
                anyhow::anyhow!("failed to decode .text at {offset:#x}: {err}")
            })?;
        if instruction.opcode == Opcode::Call
            && let Some(Either::Left(identifier)) = instruction.imm
        {
            labels.insert(offset as u64, identifier);
        }
        offset += if instruction.opcode == Opcode::Lddw { 16 } else { 8 };
    }

    Ok(labels)
}

fn render_rodata(rodata: &ROData) -> anyhow::Result<String> {
    match (&rodata.args[0], &rodata.args[1]) {
        (Token::Directive(directive, _), Token::VectorLiteral(values, _)) => {
            let bytes =
                values.iter().map(ToString::to_string).collect::<Vec<_>>();
            Ok(format!("{directive} {}", bytes.join(", ")))
        }
        (Token::Directive(directive, _), Token::StringLiteral(value, _)) => {
            Ok(format!("{directive} {:?}", value))
        }
        _ => Err(anyhow::anyhow!(
            "unsupported rodata node layout for {}",
            rodata.name
        )),
    }
}