ableos/repbuild/src/main.rs

use {
    derive_more::Display,
    error_stack::{bail, report, Context, Report, Result, ResultExt},
    fatfs::{FileSystem, FormatVolumeOptions, FsOptions, ReadWriteSeek},
    std::{fmt::Display, fs::File, io, path::Path, process::Command},
};

fn main() -> Result<(), Error> {
    let mut args = std::env::args();
    args.next();

    // let disk_meta = fs::metadata("target/disk.img").unwrap();
    // let config_meta = fs::metadata("system.toml").unwrap();

    // if disk_meta.modified().unwrap() < config_meta.modified().unwrap() {
    //     // TODO: work on adding in system.toml support
    //     // TODO: rebuild the disk
    // }

    match args.next().as_deref() {
        Some("build" | "b") => {
            let mut release = false;
            let mut target = Target::X86_64;
            for arg in args {
                if arg == "-r" || arg == "--release" {
                    release = true;
                } else if arg == "rv64" || arg == "riscv64" || arg == "riscv64-virt" {
                    target = Target::Riscv64Virt;
                } else if arg == "arm64" || arg == "aarch64" || arg == "aarch64-virt" {
                    target = Target::Aarch64;
                } else {
                    return Err(report!(Error::InvalidSubCom));
                }
            }

            assemble()?;
            build(release, target).change_context(Error::Build)
        }
        Some("run" | "r") => {
            let mut release = false;
            let mut target = Target::X86_64;
            for arg in args {
                if arg == "-r" || arg == "--release" {
                    release = true;
                } else if arg == "rv64" || arg == "riscv64" || arg == "riscv64-virt" {
                    target = Target::Riscv64Virt;
                } else if arg == "arm64" || arg == "aarch64" || arg == "aarch64-virt" {
                    target = Target::Aarch64;
                } else {
                    return Err(report!(Error::InvalidSubCom));
                }
            }

            assemble()?;
            build(release, target)?;
            run(release, target)
        }
        Some("help" | "h") => {
            println!(concat!(
                "AbleOS RepBuild\n",
                "Subcommands:\n",
                "  build (b): Build a bootable disk image\n",
                "   help (h): Print this message\n",
                "    run (r): Build and run AbleOS in QEMU\n\n",
                "Options for build and run:\n",
                "        -r: build in release mode",
                "  [target]: sets target"
            ),);
            Ok(())
        }
        _ => Err(report!(Error::InvalidSubCom)),
    }
}

fn assemble() -> Result<(), Error> {
    match std::fs::create_dir("target/test-programs") {
        Ok(_) => (),
        Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => (),
        Err(e) => return Err(Report::new(e).change_context(Error::Io)),
    }

    for entry in std::fs::read_dir("sysdata/test-programs")
        .map_err(Report::from)
        .change_context(Error::Io)?
    {
        let entry = entry.map_err(Report::from).change_context(Error::Io)?;
        if !entry
            .file_type()
            .map_err(Report::from)
            .change_context(Error::Io)?
            .is_file()
        {
            continue;
        }

        let name = entry.file_name();
        let name = name.to_string_lossy();
        let name = name.trim_end_matches(".rhai");

        let mut out = File::options()
            .write(true)
            .create(true)
            .open(Path::new("target/test-programs").join(format!("{name}.hbf")))
            .map_err(Report::from)
            .change_context(Error::Io)?;

        out.set_len(0)
            .map_err(Report::from)
            .change_context(Error::Io)?;

        hbasm::assembler(&mut out, |engine| engine.run_file(entry.path()))
            .map_err(|e| report!(Error::Assembler).attach_printable(e.to_string()))?;
    }

    Ok(())
}

fn get_fs() -> Result<FileSystem<impl ReadWriteSeek>, io::Error> {
    let mut img = File::options()
        .read(true)
        .write(true)
        .create(true)
        .open(Path::new("target/disk.img"))?;

    img.set_len(1024 * 1024 * 64)?;

    fatfs::format_volume(&mut img, FormatVolumeOptions::new())?;

    let fs = FileSystem::new(img, FsOptions::new())?;
    let bootdir = fs.root_dir().create_dir("efi")?.create_dir("boot")?;

    io::copy(
        &mut File::open("limine/BOOTX64.EFI")
            .map_err(Report::from)
            .attach_printable("Copying Limine (x86_64): have you pulled the submodule?")?,
        &mut bootdir.create_file("bootx64.efi")?,
    )?;

    io::copy(
        &mut File::open("limine/BOOTAA64.EFI")
            .map_err(Report::from)
            .attach_printable("Copying Limine (ARM): have you pulled the submodule?")?,
        &mut bootdir.create_file("bootaa64.efi")?,
    )?;

    for fpath in [
        "sysdata/limine.cfg",
        "sysdata/background.bmp",
        "target/test-programs/failure.hbf",
        "target/test-programs/ecall.hbf",
        "target/test-programs/main.hbf",
        "target/test-programs/serial_driver.hbf",
        "target/test-programs/vfs_test.hbf",
        "target/test-programs/limine_framebuffer_driver.hbf",
        "target/test-programs/keyboard_driver.hbf",
    ] {
        let path = Path::new(fpath);
        io::copy(
            &mut File::open(path)?,
            &mut fs
                .root_dir()
                .create_file(&path.file_name().unwrap().to_string_lossy())?,
        )?;
    }

    drop(bootdir);
    Ok(fs)
}

