Add function to create UEFI disk image to library

Author: phil-opp

src/lib.rs

@@ -67,7 +67,138 @@ for all possible configuration options.
 
 #![warn(missing_docs)]
 
-/// Provides a function to turn a bootloader executable into a disk image.
-///
-/// Used by the `builder` binary. Only available when the `builder` feature is enabled.
-pub mod disk_image;
+use std::{
+    fs::{self, File},
+    io::{self, Seek},
+    path::Path,
+};
+
+use anyhow::Context;
+
+pub fn create_uefi_disk_image(
+    kernel_binary: &Path,
+    out_fat_path: &Path,
+    out_gpt_path: &Path,
+) -> anyhow::Result<()> {
+    let bootloader_path = Path::new(env!("UEFI_BOOTLOADER_PATH"));
+
+    create_fat_filesystem(bootloader_path, kernel_binary, &out_fat_path)
+        .context("failed to create UEFI FAT filesystem")?;
+    create_gpt_disk(out_fat_path, out_gpt_path);
+
+    Ok(())
+}
+
+fn create_fat_filesystem(
+    bootloader_efi_file: &Path,
+    kernel_binary: &Path,
+    out_fat_path: &Path,
+) -> anyhow::Result<()> {
+    const MB: u64 = 1024 * 1024;
+
+    // retrieve size of `.efi` file and round it up
+    let efi_size = fs::metadata(&bootloader_efi_file).unwrap().len();
+    // size of a megabyte
+    // round it to next megabyte
+    let efi_size_rounded = ((efi_size - 1) / MB + 1) * MB;
+
+    // create new filesystem image file at the given path and set its length
+    let fat_file = fs::OpenOptions::new()
+        .read(true)
+        .write(true)
+        .create(true)
+        .truncate(true)
+        .open(&out_fat_path)
+        .unwrap();
+    fat_file.set_len(efi_size_rounded).unwrap();
+
+    // create new FAT file system and open it
+    let label = {
+        if let Some(name) = bootloader_efi_file.file_stem() {
+            let converted = name.to_string_lossy();
+            let name = converted.as_bytes();
+            let mut label = [0u8; 11];
+            let name = &name[..label.len()];
+            let slice = &mut label[..name.len()];
+            slice.copy_from_slice(name);
+            label
+        } else {
+            *b"MY_RUST_OS!"
+        }
+    };
+    let format_options = fatfs::FormatVolumeOptions::new().volume_label(label);
+    fatfs::format_volume(&fat_file, format_options).context("Failed to format UEFI FAT file")?;
+    let filesystem = fatfs::FileSystem::new(&fat_file, fatfs::FsOptions::new())
+        .context("Failed to open FAT file system of UEFI FAT file")?;
+
+    // copy EFI file to FAT filesystem
+    let root_dir = filesystem.root_dir();
+    root_dir.create_dir("efi").unwrap();
+    root_dir.create_dir("efi/boot").unwrap();
+    let mut bootx64 = root_dir.create_file("efi/boot/bootx64.efi").unwrap();
+    bootx64.truncate().unwrap();
+    io::copy(
+        &mut fs::File::open(&bootloader_efi_file).unwrap(),
+        &mut bootx64,
+    )
+    .unwrap();
+
+    // copy kernel to FAT filesystem
+    let mut kernel_file = root_dir.create_file("kernel-x86_64")?;
+    kernel_file.truncate()?;
+    io::copy(&mut fs::File::open(&kernel_binary)?, &mut kernel_file)?;
+
+    Ok(())
+}
+
+fn create_gpt_disk(fat_image: &Path, out_gpt_path: &Path) {
+    // create new file
+    let mut disk = fs::OpenOptions::new()
+        .create(true)
+        .truncate(true)
+        .read(true)
+        .write(true)
+        .open(&out_gpt_path)
+        .unwrap();
+
+    // set file size
+    let partition_size: u64 = fs::metadata(&fat_image).unwrap().len();
+    let disk_size = partition_size + 1024 * 64; // for GPT headers
+    disk.set_len(disk_size).unwrap();
+
+    // create a protective MBR at LBA0 so that disk is not considered
+    // unformatted on BIOS systems
+    let mbr = gpt::mbr::ProtectiveMBR::with_lb_size(
+        u32::try_from((disk_size / 512) - 1).unwrap_or(0xFF_FF_FF_FF),
+    );
+    mbr.overwrite_lba0(&mut disk).unwrap();
+
+    // create new GPT structure
+    let block_size = gpt::disk::LogicalBlockSize::Lb512;
+    let mut gpt = gpt::GptConfig::new()
+        .writable(true)
+        .initialized(false)
+        .logical_block_size(block_size)
+        .create_from_device(Box::new(&mut disk), None)
+        .unwrap();
+    gpt.update_partitions(Default::default()).unwrap();
+
+    // add new EFI system partition and get its byte offset in the file
+    let partition_id = gpt
+        .add_partition("boot", partition_size, gpt::partition_types::EFI, 0, None)
+        .unwrap();
+    let partition = gpt.partitions().get(&partition_id).unwrap();
+    let start_offset = partition.bytes_start(block_size).unwrap();
+
+    // close the GPT structure and write out changes
+    gpt.write().unwrap();
+
+    // place the FAT filesystem in the newly created partition
+    disk.seek(io::SeekFrom::Start(start_offset)).unwrap();
+    io::copy(&mut File::open(&fat_image).unwrap(), &mut disk).unwrap();
+}
+
+// Provides a function to turn a bootloader executable into a disk image.
+//
+// Used by the `builder` binary. Only available when the `builder` feature is enabled.
+// pub mod disk_image;