Skip to content

Commit 154dbab

Browse files
jasoncouturejasoncouture
committed
Replace UefiBoot and BiosBoot with DiskImageBuilder
1 parent b86fb0c commit 154dbab

File tree

2 files changed

+116
-120
lines changed

2 files changed

+116
-120
lines changed

‎src/lib.rs‎

Lines changed: 99 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -20,145 +20,141 @@ const KERNEL_FILE_NAME: &str = "kernel-x86_64";
2020
const RAMDISK_FILE_NAME: &str = "ramdisk";
2121
const BIOS_STAGE_3: &str = "boot-stage-3";
2222
const BIOS_STAGE_4: &str = "boot-stage-4";
23+
const UEFI_BOOT_FILENAME: &str = "efi/boot/bootx64.efi";
24+
const UEFI_TFTP_BOOT_FILENAME: &str = "bootloader";
2325

24-
/// Create disk images for booting on legacy BIOS systems.
25-
pub struct BiosBoot {
26-
kernel: PathBuf,
27-
ramdisk: Option<PathBuf>,
26+
struct DiskImageFile<'a> {
27+
source: &'a PathBuf,
28+
destination: &'a str,
2829
}
2930

30-
impl BiosBoot {
31-
/// Start creating a disk image for the given bootloader ELF executable.
32-
pub fn new(kernel_path: &Path) -> Self {
33-
Self {
34-
kernel: kernel_path.to_owned(),
35-
ramdisk: None,
36-
}
31+
/// DiskImageBuilder helps create disk images for a specified set of files.
32+
/// It can currently create MBR (BIOS), GPT (UEFI), and TFTP (UEFI) images.
33+
pub struct DiskImageBuilder<'a> {
34+
files: Vec<DiskImageFile<'a>>,
35+
}
36+
37+
impl<'a> DiskImageBuilder<'a> {
38+
/// Create a new instance of DiskImageBuilder, with the specified kernel.
39+
pub fn new(kernel: &'a PathBuf) -> Self {
40+
let mut obj = Self::empty();
41+
obj.set_kernel(kernel);
42+
obj
43+
}
44+
45+
/// Create a new, empty instance of DiskImageBuilder
46+
pub fn empty() -> Self {
47+
Self { files: Vec::new() }
48+
}
49+
/// Add or replace a ramdisk to be included in the final image.
50+
pub fn set_ramdisk(&mut self, path: &'a PathBuf) {
51+
self.add_or_replace_file(path, RAMDISK_FILE_NAME);
52+
}
53+
54+
/// Add or replace a kernel to be included in the final image.
55+
fn set_kernel(&mut self, path: &'a PathBuf) {
56+
self.add_or_replace_file(path, KERNEL_FILE_NAME)
3757
}
3858

39-
/// Add a ramdisk file to the image
40-
pub fn set_ramdisk(&mut self, ramdisk_path: &Path) -> &mut Self {
41-
self.ramdisk = Some(ramdisk_path.to_owned());
42-
self
59+
/// Add or replace arbitrary files.
60+
/// NOTE: You can overwrite internal files if you choose, such as EFI/BOOT/BOOTX64.EFI
61+
/// This can be useful in situations where you want to generate an image, but not use the provided bootloader.
62+
fn add_or_replace_file(&mut self, path: &'a PathBuf, target: &'a str) {
63+
self.files.insert(
64+
0,
65+
DiskImageFile::<'a> {
66+
source: &path,
67+
destination: &target,
68+
},
69+
);
4370
}
71+
fn create_fat_filesystem_image(
72+
&self,
73+
internal_files: BTreeMap<&'a str, &'a Path>,
74+
) -> anyhow::Result<NamedTempFile> {
75+
let mut local_map = BTreeMap::new();
76+
77+
for k in internal_files {
78+
local_map.insert(k.0, k.1);
79+
}
80+
81+
for f in self.files.as_slice() {
82+
local_map.insert(f.destination, &f.source.as_path());
83+
}
84+
85+
let out_file = NamedTempFile::new().context("failed to create temp file")?;
86+
fat::create_fat_filesystem(local_map, out_file.path())
87+
.context("failed to create BIOS FAT filesystem")?;
4488

45-
/// Create a bootable UEFI disk image at the given path.
46-
pub fn create_disk_image(&self, out_path: &Path) -> anyhow::Result<()> {
89+
Ok(out_file)
90+
}
91+
92+
/// Create an MBR disk image for booting on BIOS systems.
93+
pub fn create_bios_image(&self, image_filename: &Path) -> anyhow::Result<()> {
4794
let bootsector_path = Path::new(env!("BIOS_BOOT_SECTOR_PATH"));
4895
let stage_2_path = Path::new(env!("BIOS_STAGE_2_PATH"));
96+
let stage_3_path = Path::new(env!("BIOS_STAGE_3_PATH"));
97+
let stage_4_path = Path::new(env!("BIOS_STAGE_4_PATH"));
98+
let mut internal_files = BTreeMap::new();
99+
internal_files.insert(BIOS_STAGE_3, stage_3_path);
100+
internal_files.insert(BIOS_STAGE_4, stage_4_path);
49101

50102
let fat_partition = self
51-
.create_fat_partition()
103+
.create_fat_filesystem_image(internal_files)
52104
.context("failed to create FAT partition")?;
53-
54105
mbr::create_mbr_disk(
55106
bootsector_path,
56107
stage_2_path,
57108
fat_partition.path(),
58-
out_path,
109+
image_filename,
59110
)
60111
.context("failed to create BIOS MBR disk image")?;
61112

62113
fat_partition
63114
.close()
64115
.context("failed to delete FAT partition after disk image creation")?;
65-
66116
Ok(())
67117
}
68-
69-
/// Creates an BIOS-bootable FAT partition with the kernel.
70-
fn create_fat_partition(&self) -> anyhow::Result<NamedTempFile> {
71-
let stage_3_path = Path::new(env!("BIOS_STAGE_3_PATH"));
72-
let stage_4_path = Path::new(env!("BIOS_STAGE_4_PATH"));
73-
let kernel_path = self.kernel.as_path();
74-
75-
let mut files = BTreeMap::new();
76-
files.insert(KERNEL_FILE_NAME, kernel_path);
77-
files.insert(BIOS_STAGE_3, stage_3_path);
78-
files.insert(BIOS_STAGE_4, stage_4_path);
79-
if let Some(ramdisk_path) = &self.ramdisk {
80-
files.insert(RAMDISK_FILE_NAME, ramdisk_path);
81-
}
82-
let out_file = NamedTempFile::new().context("failed to create temp file")?;
83-
fat::create_fat_filesystem(files, out_file.path())
84-
.context("failed to create BIOS FAT filesystem")?;
85-
86-
Ok(out_file)
87-
}
88-
}
89-
90-
/// Create disk images for booting on UEFI systems.
91-
pub struct UefiBoot {
92-
kernel: PathBuf,
93-
ramdisk: Option<PathBuf>,
94-
}
95-
96-
impl UefiBoot {
97-
/// Start creating a disk image for the given bootloader ELF executable.
98-
pub fn new(kernel_path: &Path) -> Self {
99-
Self {
100-
kernel: kernel_path.to_owned(),
101-
ramdisk: None,
102-
}
103-
}
104-
105-
/// Add a ramdisk file to the disk image
106-
pub fn set_ramdisk(&mut self, ramdisk_path: &Path) -> &mut Self {
107-
self.ramdisk = Some(ramdisk_path.to_owned());
108-
self
109-
}
110-
111-
/// Create a bootable UEFI disk image at the given path.
112-
pub fn create_disk_image(&self, out_path: &Path) -> anyhow::Result<()> {
118+
/// Create a GPT disk image for booting on UEFI systems.
119+
pub fn create_uefi_image(&self, image_filename: &Path) -> anyhow::Result<()> {
120+
let bootloader_path = Path::new(env!("UEFI_BOOTLOADER_PATH"));
121+
let mut internal_files = BTreeMap::new();
122+
internal_files.insert(UEFI_BOOT_FILENAME, bootloader_path);
113123
let fat_partition = self
114-
.create_fat_partition()
124+
.create_fat_filesystem_image(internal_files)
115125
.context("failed to create FAT partition")?;
116-
117-
gpt::create_gpt_disk(fat_partition.path(), out_path)
126+
gpt::create_gpt_disk(fat_partition.path(), image_filename)
118127
.context("failed to create UEFI GPT disk image")?;
119-
120128
fat_partition
121129
.close()
122130
.context("failed to delete FAT partition after disk image creation")?;
123131

124132
Ok(())
125133
}
126134

127-
/// Prepare a folder for use with booting over UEFI_PXE.
128-
///
129-
/// This places the bootloader executable under the path "bootloader". The
130-
/// DHCP server should set the filename option to that path, otherwise the
131-
/// bootloader won't be found.
132-
pub fn create_pxe_tftp_folder(&self, out_path: &Path) -> anyhow::Result<()> {
135+
/// Create a folder containing the needed files for UEFI TFTP/PXE booting.
136+
pub fn create_uefi_tftp_folder(&self, tftp_path: &Path) -> anyhow::Result<()> {
133137
let bootloader_path = Path::new(env!("UEFI_BOOTLOADER_PATH"));
134-
let ramdisk_path = self.ramdisk.as_deref();
135-
pxe::create_uefi_tftp_folder(
136-
bootloader_path,
137-
self.kernel.as_path(),
138-
ramdisk_path,
139-
out_path,
140-
)
141-
.context("failed to create UEFI PXE tftp folder")?;
142-
143-
Ok(())
144-
}
145-
146-
/// Creates an UEFI-bootable FAT partition with the kernel.
147-
fn create_fat_partition(&self) -> anyhow::Result<NamedTempFile> {
148-
let bootloader_path = Path::new(env!("UEFI_BOOTLOADER_PATH"));
149-
let kernel_path = self.kernel.as_path();
150-
let mut files = BTreeMap::new();
151-
files.insert("efi/boot/bootx64.efi", bootloader_path);
152-
files.insert(KERNEL_FILE_NAME, kernel_path);
153-
154-
if let Some(ramdisk_path) = &self.ramdisk {
155-
files.insert(RAMDISK_FILE_NAME, ramdisk_path);
138+
std::fs::create_dir_all(tftp_path)
139+
.with_context(|| format!("failed to create out dir at {}", tftp_path.display()))?;
140+
141+
let to = tftp_path.join(UEFI_TFTP_BOOT_FILENAME);
142+
std::fs::copy(bootloader_path, &to).with_context(|| {
143+
format!(
144+
"failed to copy bootloader from {} to {}",
145+
bootloader_path.display(),
146+
to.display()
147+
)
148+
})?;
149+
150+
for f in self.files.as_slice() {
151+
let to = tftp_path.join(f.destination);
152+
let result = std::fs::copy(f.source, to);
153+
if result.is_err() {
154+
return Err(anyhow::Error::from(result.unwrap_err()));
155+
}
156156
}
157157

158-
let out_file = NamedTempFile::new().context("failed to create temp file")?;
159-
fat::create_fat_filesystem(files, out_file.path())
160-
.context("failed to create UEFI FAT filesystem")?;
161-
162-
Ok(out_file)
158+
Ok(())
163159
}
164160
}

‎tests/runner/src/lib.rs‎

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,8 @@
1-
use std::{io::Write, path::Path, process::Command};
1+
use std::{
2+
io::Write,
3+
path::{Path, PathBuf},
4+
process::Command,
5+
};
26

37
const QEMU_ARGS: &[&str] = &[
48
"-device",
@@ -13,33 +17,29 @@ const QEMU_ARGS: &[&str] = &[
1317
];
1418

1519
pub fn run_test_kernel(kernel_binary_path: &str, ramdisk_path: Option<&str>) {
20+
use bootloader::DiskImageBuilder;
1621
let kernel_path = Path::new(kernel_binary_path);
17-
let ramdisk_path = match ramdisk_path {
18-
Some(rdp) => Some(Path::new(rdp)),
22+
let ramdisk_path_buf = match ramdisk_path {
23+
Some(rdp) => Some(PathBuf::from(rdp)),
1924
None => None,
2025
};
26+
let ramdisk_path_buf = ramdisk_path_buf.as_ref();
2127

2228
// create an MBR disk image for legacy BIOS booting
2329
let mbr_path = kernel_path.with_extension("mbr");
24-
let mut bios_builder = bootloader::BiosBoot::new(kernel_path);
25-
26-
// create a GPT disk image for UEFI booting
2730
let gpt_path = kernel_path.with_extension("gpt");
28-
let mut uefi_builder = bootloader::UefiBoot::new(kernel_path);
31+
let tftp_path = kernel_path.with_extension(".tftp");
32+
let kernel_path_buf = kernel_path.to_path_buf();
33+
let mut image_builder = DiskImageBuilder::new(&kernel_path_buf);
2934

3035
// Set ramdisk for test, if supplied.
31-
if let Some(rdp) = ramdisk_path {
32-
bios_builder.set_ramdisk(rdp);
33-
uefi_builder.set_ramdisk(rdp);
36+
if let Some(rdp) = ramdisk_path_buf {
37+
image_builder.set_ramdisk(&rdp);
3438
}
3539

36-
bios_builder.create_disk_image(&mbr_path).unwrap();
37-
uefi_builder.create_disk_image(&gpt_path).unwrap();
38-
39-
// create a TFTP folder with the kernel executable and UEFI bootloader for
40-
// UEFI PXE booting
41-
let tftp_path = kernel_path.with_extension(".tftp");
42-
uefi_builder.create_pxe_tftp_folder(&tftp_path).unwrap();
40+
image_builder.create_bios_image(&mbr_path).unwrap();
41+
image_builder.create_uefi_image(&gpt_path).unwrap();
42+
image_builder.create_uefi_tftp_folder(&tftp_path).unwrap();
4343

4444
run_test_kernel_on_uefi(&gpt_path);
4545
run_test_kernel_on_bios(&mbr_path);

0 commit comments

Comments
 (0)