fn build(release: bool, target: Target) -> Result<(), Error> {
    let fs = get_fs().change_context(Error::Io)?;
    let mut com = Command::new("cargo");
    com.current_dir("kernel");
    com.args(["b"]);
    if release {
        com.arg("-r");
    }

    if target == Target::Riscv64Virt {
        com.args(["--target", "targets/riscv64-virt-ableos.json"]);
    }
    if target == Target::Aarch64 {
        com.args(["--target", "targets/aarch64-virt-ableos.json"]);
    }

    match com.status() {
        Ok(s) if s.code() != Some(0) => bail!(Error::Build),
        Err(e) => bail!(report!(e).change_context(Error::Build)),
        _ => (),
    }

    let mut path: String = "kernel".to_string();
    let kernel_dir = match target {
        Target::X86_64 => {
            path.push_str("_x86-64");
            "target/x86_64-ableos"
        }
        Target::Riscv64Virt => "target/riscv64-virt-ableos",
        Target::Aarch64 => {
            path.push_str("_aarch64");
            "target/aarch64-virt-ableos"
        }
    };

    (|| -> std::io::Result<_> {
        io::copy(
            &mut File::open(
                Path::new(kernel_dir)
                    .join(if release { "release" } else { "debug" })
                    .join("kernel"),
            )?,
            &mut fs.root_dir().create_file(&path)?,
        )
        .map(|_| ())
    })()
    .map_err(Report::from)
    .change_context(Error::Io)
}

fn run(release: bool, target: Target) -> Result<(), Error> {
    let mut com = match target {
        Target::X86_64 => Command::new("qemu-system-x86_64"),
        Target::Riscv64Virt => Command::new("qemu-system-riscv64"),
        Target::Aarch64 => Command::new("qemu-system-aarch64"),
    };
    let ovmf_path = fetch_ovmf(target);

    match target {
        Target::X86_64 => {
            #[rustfmt::skip]
            com.args([
                "-bios", &ovmf_path.change_context(Error::OvmfFetch)?,
                "-drive", "file=target/disk.img,format=raw",
                "-m", "4G",
                "-smp", "cores=4",
                // "-enable-kvm",
                "-cpu", "Broadwell-v4"
            ]);
        }
        Target::Riscv64Virt => {
            #[rustfmt::skip]
            com.args([
                "-M", "virt",
                "-m", "128M",
                "-serial", "stdio",
                "-kernel",
                if release {
                    "target/riscv64-virt-ableos/release/kernel"
                } else {
                    "target/riscv64-virt-ableos/debug/kernel"
                }
            ]);
        }
        Target::Aarch64 => {
            #[rustfmt::skip]
            com.args([
                "-M", "virt",
                "-cpu", "cortex-a72",
                "-device", "ramfb",
                "-device", "qemu-xhci",
                "-device", "usb-kbd",
                "-m", "2G",
                "-bios", &ovmf_path.change_context(Error::OvmfFetch)?,
                "-drive", "file=target/disk.img,format=raw",
            ]);
        }
    }

    match com
        .status()
        .map_err(Report::from)
        .change_context(Error::ProcessSpawn)?
    {
        s if s.success() => Ok(()),
        s => Err(report!(Error::Qemu(s.code()))),
    }
}

fn fetch_ovmf(target: Target) -> Result<String, OvmfFetchError> {
    let (ovmf_url, ovmf_path) = match target {
        Target::X86_64 => (
            "https://retrage.github.io/edk2-nightly/bin/RELEASEX64_OVMF.fd",
            "target/RELEASEX64_OVMF.fd",
        ),
        Target::Riscv64Virt => return Err(OvmfFetchError::Empty.into()),
        Target::Aarch64 => (
            "https://retrage.github.io/edk2-nightly/bin/RELEASEAARCH64_QEMU_EFI.fd",
            "target/RELEASEAARCH64_QEMU_EFI.fd",
        ),
    };

    let mut file = match std::fs::metadata(ovmf_path) {
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => std::fs::OpenOptions::new()
            .create(true)
            .write(true)
            .read(true)
            .open(ovmf_path)
            .map_err(Report::from)
            .change_context(OvmfFetchError::Io)?,
        Ok(_) => return Ok(ovmf_path.to_owned()),
        Err(e) => return Err(report!(e).change_context(OvmfFetchError::Io)),
    };
    let mut bytes = reqwest::blocking::get(ovmf_url)
        .map_err(Report::from)
        .change_context(OvmfFetchError::Fetch)?;

    bytes
        .copy_to(&mut file)
        .map_err(Report::from)
        .change_context(OvmfFetchError::Io)?;

    Ok(ovmf_path.to_owned())
}

#[derive(Debug, Display)]
enum OvmfFetchError {
    #[display(fmt = "Failed to fetch OVMF package")]
    Fetch,
    #[display(fmt = "No OVMF package available")]
    Empty,
    #[display(fmt = "IO Error")]
    Io,
}

impl Context for OvmfFetchError {}

#[derive(Clone, Copy, PartialEq, Eq)]
enum Target {
    X86_64,
    Riscv64Virt,
    Aarch64,
}

#[derive(Debug, Display)]
enum Error {
    #[display(fmt = "Failed to build the kernel")]
    Build,
    #[display(fmt = "Missing or invalid subcommand (available: build, run)")]
    InvalidSubCom,
    #[display(fmt = "IO Error")]
    Io,
    #[display(fmt = "Failed to spawn a process")]
    ProcessSpawn,
    #[display(fmt = "Failed to fetch UEFI firmware")]
    OvmfFetch,
    #[display(fmt = "Failed to assemble Holey Bytes code")]
    Assembler,
    #[display(fmt = "QEMU Error: {}", "fmt_qemu_err(*_0)")]
    Qemu(Option<i32>),
}

impl Context for Error {}

fn fmt_qemu_err(e: Option<i32>) -> impl Display {
    struct W(Option<i32>);
    impl Display for W {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            if let Some(c) = self.0 {
                c.fmt(f)
            } else {
                f.write_str("Interrupted by signal")
            }
        }
    }

    W(e)
}