src/starnix/modules/usbfs/testing/trace_generator.py

#!/usr/bin/env fuchsia-vendored-python
# Copyright 2026 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Host runner for recording usbfs golden traces in QEMU and generating scenarios.rs."""

import argparse
import io
import json
import os
import subprocess
import sys
import tarfile
import tempfile
from collections.abc import Sequence


def _rust_raw_str(content: str) -> str:
    hashes_count = 1
    while f'"{("#" * hashes_count)}' in content:
        hashes_count += 1
    hashes = "#" * hashes_count
    return f'r{hashes}"{content}"{hashes}'


def _step_uses_arch32(step: dict) -> bool:
    if step.get("arch32") is True:
        return True
    inner = step.get("inner")
    if isinstance(inner, dict) and _step_uses_arch32(inner):
        return True
    return False


def generate_rs(
    output_path: str,
    trace_dir: str | None,
    divergences_path: str | None,
    names: Sequence[str],
):
    divergences_content = '{"divergences": []}'
    if divergences_path and os.path.exists(divergences_path):
        with open(divergences_path, "r", encoding="utf-8") as df:
            divergences_content = df.read().strip()

    sorted_names = sorted(names)
    traces: list[tuple[str, str, bool]] = []
    for name in sorted_names:
        trace_content = ""
        if trace_dir:
            trace_path = os.path.join(trace_dir, f"{name}.json")
            if os.path.exists(trace_path):
                with open(trace_path, "r", encoding="utf-8") as tf:
                    trace_content = tf.read().strip()
        if not trace_content:
            raise ValueError(
                f"Trace content for scenario '{name}' is empty or missing. "
                f"Did you forget to run trace_generator.py? "
                f"See src/starnix/modules/usbfs/testing/README.md for instructions."
            )
        trace_obj = json.loads(trace_content)
        steps = trace_obj.get("steps")
        if not isinstance(steps, list) or not steps:
            raise ValueError(
                f"Recorded trace for scenario '{name}' has no steps. "
                f"Did you forget to run trace_generator.py? "
                f"See src/starnix/modules/usbfs/testing/README.md for instructions."
            )
        uses_arch32 = any(
            isinstance(s, dict)
            and isinstance(s.get("step"), dict)
            and _step_uses_arch32(s["step"])
            for s in steps
        )
        traces.append((name, trace_content, uses_arch32))

    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        f.write("// Copyright 2026 The Fuchsia Authors. All rights reserved.\n")
        f.write(
            "// Use of this source code is governed by a BSD-style license that can be\n"
        )
        f.write("// found in the LICENSE file.\n\n")
        f.write("// Generated by trace_generator.py. DO NOT EDIT.\n\n")
        f.write(
            f"pub const DIVERGENCES_JSON: &str = {_rust_raw_str(divergences_content)};\n\n"
        )
        f.write("pub const SCENARIOS: &[(&str, &str)] = &[\n")
        for name, trace_content, _ in traces:
            f.write(f'    ("{name}", {_rust_raw_str(trace_content)}),\n')
        f.write("];\n\n")
        f.write("#[macro_export]\n")
        f.write(
            "/// Generates test cases replaying golden traces.\n"
            "///\n"
            "/// Optional caller test attributes must be fully qualified (e.g. `#[serial_test::serial]`).\n"
        )
        f.write("macro_rules! usbfs_trace_tests {\n")
        f.write("    ($replayer:path $(, $(#[$attr:meta])* )* $(,)?) => {\n")
        if traces:
            std_traces = [name for name, _, arch32 in traces if not arch32]
            arch32_traces = [name for name, _, arch32 in traces if arch32]
            f.write("        #[cfg(test)]\n")
            f.write("        mod generated_usbfs_trace_tests {\n")
            f.write("            use super::*;\n")
            f.write("            use ::test_case::test_case;\n\n")
            if std_traces:
                for name in std_traces:
                    f.write(f'            #[test_case("{name}"; "{name}")]\n')
                f.write("            $( $(#[$attr])* )*\n")
                f.write("            fn test_replay_trace(name: &str) {\n")
                f.write(
                    "                let (_, json_data) = $crate::scenarios::SCENARIOS\n"
                    "                    .iter()\n"
                    "                    .find(|(n, _)| *n == name)\n"
                    '                    .expect("scenario trace exists");\n'
                    "                $replayer(name, json_data);\n"
                )
                f.write("            }\n")
            if arch32_traces:
                if std_traces:
                    f.write("\n")
                for name in arch32_traces:
                    f.write(f'            #[test_case("{name}"; "{name}")]\n')
                f.write(
                    '            #[cfg_attr(not(target_arch = "aarch64"), ignore)]\n'
                )
                f.write("            $( $(#[$attr])* )*\n")
                f.write(
                    "            fn test_replay_trace_arch32(name: &str) {\n"
                )
                f.write(
                    "                let (_, json_data) = $crate::scenarios::SCENARIOS\n"
                    "                    .iter()\n"
                    "                    .find(|(n, _)| *n == name)\n"
                    '                    .expect("scenario trace exists");\n'
                    "                $replayer(name, json_data);\n"
                )
                f.write("            }\n")
            f.write("        }\n")
        else:
            f.write("        use ::test_case::test_case as _;\n")
            f.write("        const _: fn(&str, &str) = $replayer;\n")
        f.write("    };\n")
        f.write("}\n")


