build single blob for stages 2 and 3

xtask/src/build.rs

@@ -1,6 +1,7 @@
 use std::{
     ffi::OsStr,
-    fs,
+    fs::{self, File, OpenOptions},
+    io::{self, Seek, Write},
     path::{Path, PathBuf},
     process::{Command, Output, Stdio},
 };
@@ -9,29 +10,39 @@ use eyre::{WrapErr, eyre};
 
 use crate::Context;
 
+const SECTOR_SIZE: usize = 512;
+
 const NASM_COMMON_FLAGS: &[&str] = &["-werror"];
 
 pub fn build(ctx: &Context) -> Result<(), eyre::Error> {
-    let stage_1_src_dir = ctx.workspace.join("stage_1");
-    let stage_2_src_dir = ctx.workspace.join("stage_2");
-    let stage_3_src_dir = ctx.workspace.join("stage_3");
+    let s1_src_dir = ctx.workspace.join("stage_1");
+    let s2_src_dir = ctx.workspace.join("stage_2");
+    let s3_src_dir = ctx.workspace.join("stage_3");
 
     let build_dir = ctx.workspace.join("build");
-    let stage_1_build_dir = build_dir.join("stage_1");
-    let stage_2_build_dir = build_dir.join("stage_2");
+    let s1_build_dir = build_dir.join("stage_1");
+    let s2_build_dir = build_dir.join("stage_2");
 
     mkdir_if_missing(&build_dir)?;
-    mkdir_if_missing(&stage_1_build_dir)?;
-    mkdir_if_missing(&stage_2_build_dir)?;
+    mkdir_if_missing(&s1_build_dir)?;
+    mkdir_if_missing(&s2_build_dir)?;
 
-    build_stage_1(ctx, &stage_1_build_dir, &stage_1_src_dir)?;
-    build_stage_2(ctx, &stage_2_build_dir, &stage_2_src_dir)?;
-    build_stage_3(ctx, &stage_3_src_dir)?;
+    println!("building stage 1");
+    let s1_bin_path = build_stage_1(ctx, &s1_build_dir, &s1_src_dir)?;
+
+    println!("building stage 2");
+    let s2_bin_path = build_stage_2(ctx, &s2_build_dir, &s2_src_dir)?;
+
+    println!("building stage 3");
+    let s3_bin_path = build_stage_3(ctx, &s3_src_dir)?;
+
+    println!("creating stage 2/3 blob");
+    make_s23_blob(&build_dir, &s2_bin_path, &s3_bin_path)?;
 
     Ok(())
 }
 
-fn build_stage_1(ctx: &Context, build_dir: &Path, src_dir: &Path) -> Result<(), eyre::Error> {
+fn build_stage_1(ctx: &Context, build_dir: &Path, src_dir: &Path) -> Result<PathBuf, eyre::Error> {
     let src_paths = ls_with_extension(src_dir, "s")?;
 
     let [src_path] = &*src_paths else {
@@ -41,7 +52,7 @@ fn build_stage_1(ctx: &Context, build_dir: &Path, src_dir: &Path) -> Result<(),
         ));
     };
 
-    let out_path = build_dir.join("s1.bin");
+    let bin_path = build_dir.join("s1.bin");
 
     let include_dir = ctx.workspace.join("include");
 
@@ -51,16 +62,16 @@ fn build_stage_1(ctx: &Context, build_dir: &Path, src_dir: &Path) -> Result<(),
         .arg("-I")
         .arg(&include_dir)
         .arg("-o")
-        .arg(&out_path)
+        .arg(&bin_path)
         .arg(src_path)
         .stdout(Stdio::inherit())
         .stderr(Stdio::inherit())
         .run_ok()?;
 
-    Ok(())
+    Ok(bin_path)
 }
 
