Add disk wrapper to ensure block aligned I/O

This commit is contained in:
Jeremy Soller
2022-09-07 10:29:15 -06:00
parent 9666c62036
commit 8d2d4cd62d
2 changed files with 204 additions and 104 deletions
+125
View File
@@ -0,0 +1,125 @@
use std::{
cmp,
convert::TryInto,
fs::{File, OpenOptions},
io::{Read, Result, Seek, SeekFrom, Write},
path::Path,
};
#[derive(Debug)]
pub struct DiskWrapper {
disk: File,
size: u64,
block: Box<[u8]>,
seek: u64,
}
impl DiskWrapper {
pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
let disk = OpenOptions::new()
.read(true)
.write(true)
.open(path)?;
let metadata = disk.metadata()?;
let size = metadata.len();
// TODO: get real block size: disk_metadata.blksize() works on disks but not image files
let block_size = 512;
let block = vec![0u8; block_size].into_boxed_slice();
Ok(Self {
disk,
size,
block,
seek: 0,
})
}
pub fn block_size(&self) -> usize {
self.block.len()
}
pub fn size(&self) -> u64 {
self.size
}
}
impl Read for DiskWrapper {
//TODO: improve performance by directly using block aligned parts of buf
fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
let mut i = 0;
while i < buf.len() {
let block_len: u64 = self.block.len().try_into().unwrap();
let block = self.seek / block_len;
let offset: usize = (self.seek % block_len).try_into().unwrap();
let remaining = buf.len().checked_sub(i).unwrap();
let len = cmp::min(
remaining,
self.block.len().checked_sub(offset.try_into().unwrap()).unwrap()
);
self.disk.seek(SeekFrom::Start(block.checked_mul(block_len).unwrap()))?;
self.disk.read_exact(&mut self.block)?;
buf[i..i.checked_add(len).unwrap()].copy_from_slice(
&self.block[offset..offset.checked_add(len).unwrap()]
);
i = i.checked_add(len).unwrap();
self.seek = self.seek.checked_add(len.try_into().unwrap()).unwrap();
}
Ok(i)
}
}
impl Seek for DiskWrapper {
fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
let current: i64 = self.seek.try_into().unwrap();
let end: i64 = self.size.try_into().unwrap();
self.seek = match pos {
SeekFrom::Start(offset) => {
cmp::min(self.size, offset)
},
SeekFrom::End(offset) => {
cmp::max(0, cmp::min(end, end.wrapping_add(offset))) as u64
},
SeekFrom::Current(offset) => {
cmp::max(0, cmp::min(end, current.wrapping_add(offset))) as u64
}
};
Ok(self.seek)
}
}
impl Write for DiskWrapper {
//TODO: improve performance by directly using block aligned parts of buf
fn write(&mut self, buf: &[u8]) -> Result<usize> {
let mut i = 0;
while i < buf.len() {
let block_len: u64 = self.block.len().try_into().unwrap();
let block = self.seek / block_len;
let offset: usize = (self.seek % block_len).try_into().unwrap();
let remaining = buf.len().checked_sub(i).unwrap();
let len = cmp::min(
remaining,
self.block.len().checked_sub(offset.try_into().unwrap()).unwrap()
);
self.disk.seek(SeekFrom::Start(block.checked_mul(block_len).unwrap()))?;
self.disk.read_exact(&mut self.block)?;
self.block[offset..offset.checked_add(len).unwrap()].copy_from_slice(
&buf[i..i.checked_add(len).unwrap()]
);
self.disk.seek(SeekFrom::Start(block.checked_mul(block_len).unwrap()))?;
self.disk.write_all(&mut self.block)?;
i = i.checked_add(len).unwrap();
self.seek = self.seek.checked_add(len.try_into().unwrap()).unwrap();
}
Ok(i)
}
fn flush(&mut self) -> Result<()> {
self.disk.flush()
}
}
+79 -104
View File
@@ -12,10 +12,12 @@ extern crate syscall;
extern crate termion;
mod config;
mod disk_wrapper;
pub use config::Config;
pub use config::file::FileConfig;
pub use config::package::PackageConfig;
use disk_wrapper::DiskWrapper;
use failure::{Error, err_msg};
use rand::{RngCore, rngs::OsRng};
@@ -27,8 +29,7 @@ use std::{
collections::BTreeMap,
env,
fs,
io::{self, stderr, Read, Seek, SeekFrom, Write},
os::unix::fs::MetadataExt,
io::{self, stderr, Seek, SeekFrom, Write},
path::Path,
process::{self, Command},
str::FromStr,
@@ -406,14 +407,9 @@ pub fn with_whole_disk<P, F, T>(disk_path: P, bootloader_bios: &[u8], bootloader
// Open disk and read metadata
eprintln!("Opening disk {}", disk_path.as_ref().display());
let mut disk_file = fs::OpenOptions::new()
.read(true)
.write(true)
.open(disk_path.as_ref())?;
let disk_metadata = disk_file.metadata()?;
let disk_size = disk_metadata.len();
// TODO: get real block size: disk_metadata.blksize() works on disks but not image files
let block_size = 512;
let mut disk_file = DiskWrapper::open(disk_path.as_ref())?;
let disk_size = disk_file.size();
let block_size = disk_file.block_size() as u64;
let gpt_block_size = match block_size {
512 => gpt::disk::LogicalBlockSize::Lb512,
_ => {
@@ -422,6 +418,7 @@ pub fn with_whole_disk<P, F, T>(disk_path: P, bootloader_bios: &[u8], bootloader
}
};
// Calculate partition offsets
let gpt_reserved = 34 * 512; // GPT always reserves 34 512-byte sectors
@@ -440,117 +437,95 @@ pub fn with_whole_disk<P, F, T>(disk_path: P, bootloader_bios: &[u8], bootloader
// Format and install BIOS partition
{
let mut bios_partition = {
// We write the BIOS partition and partition tables in memory to avoid non block aligned I/O
let mut bios_partition = Box::new(io::Cursor::new(vec![0u8; bios_size as usize]));
// Write BIOS bootloader to disk
eprintln!("Write bootloader with size {:#x}", bootloader_bios.len());
disk_file.seek(SeekFrom::Start(0))?;
disk_file.write_all(&bootloader_bios)?;
// Write BIOS bootloader to disk
eprintln!("Write bootloader with size {:#x}", bootloader_bios.len());
bios_partition.write_all(&bootloader_bios)?;
// Replace MBR tables with protective MBR
let mbr_blocks = ((disk_size + block_size - 1) / block_size) - 1;
eprintln!("Writing protective MBR with disk blocks {:#x}", mbr_blocks);
gpt::mbr::ProtectiveMBR::with_lb_size(mbr_blocks as u32)
.update_conservative(&mut disk_file)?;
// Replace MBR tables with protective MBR
let mbr_blocks = ((disk_size + block_size - 1) / block_size) - 1;
eprintln!("Writing protective MBR with disk blocks {:#x}", mbr_blocks);
gpt::mbr::ProtectiveMBR::with_lb_size(mbr_blocks as u32)
.update_conservative(&mut bios_partition)?;
// Open disk, mark it as not initialized
let mut gpt_disk = gpt::GptConfig::new()
.initialized(false)
.writable(true)
.logical_block_size(gpt_block_size)
.create_from_device(Box::new(&mut disk_file), None)?;
// Open disk, mark it as not initialized
let mut gpt_disk = gpt::GptConfig::new()
.initialized(false)
.writable(true)
.logical_block_size(gpt_block_size)
.create_from_device(bios_partition, None)?;
// Add BIOS boot partition
let mut partitions = BTreeMap::new();
let mut partition_id = 1;
partitions.insert(partition_id, gpt::partition::Partition {
part_type_guid: gpt::partition_types::BIOS,
part_guid: uuid::Uuid::new_v4(),
first_lba: bios_start,
last_lba: bios_end,
flags: 0, // TODO
name: "BIOS".to_string(),
});
partition_id += 1;
// Add BIOS boot partition
let mut partitions = BTreeMap::new();
let mut partition_id = 1;
partitions.insert(partition_id, gpt::partition::Partition {
part_type_guid: gpt::partition_types::BIOS,
part_guid: uuid::Uuid::new_v4(),
first_lba: bios_start,
last_lba: bios_end,
flags: 0, // TODO
name: "BIOS".to_string(),
});
partition_id += 1;
// Add RedoxFS partition
partitions.insert(partition_id, gpt::partition::Partition {
//TODO: Use REDOX_REDOXFS type (needs GPT crate changes)
part_type_guid: gpt::partition_types::LINUX_FS,
part_guid: uuid::Uuid::new_v4(),
first_lba: redoxfs_start,
last_lba: redoxfs_end,
flags: 0,
name: "REDOX".to_string(),
});
partition_id += 1;
// Add RedoxFS partition
partitions.insert(partition_id, gpt::partition::Partition {
//TODO: Use REDOX_REDOXFS type (needs GPT crate changes)
part_type_guid: gpt::partition_types::LINUX_FS,
part_guid: uuid::Uuid::new_v4(),
first_lba: redoxfs_start,
last_lba: redoxfs_end,
flags: 0,
name: "REDOX".to_string(),
});
partition_id += 1;
// Add EFI boot partition
partitions.insert(partition_id, gpt::partition::Partition {
part_type_guid: gpt::partition_types::EFI,
part_guid: uuid::Uuid::new_v4(),
first_lba: efi_start,
last_lba: efi_end,
flags: 0, // TODO
name: "EFI".to_string(),
});
// Add EFI boot partition
partitions.insert(partition_id, gpt::partition::Partition {
part_type_guid: gpt::partition_types::EFI,
part_guid: uuid::Uuid::new_v4(),
first_lba: efi_start,
last_lba: efi_end,
flags: 0, // TODO
name: "EFI".to_string(),
});
eprintln!("Writing GPT tables: {:#?}", partitions);
eprintln!("Writing GPT tables: {:#?}", partitions);
// Initialize GPT table
gpt_disk.update_partitions(partitions)?;
// Initialize GPT table
gpt_disk.update_partitions(partitions)?;
// Write partition layout, returning disk file
gpt_disk.write()?
};
eprintln!("Copying BIOS partition to disk with size {:#x}", bios_size);
for block in 0..bios_size / block_size {
let offset = block * block_size;
let mut block = vec![0; block_size as usize];
bios_partition.seek(SeekFrom::Start(offset))?;
bios_partition.read_exact(&mut block)?;
disk_file.seek(SeekFrom::Start(offset))?;
disk_file.write_all(&block)?;
}
// Write partition layout, returning disk file
gpt_disk.write()?;
}
// Format and install EFI partition
{
// We write the EFI partition and partition tables in memory to avoid non block aligned I/O
let mut efi_partition = vec![0u8; efi_size as usize];
{
eprintln!("Formatting EFI partition with size {:#x}", efi_size);
fatfs::format_volume(io::Cursor::new(&mut efi_partition), fatfs::FormatVolumeOptions::new())?;
eprintln!("Opening EFI partition");
let fs = fatfs::FileSystem::new(io::Cursor::new(&mut efi_partition), fatfs::FsOptions::new())?;
eprintln!("Creating EFI directory");
let root_dir = fs.root_dir();
root_dir.create_dir("EFI")?;
eprintln!("Creating EFI/BOOT directory");
let efi_dir = root_dir.open_dir("EFI")?;
efi_dir.create_dir("BOOT")?;
eprintln!("Writing EFI/BOOT/{} file with size {:#x}", bootloader_efi_name, bootloader_efi.len());
let boot_dir = efi_dir.open_dir("BOOT")?;
let mut file = boot_dir.create_file(bootloader_efi_name)?;
file.truncate()?;
file.write_all(&bootloader_efi)?;
}
eprintln!("Copying EFI partition to disk with size {:#x}", efi_size);
let mut disk_efi = fscommon::StreamSlice::new(
&mut disk_file,
efi_start * block_size,
(efi_end + 1) * block_size,
)?;
disk_efi.write_all(&efi_partition)?;
eprintln!("Formatting EFI partition with size {:#x}", efi_size);
fatfs::format_volume(&mut disk_efi, fatfs::FormatVolumeOptions::new())?;
eprintln!("Opening EFI partition");
let fs = fatfs::FileSystem::new(&mut disk_efi, fatfs::FsOptions::new())?;
eprintln!("Creating EFI directory");
let root_dir = fs.root_dir();
root_dir.create_dir("EFI")?;
eprintln!("Creating EFI/BOOT directory");
let efi_dir = root_dir.open_dir("EFI")?;
efi_dir.create_dir("BOOT")?;
eprintln!("Writing EFI/BOOT/{} file with size {:#x}", bootloader_efi_name, bootloader_efi.len());
let boot_dir = efi_dir.open_dir("BOOT")?;
let mut file = boot_dir.create_file(bootloader_efi_name)?;
file.truncate()?;
file.write_all(&bootloader_efi)?;
}
// Format and install RedoxFS partition