Merge pull request #535 from Wasabi375/disk_image

Expose data for custom boot image creation

Author: Freax13

src/lib.rs

@@ -111,13 +111,15 @@ impl DiskImageBuilder {
     pub fn create_bios_image(&self, image_path: &Path) -> anyhow::Result<()> {
         const BIOS_STAGE_3_NAME: &str = "boot-stage-3";
         const BIOS_STAGE_4_NAME: &str = "boot-stage-4";
+
+        let fat_partition = NamedTempFile::new().context("failed to create temp file")?;
+
         let stage_3 = FileDataSource::Bytes(BIOS_STAGE_3);
         let stage_4 = FileDataSource::Bytes(BIOS_STAGE_4);
         let mut internal_files = BTreeMap::new();
         internal_files.insert(BIOS_STAGE_3_NAME, stage_3);
         internal_files.insert(BIOS_STAGE_4_NAME, stage_4);
-        let fat_partition = self
-            .create_fat_filesystem_image(internal_files)
+        self.create_fat_filesystem_image(internal_files, fat_partition.path())
             .context("failed to create FAT partition")?;
         mbr::create_mbr_disk(
             BIOS_BOOT_SECTOR,
@@ -136,13 +138,10 @@ impl DiskImageBuilder {
     #[cfg(feature = "uefi")]
     /// Create a GPT disk image for booting on UEFI systems.
     pub fn create_uefi_image(&self, image_path: &Path) -> anyhow::Result<()> {
-        const UEFI_BOOT_FILENAME: &str = "efi/boot/bootx64.efi";
-
-        let mut internal_files = BTreeMap::new();
-        internal_files.insert(UEFI_BOOT_FILENAME, FileDataSource::Bytes(UEFI_BOOTLOADER));
-        let fat_partition = self
-            .create_fat_filesystem_image(internal_files)
+        let fat_partition = NamedTempFile::new().context("failed to create temp file")?;
+        self.create_uefi_fat_partition(fat_partition.path())
             .context("failed to create FAT partition")?;
+
         gpt::create_gpt_disk(fat_partition.path(), image_path)
             .context("failed to create UEFI GPT disk image")?;
         fat_partition
@@ -152,6 +151,19 @@ impl DiskImageBuilder {
         Ok(())
     }
 
+    #[cfg(feature = "uefi")]
+    /// Create a bootable FAT partition for a UEFI system.
+    pub fn create_uefi_fat_partition(&self, partition_path: &Path) -> anyhow::Result<()> {
+        const UEFI_BOOT_FILENAME: &str = "efi/boot/bootx64.efi";
+
+        let mut internal_files = BTreeMap::new();
+        internal_files.insert(UEFI_BOOT_FILENAME, FileDataSource::Bytes(UEFI_BOOTLOADER));
+        self.create_fat_filesystem_image(internal_files, partition_path)
+            .context("failed to create FAT partition")?;
+
+        Ok(())
+    }
+
     #[cfg(feature = "uefi")]
     /// Create a folder containing the needed files for UEFI TFTP/PXE booting.
     pub fn create_uefi_tftp_folder(&self, tftp_path: &Path) -> anyhow::Result<()> {
@@ -198,7 +210,8 @@ impl DiskImageBuilder {
     fn create_fat_filesystem_image(
         &self,
         internal_files: BTreeMap<&str, FileDataSource>,
-    ) -> anyhow::Result<NamedTempFile> {
+        partition_path: &Path,
+    ) -> anyhow::Result<()> {
         let mut local_map: BTreeMap<&str, _> = BTreeMap::new();
 
         for (name, source) in &self.files {
@@ -214,10 +227,9 @@ impl DiskImageBuilder {
             }
         }
 
-        let out_file = NamedTempFile::new().context("failed to create temp file")?;
-        fat::create_fat_filesystem(local_map, out_file.path())
+        fat::create_fat_filesystem(local_map, partition_path)
             .context("failed to create FAT filesystem")?;
 
-        Ok(out_file)
+        Ok(())
     }
 }