-fn build_stage_2(ctx: &Context, build_dir: &Path, src_dir: &Path) -> Result<(), eyre::Error> {
+fn build_stage_2(ctx: &Context, build_dir: &Path, src_dir: &Path) -> Result<PathBuf, eyre::Error> {
     let include_dir = ctx.workspace.join("include");
     let src_paths = ls_with_extension(src_dir, "s")?;
     let mut obj_paths = Vec::new();
@@ -95,16 +106,16 @@ fn build_stage_2(ctx: &Context, build_dir: &Path, src_dir: &Path) -> Result<(),
         .arg("-T")
         .arg(&linker_file_path)
         .arg("-o")
-        .arg(bin_path)
+        .arg(&bin_path)
         .args(&obj_paths)
         .stdout(Stdio::inherit())
         .stderr(Stdio::inherit())
         .run_ok()?;
 
-    Ok(())
+    Ok(bin_path)
 }
 
-fn build_stage_3(ctx: &Context, stage_3_src_dir: &Path) -> Result<(), eyre::Error> {
+fn build_stage_3(ctx: &Context, stage_3_src_dir: &Path) -> Result<PathBuf, eyre::Error> {
     Command::new(ctx.cargo)
         .arg("build")
         .arg("--release")
@@ -113,29 +124,72 @@ fn build_stage_3(ctx: &Context, stage_3_src_dir: &Path) -> Result<(), eyre::Erro
         .stderr(Stdio::inherit())
         .run_ok()?;
 
+    let bin_path = {
+        let mut bin_path = stage_3_src_dir.to_owned();
+        bin_path.extend(["target", "protected32", "release", "stage_3"]);
+        bin_path
+    };
+
+    Ok(bin_path)
+}
+
+fn make_s23_blob(
+    build_dir: &Path, s2_bin_path: &Path, s3_bin_path: &Path,
+) -> Result<(), eyre::Error> {
+    let out_path = build_dir.join("s23.bin");
+
+    let mut out_file = OpenOptions::new()
+        .write(true)
+        .truncate(true)
+        .create(true)
+        .open(&out_path)
+        .wrap_io_err(IoOp::Open, &out_path)?;
+
+    let mut buf = [0u8; SECTOR_SIZE];
+
+    let s2_len_sectors = read_sectors_from_to(s2_bin_path, &mut out_file, &out_path, &mut buf)?;
+    if s2_len_sectors == 0 {
+        return Err(eyre!("empty stage 2"));
+    }
+    let s3_offset_sectors_16 = u16::try_from(s2_len_sectors).wrap_err("stage 3 offset overflow")?;
+
+    let s3_len_sectors = read_sectors_from_to(s3_bin_path, &mut out_file, &out_path, &mut buf)?;
+    if s3_len_sectors == 0 {
+        return Err(eyre!("empty stage 3"));
+    }
+    let s3_len_sectors_16 = u16::try_from(s3_len_sectors).wrap_err("stage 3 length overflow")?;
+
+    out_file
+        .seek(io::SeekFrom::Start(0))
+        .wrap_io_err(IoOp::Seek, &out_path)?;
+
+    // Write the stage 2 data section fields
+    out_file
+        .write_all(&s3_offset_sectors_16.to_le_bytes())
+        .wrap_io_err(IoOp::Write, &out_path)?;
+    out_file
+        .write_all(&s3_len_sectors_16.to_le_bytes())
+        .wrap_io_err(IoOp::Write, &out_path)?;
+
+    out_file.flush().wrap_io_err(IoOp::Flush, &out_path)?;
+
     Ok(())
 }
 
 fn mkdir_if_missing(path: &Path) -> Result<(), eyre::Error> {
-    let exists = fs::exists(path)
-        .wrap_err_with(|| format!("failed to check if {} exists", path.to_string_lossy()))?;
-
-    if !exists {
-        fs::create_dir(path)
-            .wrap_err_with(|| format!("failed to create {}", path.to_string_lossy()))?;
+    if !fs::exists(path).wrap_io_err(IoOp::Stat, path)? {
+        fs::create_dir(path).wrap_io_err(IoOp::Open, path)?;
     }
-
     Ok(())
 }
 
 fn ls_with_extension(path: &Path, ext: &str) -> Result<Vec<PathBuf>, eyre::Error> {
-    let read_dir = fs::read_dir(path)
-        .wrap_err_with(|| format!("failed to read {}", path.to_string_lossy()))?;
+    let read_dir = fs::read_dir(path).wrap_io_err(IoOp::Read, path)?;
 
     let mut paths = Vec::new();
 
     for dir in read_dir {
-        let dir = dir.wrap_err_with(|| format!("failed to read {}", path.to_string_lossy()))?;
+        let dir = dir.wrap_io_err(IoOp::Read, path)?;
 
         if dir.path().extension() == Some(OsStr::new(ext)) {
             paths.push(dir.path());
@@ -145,6 +199,57 @@ fn ls_with_extension(path: &Path, ext: &str) -> Result<Vec<PathBuf>, eyre::Error
     Ok(paths)
 }
 
+fn read_retry_eintr<T>(r: &mut T, buf: &mut [u8]) -> Result<usize, io::Error>
+where
+    T: io::Read,
+{
+    loop {
+        match r.read(buf) {
+            Ok(n) => return Ok(n),
+            Err(err) => match err.kind() {
+                io::ErrorKind::Interrupted => (),
+                _ => return Err(err),
+            },
+        }
+    }
+}
+
+fn read_sectors_from_to(
+    src_path: &Path, dst_file: &mut File, dst_path: &Path, buf: &mut [u8],
+) -> Result<u64, eyre::Error> {
+    let mut total_read = 0u64;
+
+    let mut src_file = File::open(src_path).wrap_io_err(IoOp::Open, src_path)?;
+
+    loop {
+        match read_retry_eintr(&mut src_file, buf).wrap_io_err(IoOp::Read, src_path)? {
+            0 => break,
+            n => {
+                total_read += n as u64;
+                dst_file
+                    .write_all(&buf[..n])
+                    .wrap_io_err(IoOp::Write, dst_path)?;
+            }
+        }
+    }
+
+    let remainder = (total_read % (SECTOR_SIZE as u64)) as usize;
+    let padding = (SECTOR_SIZE - remainder) % SECTOR_SIZE;
+
+    if padding != 0 {
+        // Pad with 0xf4 since it encodes HLT, in case we accidentally jump to it.
+        buf[..padding].fill(0xf4);
+        dst_file
+            .write_all(&buf[..padding])
+            .wrap_io_err(IoOp::Write, dst_path)?;
+    }
+
+    let total_bytes = total_read + padding as u64;
+    assert!(total_bytes % SECTOR_SIZE as u64 == 0);
+
+    Ok(total_bytes / SECTOR_SIZE as u64)
+}
+
 trait CommandExt {
     fn run_ok(&mut self) -> Result<Output, eyre::Error>;
 }
@@ -162,3 +267,39 @@ impl CommandExt for Command {
         Ok(output)
     }
 }
+
+trait WrapIoError<T> {
+    fn wrap_io_err(self, op: IoOp, path: &Path) -> Result<T, eyre::Report>;
+}
+
+impl<T, E> WrapIoError<T> for Result<T, E>
+where
+    Self: WrapErr<T, E>,
+{
+    fn wrap_io_err(self, op: IoOp, path: &Path) -> Result<T, eyre::Report> {
+        self.wrap_err_with(|| format!("failed to {} {}", op.name(), path.to_string_lossy()))
+    }
+}
+
+#[derive(Clone, Copy)]
+enum IoOp {
+    Read,
+    Write,
+    Stat,
+    Open,
+    Seek,
+    Flush,
+}
+
+impl IoOp {
+    fn name(self) -> &'static str {
+        match self {
+            IoOp::Read => "read",
+            IoOp::Write => "write",
+            IoOp::Stat => "stat",
+            IoOp::Open => "open",
+            IoOp::Seek => "seek",
+            IoOp::Flush => "flush",
+        }
+    }
+}