Add support for adding a ramdisk to the disk image

src/lib.rs

@@ -17,22 +17,31 @@ mod mbr;
 mod pxe;
 
 const KERNEL_FILE_NAME: &str = "kernel-x86_64";
+const RAMDISK_FILE_NAME: &str = "ramdisk-x86_64";
 const BIOS_STAGE_3: &str = "boot-stage-3";
 const BIOS_STAGE_4: &str = "boot-stage-4";
 
 /// Create disk images for booting on legacy BIOS systems.
 pub struct BiosBoot {
     kernel: PathBuf,
+    ramdisk: Option<PathBuf>,
 }
 
 impl BiosBoot {
     /// Start creating a disk image for the given bootloader ELF executable.
     pub fn new(kernel_path: &Path) -> Self {
         Self {
             kernel: kernel_path.to_owned(),
+            ramdisk: None,
         }
     }
 
+    /// Add a ramdisk file to the image
+    pub fn set_ramdisk(&mut self, ramdisk_path: &Path) -> &mut Self {
+        self.ramdisk = Some(ramdisk_path.to_owned());
+        self
+    }
+
     /// Create a bootable UEFI disk image at the given path.
     pub fn create_disk_image(&self, out_path: &Path) -> anyhow::Result<()> {
         let bootsector_path = Path::new(env!("BIOS_BOOT_SECTOR_PATH"));
@@ -61,12 +70,18 @@ impl BiosBoot {
     fn create_fat_partition(&self) -> anyhow::Result<NamedTempFile> {
         let stage_3_path = Path::new(env!("BIOS_STAGE_3_PATH"));
         let stage_4_path = Path::new(env!("BIOS_STAGE_4_PATH"));
+        let has_rd_path = self.ramdisk.is_some();
+        let binding = self.ramdisk.as_deref();
+        let ramdisk_path = binding.unwrap_or(Path::new("no-such-file"));
+        let kernel_path = self.kernel.as_path();
 
         let mut files = BTreeMap::new();
-        files.insert(KERNEL_FILE_NAME, self.kernel.as_path());
+        files.insert(KERNEL_FILE_NAME, kernel_path);
         files.insert(BIOS_STAGE_3, stage_3_path);
         files.insert(BIOS_STAGE_4, stage_4_path);
-
+        if let Some(ramdisk_path) = &self.ramdisk {
+            files.insert(RAMDISK_FILE_NAME, ramdisk_path);
+        }
         let out_file = NamedTempFile::new().context("failed to create temp file")?;
         fat::create_fat_filesystem(files, out_file.path())
             .context("failed to create BIOS FAT filesystem")?;
@@ -78,13 +93,23 @@ impl BiosBoot {
 /// Create disk images for booting on UEFI systems.
 pub struct UefiBoot {
     kernel: PathBuf,
+    ramdisk: Option<PathBuf>,
 }
 
 impl UefiBoot {
     /// Start creating a disk image for the given bootloader ELF executable.
     pub fn new(kernel_path: &Path) -> Self {
         Self {
             kernel: kernel_path.to_owned(),
+            ramdisk: None,
+        }
+    }
+
+    /// Add a ramdisk file to the disk image
+    pub fn with_ramdisk(&self, ramdisk_path: &Path) -> Self {
+        Self {
+            kernel: self.kernel.clone(),
+            ramdisk: Some(ramdisk_path.to_owned()),
         }
     }
 
@@ -121,10 +146,17 @@ impl UefiBoot {
     /// Creates an UEFI-bootable FAT partition with the kernel.
     fn create_fat_partition(&self) -> anyhow::Result<NamedTempFile> {
         let bootloader_path = Path::new(env!("UEFI_BOOTLOADER_PATH"));
-
+        let has_rd_path = self.ramdisk.is_some();
+        let binding = self.ramdisk.as_deref();
+        let ramdisk_path = binding.unwrap_or(Path::new("no-such-file"));
+        let kernel_path = self.kernel.as_path();
         let mut files = BTreeMap::new();
         files.insert("efi/boot/bootx64.efi", bootloader_path);
-        files.insert(KERNEL_FILE_NAME, self.kernel.as_path());
+        files.insert(KERNEL_FILE_NAME, kernel_path);
+
+        if has_rd_path {
+            files.insert(RAMDISK_FILE_NAME, ramdisk_path);
+        }
 
         let out_file = NamedTempFile::new().context("failed to create temp file")?;
         fat::create_fat_filesystem(files, out_file.path())