def find_build_dir(fuchsia_dir: str) -> str:
    fx_build_dir_file = os.path.join(fuchsia_dir, ".fx-build-dir")
    if os.path.exists(fx_build_dir_file):
        with open(fx_build_dir_file, "r") as f:
            rel = f.read().strip()
            if rel:
                return os.path.join(fuchsia_dir, rel)
    return os.path.join(fuchsia_dir, "out/default")


def add_bytes_to_tar(
    tf: tarfile.TarFile, arcname: str, data: bytes, mode: int = 0o644
):
    info = tarfile.TarInfo(name=arcname)
    info.size = len(data)
    info.mode = mode
    tf.addfile(info, io.BytesIO(data))


def record_in_qemu(
    scenarios_filter: Sequence[str],
    no_build: bool = False,
    recorder_bin_override: str | None = None,
) -> int:
    script_dir = os.path.dirname(os.path.abspath(__file__))
    fuchsia_dir = os.path.abspath(os.path.join(script_dir, "../../../../.."))
    build_dir = find_build_dir(fuchsia_dir)

    if not no_build and not recorder_bin_override:
        fx_bin = os.path.join(fuchsia_dir, ".jiri_root/bin/fx")
        target = "//src/starnix/modules/usbfs/testing:usbfs_trace_recorder(//build/toolchain:linux_x64)"
        print(f"[trace_generator] Building {target}...")
        subprocess.run([fx_bin, "build", target], cwd=fuchsia_dir, check=True)

    recorder_bin = recorder_bin_override or os.path.join(
        build_dir, "linux_x64/usbfs_trace_recorder"
    )
    if not os.path.exists(recorder_bin):
        raise FileNotFoundError(
            f"Missing recorder binary at {recorder_bin}. Run fx build first."
        )

    qemu_bin = os.path.join(
        fuchsia_dir,
        "prebuilt",
        "third_party/qemu/linux-x64/bin/qemu-system-x86_64",
    )
    debian_img_dir = os.path.join(
        fuchsia_dir,
        "prebuilt",
        "virtualization/packages/debian_guest/images/x64",
    )
    vmlinuz = os.path.join(debian_img_dir, "vmlinuz")
    initrd = os.path.join(debian_img_dir, "initrd.img")
    rootfs = os.path.join(debian_img_dir, "rootfs.qcow2")

    list_path = os.path.join(script_dir, "scenarios_list.json")
    with open(list_path, "r") as f:
        all_scenarios: list[str] = json.load(f)

    selected = scenarios_filter if scenarios_filter else all_scenarios
    if not selected:
        print("[trace_generator] No scenarios to record.")
        return 0

    scenarios_dir = os.path.join(script_dir, "scenarios")
    generated_dir = os.path.join(script_dir, "generated")
    os.makedirs(generated_dir, exist_ok=True)

    if not scenarios_filter:
        expected_files = {f"{name}.json" for name in all_scenarios}
        for existing in os.listdir(generated_dir):
            if existing.endswith(".json") and existing not in expected_files:
                os.remove(os.path.join(generated_dir, existing))

    with tempfile.TemporaryDirectory(prefix="usbfs_trace_gen_") as tmpdir:
        in_tar_path = os.path.join(tmpdir, "in.tar")
        out_tar_path = os.path.join(tmpdir, "out.tar")

        run_sh = """#!/bin/sh
set -e
mount -t proc proc /proc 2>/dev/null || true
mount -t sysfs sys /sys 2>/dev/null || true
mount -t devtmpfs dev /dev 2>/dev/null || true
mkdir -p /sys/kernel/config /dev/ffs /tmp/out
mount -t configfs none /sys/kernel/config 2>/dev/null || true
modprobe dummy_hcd 2>/dev/null || true
modprobe usb_f_fs 2>/dev/null || true
ip link set lo up 2>/dev/null || ifconfig lo up 2>/dev/null || true
/tmp/in/usbfs_trace_recorder --scenarios-dir /tmp/in/scenarios --list /tmp/in/scenarios_list.json --out-dir /tmp/out
echo OK > /tmp/out/_RECORD_OK
tar -cf /dev/vdc -C /tmp/out .
sync
echo USBFS_TRACE_RECORDER_SUCCESS
sleep 0.2
echo o > /proc/sysrq-trigger
"""

        with tarfile.open(in_tar_path, "w") as tf:
            tf.add(recorder_bin, arcname="usbfs_trace_recorder")
            add_bytes_to_tar(
                tf,
                "scenarios_list.json",
                json.dumps(list(selected), indent=2).encode("utf-8"),
                mode=0o644,
            )
            add_bytes_to_tar(tf, "run.sh", run_sh.encode("utf-8"), mode=0o755)
            for name in selected:
                s_path = os.path.join(scenarios_dir, f"{name}.json")
                if not os.path.exists(s_path):
                    raise FileNotFoundError(
                        f"Scenario file not found: {s_path}"
                    )
                tf.add(s_path, arcname=f"scenarios/{name}.json")

        # Pad in.tar to 4096-byte boundary for virtio_blk
        in_size = os.path.getsize(in_tar_path)
        if in_size % 4096 != 0:
            with open(in_tar_path, "ab") as f:
                f.write(b"\x00" * (4096 - (in_size % 4096)))

        # Pre-allocate 10 MiB zero-filled out.tar for /dev/vdc
        with open(out_tar_path, "wb") as f:
            f.truncate(10 * 1024 * 1024)

        init_cmd = (
            "mkdir$IFS-p$IFS/tmp/in&&"
            "tar$IFS-xf$IFS/dev/vdb$IFS-C$IFS/tmp/in&&"
            "/tmp/in/run.sh"
        )
        qemu_cmd = [
            qemu_bin,
            "-enable-kvm",
            "-cpu",
            "host",
            "-smp",
            "4",
            "-m",
            "1024M",
            "-nographic",
            "-no-reboot",
            "-kernel",
            vmlinuz,
            "-initrd",
            initrd,
            "-drive",
            f"file={rootfs},format=qcow2,if=virtio,index=0,snapshot=on",
            "-drive",
            f"file={in_tar_path},format=raw,if=virtio,readonly=on,index=1",
            "-drive",
            f"file={out_tar_path},format=raw,if=virtio,index=2",
            "-append",
            f"root=/dev/vda rw console=ttyS0 quiet panic=-1 init=/bin/sh -- -c {init_cmd}",
        ]

        print(
            f"[trace_generator] Booting QEMU/KVM to record {len(selected)} scenario(s)..."
        )
        try:
            proc = subprocess.run(
                qemu_cmd,
                stdin=subprocess.DEVNULL,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                timeout=120,
                check=False,
            )
        except subprocess.TimeoutExpired as e:
            if e.stdout:
                out_str = (
                    e.stdout.decode("utf-8", errors="replace")
                    if isinstance(e.stdout, bytes)
                    else e.stdout
                )
                sys.stderr.write(out_str)
            raise RuntimeError(
                f"QEMU trace recording timed out after {e.timeout}s"
            ) from e
        has_tar_ok = False
        try:
            with tarfile.open(out_tar_path, "r:") as tf_check:
                names = tf_check.getnames()
                has_tar_ok = "./_RECORD_OK" in names or "_RECORD_OK" in names
        except Exception:
            has_tar_ok = False

        if not has_tar_ok or "USBFS_TRACE_RECORDER_SUCCESS" not in proc.stdout:
            sys.stderr.write(proc.stdout)
            raise RuntimeError(
                f"QEMU trace recording failed (exit code {proc.returncode})"
            )

        with tarfile.open(out_tar_path, "r:") as tf:
            for name in selected:
                member_name = f"./{name}.json"
                if (
                    member_name not in tf.getnames()
                    and f"{name}.json" in tf.getnames()
                ):
                    member_name = f"{name}.json"
                extracted = tf.extractfile(member_name)
                if extracted is None:
                    raise RuntimeError(
                        f"Missing {name}.json in /dev/vdc output tar"
                    )
                trace_obj = json.loads(extracted.read().decode("utf-8"))
                dest_path = os.path.join(generated_dir, f"{name}.json")
                with open(dest_path, "w") as out_f:
                    json.dump(trace_obj, out_f, indent=2)
                    out_f.write("\n")
                print(f"[trace_generator] Wrote {dest_path}")

    return 0


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Record usbfs golden traces or generate scenarios.rs."
    )
    parser.add_argument(
        "--generate-rs",
        action="store_true",
        help="Generate scenarios.rs for GN build.",
    )
    parser.add_argument(
        "--output",
        help="Output file path for --generate-rs.",
    )
    parser.add_argument(
        "--trace-dir",
        help="Directory containing generated/*.json golden trace files.",
    )
    parser.add_argument(
        "--divergences",
        help="Path to divergences.json.",
    )
    parser.add_argument(
        "--names",
        nargs="*",
        default=[],
        help="List of scenario names for --generate-rs.",
    )
    parser.add_argument(
        "--scenarios",
        nargs="*",
        default=[],
        help="Optional subset of scenarios to record in QEMU.",
    )
    parser.add_argument(
        "--no-build",
        action="store_true",
        help="Skip running fx build before launching QEMU.",
    )
    parser.add_argument(
        "--recorder-bin",
        help="Explicit path to linux_x64/usbfs_trace_recorder.",
    )
    args = parser.parse_args()

    if args.generate_rs:
        if not args.output:
            parser.error("--output is required for --generate-rs")
        generate_rs(
            output_path=args.output,
            trace_dir=args.trace_dir,
            divergences_path=args.divergences,
            names=args.names,
        )
        return 0

    return record_in_qemu(
        scenarios_filter=args.scenarios,
        no_build=args.no_build,
        recorder_bin_override=args.recorder_bin,
    )


if __name__ == "__main__":
    sys.exit(main())