Create a GPT disk image with an UEFI boot partition

Author: phil-opp

Cargo.lock

@@ -47,6 +47,7 @@ thiserror = { version = "1.0.20", optional = true }
 json = { version = "0.12.4", optional = true }
 rsdp = { version = "1.0.0", optional = true }
 fatfs = { version = "0.3.4", optional = true }
+gpt = { version = "2.0.0", optional = true }
 
 [dependencies.font8x8]
 version = "0.2.5"
@@ -61,7 +62,7 @@ serde = { version = "1.0", features = ["derive"], optional = true}
 
 [features]
 default = []
-builder = ["argh", "thiserror", "displaydoc", "anyhow", "llvm-tools", "json", "fatfs"]
+builder = ["argh", "thiserror", "displaydoc", "anyhow", "llvm-tools", "json", "fatfs", "gpt"]
 runner = ["anyhow"]
 bios_bin = ["binary", "vga_320x200", "rsdp"]
 uefi_bin = ["binary", "uefi", "font8x8"]

Cargo.toml

@@ -1,8 +1,10 @@
-use anyhow::{anyhow, Context};
+use anyhow::{anyhow, bail, Context};
 use argh::FromArgs;
 use bootloader::disk_image::create_disk_image;
 use std::{
-    fs, io,
+    convert::TryFrom,
+    fs::{self, File},
+    io::{self, Seek},
     path::{Path, PathBuf},
     process::Command,
     str::FromStr,
@@ -129,55 +131,32 @@ fn main() -> anyhow::Result<()> {
 
         assert_eq!(executables.len(), 1);
         let executable_path = executables.pop().unwrap();
-        let executable_name = executable_path.file_stem().unwrap().to_str().unwrap();
-        let kernel_name = args.kernel_binary.file_name().unwrap().to_str().unwrap();
+
+        let executable_name = executable_path
+            .file_stem()
+            .and_then(|stem| stem.to_str())
+            .ok_or_else(|| {
+                anyhow!(
+                    "executable path `{}` has invalid file stem",
+                    executable_path.display()
+                )
+            })?;
+        let kernel_name = args
+            .kernel_binary
+            .file_name()
+            .and_then(|name| name.to_str())
+            .ok_or_else(|| {
+                anyhow!(
+                    "kernel binary path `{}` has invalid file name",
+                    args.kernel_binary.display()
+                )
+            })?;
 
         if let Some(out_dir) = &args.out_dir {
             let efi_file =
                 out_dir.join(format!("bootimage-{}-{}.efi", executable_name, kernel_name));
-            fs::copy(&executable_path, &efi_file).context("failed to copy efi file to out dir")?;
-
-            let efi_size = fs::metadata(&efi_file)
-                .context("failed to read metadata of efi file")?
-                .len();
-
-            // create fat partition
-            {
-                const MB: u64 = 1024 * 1024;
-
-                let fat_path = efi_file.with_extension("fat");
-                dbg!(&fat_path);
-                let fat_file = fs::OpenOptions::new()
-                    .read(true)
-                    .write(true)
-                    .create(true)
-                    .truncate(true)
-                    .open(&fat_path)
-                    .context("Failed to create UEFI FAT file")?;
-                let efi_size_rounded = ((efi_size - 1) / MB + 1) * MB;
-                fat_file
-                    .set_len(dbg!(efi_size_rounded))
-                    .context("failed to set UEFI FAT file length")?;
-
-                // create new FAT partition
-                fatfs::format_volume(&fat_file, fatfs::FormatVolumeOptions::new())
-                    .context("Failed to format UEFI FAT file")?;
-
-                // copy EFI file to FAT filesystem
-                let partition = fatfs::FileSystem::new(&fat_file, fatfs::FsOptions::new())
-                    .context("Failed to open FAT file system of UEFI FAT file")?;
-                let root_dir = partition.root_dir();
-                root_dir.create_dir("efi")?;
-                root_dir.create_dir("efi/boot")?;
-                let mut bootx64 = root_dir.create_file("efi/boot/bootx64.efi")?;
-                bootx64.truncate()?;
-                io::copy(&mut fs::File::open(&executable_path)?, &mut bootx64)?;
-            }
-
-            // create gpt disk
-            {
-                //todo!()
-            }
+            create_uefi_disk_image(&executable_path, &efi_file)
+                .context("failed to create UEFI disk image")?;
         }
     }
 
@@ -228,7 +207,7 @@ fn main() -> anyhow::Result<()> {
         let mut output_bin_path = executable_path
             .parent()
             .unwrap()
-            .join(format!("bootimage-{}-{}.bin", executable_name, kernel_name));
+            .join(format!("bootimage-{}-{}.img", executable_name, kernel_name));
 
         create_disk_image(&executable_path, &output_bin_path)
             .context("Failed to create bootable disk image")?;
@@ -254,6 +233,132 @@ fn main() -> anyhow::Result<()> {
     Ok(())
 }
 
+fn create_uefi_disk_image(executable_path: &Path, efi_file: &Path) -> anyhow::Result<()> {
+    fs::copy(&executable_path, &efi_file).context("failed to copy efi file to out dir")?;
+
+    let efi_size = fs::metadata(&efi_file)
+        .context("failed to read metadata of efi file")?
+        .len();
+
+    // create fat partition
+    let fat_file_path = {
+        const MB: u64 = 1024 * 1024;
+
+        let fat_path = efi_file.with_extension("fat");
+        let fat_file = fs::OpenOptions::new()
+            .read(true)
+            .write(true)
+            .create(true)
+            .truncate(true)
+            .open(&fat_path)
+            .context("Failed to create UEFI FAT file")?;
+        let efi_size_rounded = ((efi_size - 1) / MB + 1) * MB;
+        fat_file
+            .set_len(efi_size_rounded)
+            .context("failed to set UEFI FAT file length")?;
+
+        // create new FAT partition
+        let format_options = fatfs::FormatVolumeOptions::new().volume_label(*b"FOOO       ");
+        fatfs::format_volume(&fat_file, format_options)
+            .context("Failed to format UEFI FAT file")?;
+
+        // copy EFI file to FAT filesystem
+        let partition = fatfs::FileSystem::new(&fat_file, fatfs::FsOptions::new())
+            .context("Failed to open FAT file system of UEFI FAT file")?;
+        let root_dir = partition.root_dir();
+        root_dir.create_dir("efi")?;
+        root_dir.create_dir("efi/boot")?;
+        let mut bootx64 = root_dir.create_file("efi/boot/bootx64.efi")?;
+        bootx64.truncate()?;
+        io::copy(&mut fs::File::open(&executable_path)?, &mut bootx64)?;
+
+        fat_path
+    };
+
+    // create gpt disk
+    {
+        let image_path = efi_file.with_extension("img");
+        let mut image = fs::OpenOptions::new()
+            .create(true)
+            .truncate(true)
+            .read(true)
+            .write(true)
+            .open(&image_path)
+            .context("failed to create UEFI disk image")?;
+
+        let partition_size: u64 = fs::metadata(&fat_file_path)
+            .context("failed to read metadata of UEFI FAT partition")?
+            .len();
+        let image_size = partition_size + 1024 * 64;
+        image
+            .set_len(image_size)
+            .context("failed to set length of UEFI disk image")?;
+
+        // Create a protective MBR at LBA0
+        let mbr = gpt::mbr::ProtectiveMBR::with_lb_size(
+            u32::try_from((image_size / 512) - 1).unwrap_or(0xFF_FF_FF_FF),
+        );
+        mbr.overwrite_lba0(&mut image)
+            .context("failed to write protective MBR")?;
+
+        // create new GPT in image file
+        let block_size = gpt::disk::LogicalBlockSize::Lb512;
+        let block_size_bytes: u64 = block_size.into();
+        let mut disk = gpt::GptConfig::new()
+            .writable(true)
+            .initialized(false)
+            .logical_block_size(block_size)
+            .create_from_device(Box::new(&mut image), None)
+            .context("failed to open UEFI disk image")?;
+        disk.update_partitions(Default::default())
+            .context("failed to initialize GPT partition table")?;
+
+        // add add EFI system partition
+        let partition_id = disk
+            .add_partition("boot", partition_size, gpt::partition_types::EFI, 0)
+            .context("failed to add boot partition")?;
+
+        let partition = disk
+            .partitions()
+            .get(&partition_id)
+            .ok_or_else(|| anyhow!("Partition doesn't exist after adding it"))?;
+        let created_partition_size: u64 =
+            (partition.last_lba - partition.first_lba + 1u64) * block_size_bytes;
+        if created_partition_size != partition_size {
+            bail!(
+                "Created partition has invalid size (size is {:?}, expected {})",
+                created_partition_size,
+                partition_size
+            );
+        }
+        let start_offset = partition
+            .bytes_start(block_size)
+            .context("failed to retrieve partition start offset")?;
+
+        // Write the partition table
+        disk.write()
+            .context("failed to write GPT partition table to UEFI image file")?;
+
+        image
+            .seek(io::SeekFrom::Start(start_offset))
+            .context("failed to seek to boot partiiton start")?;
+        let bytes_written = io::copy(
+            &mut File::open(&fat_file_path).context("failed to open fat image")?,
+            &mut image,
+        )
+        .context("failed to write boot partition content")?;
+        if bytes_written != partition_size {
+            bail!(
+                "Invalid number of partition bytes written (expected {}, got {})",
+                partition_size,
+                bytes_written
+            );
+        }
+    }
+
+    Ok(())
+}
+
 fn bios_run(bin_path: &Path) -> anyhow::Result<Option<ExitCode>> {
     let mut qemu = Command::new("qemu-system-x86_64");
     qemu.arg("-drive")