diff options
author | Dov Shlachter <dovs@google.com> | 2024-04-30 15:15:12 -0700 |
---|---|---|
committer | Dov Shlachter <dovs@google.com> | 2024-05-06 10:16:03 -0700 |
commit | ef53ff51ce0ad0a02f66192645badd9ac7728257 (patch) | |
tree | 978e2b216b41c70c1d30f40947b6f7b4574ce64e | |
parent | 611dbc99bb47faafcd5d8356f455d122660e80d5 (diff) | |
download | libbootloader-master.tar.gz |
Add functionality to TestBlockDeviceBuilder to define a block device
as a collection of partitions with either associated data or just a size.
Partition sizes are rounded up to the block size and padded with
trailing zeroes.
If the device is built via partitions, a valid GPT header and backup
header are added to the device.
Tests: new unit test passes
Change-Id: I614f7d0c48e56b34ae9eaf08e5b0c953355f0cb0
-rw-r--r-- | gbl/libgbl/tests/integration_tests.rs | 4 | ||||
-rw-r--r-- | gbl/libstorage/BUILD | 7 | ||||
-rw-r--r-- | gbl/libstorage/src/gpt.rs | 32 | ||||
-rw-r--r-- | gbl/libstorage/src/lib.rs | 2 | ||||
-rw-r--r-- | gbl/libstorage/src/testlib.rs | 239 | ||||
-rw-r--r-- | gbl/tests/BUILD | 1 |
6 files changed, 245 insertions, 40 deletions
diff --git a/gbl/libgbl/tests/integration_tests.rs b/gbl/libgbl/tests/integration_tests.rs index 082ac44..e4764fe 100644 --- a/gbl/libgbl/tests/integration_tests.rs +++ b/gbl/libgbl/tests/integration_tests.rs @@ -45,8 +45,8 @@ impl TestGblOps<'_> { data: T, ) { self.block_io.push(GblTestBlockIo { - io: TestBlockIo::new(alignment, block_size, data), - max_gpt_entries: max_gpt_entries, + io: TestBlockIo::new(alignment, block_size, data.as_ref().into()), + max_gpt_entries, }) } } diff --git a/gbl/libstorage/BUILD b/gbl/libstorage/BUILD index 9f0aed0..74f8182 100644 --- a/gbl/libstorage/BUILD +++ b/gbl/libstorage/BUILD @@ -43,7 +43,9 @@ rust_library( edition = "2021", deps = [ ":libstorage", + "@crc32fast", "@gbl//libsafemath", + "@zerocopy", ], ) @@ -56,6 +58,11 @@ rust_test( ], ) +rust_test( + name = "libstorage_testlib_test", + crate = ":libstorage_testlib", +) + rust_doc_test( name = "libstorage_doc_test", crate = ":libstorage", diff --git a/gbl/libstorage/src/gpt.rs b/gbl/libstorage/src/gpt.rs index c5a44ed..92a4ba3 100644 --- a/gbl/libstorage/src/gpt.rs +++ b/gbl/libstorage/src/gpt.rs @@ -25,7 +25,7 @@ pub const GPT_NAME_LEN_U16: usize = 36; #[repr(C, packed)] #[derive(Debug, Default, Copy, Clone, AsBytes, FromBytes, FromZeroes)] -struct GptHeader { +pub struct GptHeader { pub magic: u64, pub revision: u32, pub size: u32, @@ -49,7 +49,7 @@ impl GptHeader { } /// Update the header crc32 value. - fn update_crc(&mut self) { + pub fn update_crc(&mut self) { self.crc32 = 0; self.crc32 = crc32(self.as_bytes()); } @@ -114,11 +114,10 @@ const GPT_CRC32_OFFSET: u64 = 16; const GPT_ENTRY_ALIGNMENT: u64 = align_of::<GptEntry>() as u64; const GPT_ENTRY_SIZE: u64 = size_of::<GptEntry>() as u64; const GPT_MAX_NUM_ENTRIES: u64 = 128; -const GPT_MAX_ENTRIES_SIZE: u64 = GPT_MAX_NUM_ENTRIES * GPT_ENTRY_SIZE; const GPT_HEADER_SIZE: u64 = size_of::<GptHeader>() as u64; // 92 bytes. const GPT_HEADER_SIZE_PADDED: u64 = (GPT_HEADER_SIZE + GPT_ENTRY_ALIGNMENT - 1) / GPT_ENTRY_ALIGNMENT * GPT_ENTRY_ALIGNMENT; -const GPT_MAGIC: u64 = 0x5452415020494645; +pub const GPT_MAGIC: u64 = 0x5452415020494645; enum HeaderType { Primary, @@ -128,7 +127,13 @@ enum HeaderType { #[repr(C)] #[derive(Debug, Default, Copy, Clone, AsBytes, FromBytes, FromZeroes)] struct GptInfo { + // The number of valid entries in the entries array. + // May change as partitions are added or removed. num_valid_entries: Option<NonZeroU64>, + // The maximum number of elements available in the entries array. + // Note: this is GREATER THAN OR EQUAL TO the number of valid entries + // and LESS THAN OR EQUAL TO the value of GPT_MAX_NUM_ENTRIES. + // Values other than GPT_MAX_NUM_ENTRIES are mostly used in unit tests. max_entries: u64, } @@ -308,14 +313,11 @@ impl<'a> Gpt<'a> { // Entries position for restoring. let primary_entries_blk = 2; let primary_entries_pos = SafeNum::from(primary_entries_blk) * block_size; - let secondary_entries_pos = secondary_header_pos - GPT_MAX_ENTRIES_SIZE; - let secondary_entries_blk = secondary_entries_pos / block_size; - let primary_valid = self.validate_gpt(blk_dev, scratch, HeaderType::Primary)?; let secondary_valid = self.validate_gpt(blk_dev, scratch, HeaderType::Secondary)?; let primary_header = GptHeader::from_bytes(self.primary_header); - let secondary_header = GptHeader::from_bytes(self.secondary_header.as_mut()); + let secondary_header = GptHeader::from_bytes(self.secondary_header); if !primary_valid { if !secondary_valid { return Err(StorageError::NoValidGpt); @@ -327,6 +329,7 @@ impl<'a> Gpt<'a> { primary_header.backup = secondary_header_blk.try_into()?; primary_header.entries = primary_entries_blk; primary_header.update_crc(); + write_bytes_mut(blk_dev, primary_header_pos, primary_header.as_bytes_mut(), scratch)?; write_bytes_mut( blk_dev, @@ -336,12 +339,17 @@ impl<'a> Gpt<'a> { )? } else if !secondary_valid { // Restore to secondary + let secondary_entries_pos = secondary_header_pos + - (SafeNum::from(self.info.max_entries) * core::mem::size_of::<GptEntry>()); + let secondary_entries_blk = secondary_entries_pos / block_size; + secondary_header.as_bytes_mut().clone_from_slice(primary_header.as_bytes()); self.secondary_entries.clone_from_slice(&self.primary_entries); secondary_header.current = secondary_header_blk.try_into()?; secondary_header.backup = primary_header_blk; secondary_header.entries = secondary_entries_blk.try_into()?; secondary_header.update_crc(); + write_bytes_mut( blk_dev, secondary_header_pos.try_into()?, @@ -561,6 +569,7 @@ pub(crate) mod test { // Create a header with non-max entries_count let disk = include_bytes!("../test/gpt_test_1.bin"); let mut dev = TestBlockDeviceBuilder::new().set_data(disk).build(); + let block_size: usize = dev.io.block_size.try_into().unwrap(); dev.sync_gpt().unwrap(); let gpt = gpt(&mut dev); @@ -573,14 +582,15 @@ pub(crate) mod test { gpt_header.update_crc(); // Update to primary. let primary_header = Vec::from(primary_header); - dev.io.storage[512..512 + primary_header.len()].clone_from_slice(&primary_header); + dev.io.storage[block_size..block_size + primary_header.len()] + .clone_from_slice(&primary_header); // Corrupt secondary. Sync ok - dev.io.storage[disk.len() - 512..].fill(0); + dev.io.storage[disk.len() - block_size..].fill(0); dev.sync_gpt().unwrap(); // Corrup primary. Sync ok - dev.io.storage[512..1024].fill(0); + dev.io.storage[block_size..(block_size * 2)].fill(0); dev.sync_gpt().unwrap(); } diff --git a/gbl/libstorage/src/lib.rs b/gbl/libstorage/src/lib.rs index 871540b..40d4a1c 100644 --- a/gbl/libstorage/src/lib.rs +++ b/gbl/libstorage/src/lib.rs @@ -111,7 +111,7 @@ use core::cmp::min; // Selective export of submodule types. mod gpt; use gpt::Gpt; -pub use gpt::{GptEntry, GPT_NAME_LEN_U16}; +pub use gpt::{GptEntry, GptHeader, GPT_MAGIC, GPT_NAME_LEN_U16}; use safemath::SafeNum; diff --git a/gbl/libstorage/src/testlib.rs b/gbl/libstorage/src/testlib.rs index e79a62d..21c776f 100644 --- a/gbl/libstorage/src/testlib.rs +++ b/gbl/libstorage/src/testlib.rs @@ -13,10 +13,13 @@ // limitations under the License. pub use gbl_storage::{ alignment_scratch_size, is_aligned, is_buffer_aligned, required_scratch_size, AsBlockDevice, - AsMultiBlockDevices, BlockIo, + AsMultiBlockDevices, BlockIo, GptEntry, GptHeader, GPT_MAGIC, GPT_NAME_LEN_U16, }; +use crc32fast::Hasher; use safemath::SafeNum; +use std::collections::BTreeMap; +use zerocopy::AsBytes; /// Helper `gbl_storage::BlockIo` struct for TestBlockDevice. pub struct TestBlockIo { @@ -33,14 +36,8 @@ pub struct TestBlockIo { } impl TestBlockIo { - pub fn new<T: AsRef<[u8]>>(block_size: u64, alignment: u64, data: T) -> Self { - Self { - block_size: block_size, - alignment: alignment, - storage: Vec::from(data.as_ref()), - num_writes: 0, - num_reads: 0, - } + pub fn new(block_size: u64, alignment: u64, data: Vec<u8>) -> Self { + Self { block_size, alignment, storage: data, num_writes: 0, num_reads: 0 } } fn check_alignment(&mut self, buffer: &[u8]) -> bool { @@ -114,14 +111,29 @@ impl Default for TestBlockDevice { } } -/// A description of the backing data store for a block device. -/// Can either describe explicit data the device is initialized with -/// OR a size in bytes if the device can be initialized in a blank state. -enum BackingStore<'a> { +/// A description of the backing data store for a block device or partition. +/// Can either describe explicit data the device or partition is initialized with +/// OR a size in bytes if the device or partition can be initialized in a blank state. +#[derive(Copy, Clone)] +pub enum BackingStore<'a> { Data(&'a [u8]), Size(usize), } +impl<'a> BackingStore<'a> { + fn size(&self) -> usize { + match self { + Self::Data(slice) => slice.len(), + Self::Size(size) => *size, + } + } +} + +enum DiskDescription<'a> { + Disk(BackingStore<'a>), + Partitions(BTreeMap<&'static str, BackingStore<'a>>), +} + /// Builder struct for TestBlockDevice. /// Most tests will want either: /// 1) A blank device of a reasonable size OR @@ -141,7 +153,7 @@ pub struct TestBlockDeviceBuilder<'a> { block_size: u64, max_gpt_entries: u64, alignment: u64, - backing_store: BackingStore<'a>, + disk_description: DiskDescription<'a>, scratch_size: Option<usize>, } @@ -159,7 +171,9 @@ impl<'a> TestBlockDeviceBuilder<'a> { block_size: Self::DEFAULT_BLOCK_SIZE, max_gpt_entries: Self::MAX_GPT_ENTRIES, alignment: Self::DEFAULT_ALIGNMENT, - backing_store: BackingStore::Size((Self::DEFAULT_BLOCK_SIZE * 32) as usize), + disk_description: DiskDescription::Disk(BackingStore::Size( + (Self::DEFAULT_BLOCK_SIZE * 32) as usize, + )), scratch_size: None, } } @@ -192,21 +206,45 @@ impl<'a> TestBlockDeviceBuilder<'a> { /// When built, the TestBlockDevice will have a blank backing store of size `size`. /// The default is `DEFAULT_BLOCK_SIZE` * 32. /// - /// Note: This option is mutually exclusive with set_data. - /// If set_data has been called, set_size overrides - /// that customization. + /// Note: This option is mutually exclusive with `set_data` and `add_partition`. + /// If `set_data` or `add_partition` have been called, `set_size` overrides + /// those customizations. pub fn set_size(mut self, size: usize) -> Self { - self.backing_store = BackingStore::Size(size); + self.disk_description = DiskDescription::Disk(BackingStore::Size(size)); self } /// Sets the block device's backing data to the provided slice. /// - /// Note: This option is mutually exclusive with set_size. - /// If set_size has been called, set_data overrides - /// that customization. + /// Note: This option is mutually exclusive with `set_size` and `add_partition`. + /// If `set_size` or `add_partition` have been called, `set_data` overrides + /// those customizations. pub fn set_data(mut self, data: &'a [u8]) -> Self { - self.backing_store = BackingStore::Data(data); + self.disk_description = DiskDescription::Disk(BackingStore::Data(data)); + self + } + + /// Adds a partition description. + /// Partitions can be defined either with a specific backing store + /// from a slice OR from a specific size in bytes. + /// Partition sizes are rounded up to full blocks. + /// If the same partition name is added multiple times, + /// the last definition is used. + /// + /// Note: explicitly added partitions are mutually exclusive with + /// `set_size` and `set_data`. + /// If either have been called, `add_partition` overrides that customization. + pub fn add_partition(mut self, name: &'static str, backing: BackingStore<'a>) -> Self { + match self.disk_description { + DiskDescription::Disk(_) => { + let mut map = BTreeMap::new(); + map.insert(name, backing); + self.disk_description = DiskDescription::Partitions(map); + } + DiskDescription::Partitions(ref mut map) => { + map.insert(name, backing); + } + }; self } @@ -224,9 +262,12 @@ impl<'a> TestBlockDeviceBuilder<'a> { /// Consumes the builder and generates a TestBlockDevice /// with the desired customizations. pub fn build(self) -> TestBlockDevice { - let storage = match self.backing_store { - BackingStore::Data(slice) => Vec::from(slice), - BackingStore::Size(size) => vec![0u8; size], + let storage = match self.disk_description { + DiskDescription::Disk(BackingStore::Data(slice)) => Vec::from(slice), + DiskDescription::Disk(BackingStore::Size(size)) => vec![0u8; size], + DiskDescription::Partitions(partitions) => { + partitions_to_disk_data(&partitions, self.block_size as usize) + } }; assert!(storage.len() % (self.block_size as usize) == 0); let mut io = TestBlockIo::new(self.block_size, self.alignment, storage); @@ -242,6 +283,130 @@ impl<'a> TestBlockDeviceBuilder<'a> { } } +fn str_to_utf16_entry_name(name: &str) -> [u16; GPT_NAME_LEN_U16] { + assert!(name.len() < GPT_NAME_LEN_U16); + let mut data = [0; GPT_NAME_LEN_U16]; + let tmp: Vec<u16> = name.encode_utf16().collect(); + for (d, t) in std::iter::zip(data.iter_mut(), tmp) { + *d = t; + } + data +} + +fn pad_to_block_size(store: &mut Vec<u8>, block_size: usize) { + let delta = (block_size - store.len() % block_size) % block_size; + for _ in 0..delta { + store.push(0); + } +} + +fn add_blocks(store: &mut Vec<u8>, data: &[u8], block_size: usize) { + store.extend(data.iter()); + pad_to_block_size(store, block_size); +} + +fn pad_bytes(store: &mut Vec<u8>, size: usize, block_size: usize) { + for _ in 0..size { + store.push(0); + } + pad_to_block_size(store, block_size); +} + +fn partitions_to_disk_data( + partitions: &BTreeMap<&'static str, BackingStore>, + block_size: usize, +) -> Vec<u8> { + let gpt_max_entries = 128; + assert!(partitions.len() <= gpt_max_entries); + let entry_blocks: u64 = ((SafeNum::from(partitions.len()) * std::mem::size_of::<GptEntry>()) + .round_up(block_size) + / block_size) + .try_into() + .unwrap(); + let mut block = entry_blocks + + 1 // Protective MBR + + 1 // Primary GPT header + ; + // Leading mbr + let mut store = vec![0; block_size]; + let mut header = GptHeader { + magic: GPT_MAGIC, + current: 1, + size: std::mem::size_of::<GptHeader>() as u32, + first: block, + entries: 2, + entries_count: std::cmp::min(partitions.len(), gpt_max_entries) as u32, + entries_size: std::mem::size_of::<GptEntry>() as u32, + ..Default::default() + }; + + // Define gpt entry structures + let entries: Vec<GptEntry> = partitions + .iter() + .take(gpt_max_entries) + .map(|(k, v)| { + let last = (SafeNum::from(v.size()).round_up(block_size) / block_size + block - 1) + .try_into() + .unwrap(); + let mut entry = GptEntry { + part_type: Default::default(), + guid: Default::default(), + first: block, + last, + flags: 0, + name: str_to_utf16_entry_name(k), + }; + entry.guid[0] = block as u8; + block = last + 1; + entry + }) + .collect(); + + // Patch last fields of header + header.last = block - 1; + header.backup = block + entry_blocks; + header.entries_crc = entries + .iter() + .fold(Hasher::new(), |mut h, e| { + h.update(e.as_bytes()); + h + }) + .finalize(); + header.update_crc(); + + // Primary header + add_blocks(&mut store, header.as_bytes(), block_size); + + // Primary entries + for e in &entries { + store.extend(e.as_bytes()); + } + pad_to_block_size(&mut store, block_size); + + // Partition store + for p in partitions.values() { + match p { + BackingStore::Data(d) => add_blocks(&mut store, d, block_size), + BackingStore::Size(s) => pad_bytes(&mut store, *s, block_size), + }; + } + + // Backup entries + let backup_entries_block = store.len() / block_size; + for e in entries { + store.extend(e.as_bytes()); + } + pad_to_block_size(&mut store, block_size); + // Tweak header to make it the backup. + header.current = header.backup; + header.backup = 1; + header.entries = backup_entries_block.try_into().unwrap(); + header.update_crc(); + add_blocks(&mut store, header.as_bytes(), block_size); + + store +} + /// Simple RAM based multi-block device used for unit tests. pub struct TestMultiBlockDevices(pub Vec<TestBlockDevice>); @@ -258,3 +423,25 @@ impl AsMultiBlockDevices for TestMultiBlockDevices { Ok(()) } } + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_builder_partitions() { + let data: [u8; 8] = [1, 2, 3, 4, 5, 6, 7, 8]; + let mut actual: [u8; 8] = Default::default(); + let mut block_dev = TestBlockDeviceBuilder::new() + .add_partition("squid", BackingStore::Data(&data)) + .add_partition("clam", BackingStore::Size(28)) + .build(); + + assert!(block_dev.sync_gpt().is_ok()); + assert!(block_dev.read_gpt_partition("squid", 0, actual.as_mut_slice()).is_ok()); + assert_eq!(actual, data); + + assert!(block_dev.read_gpt_partition("clam", 0, actual.as_mut_slice()).is_ok()); + assert_eq!(actual, [0u8; 8]); + } +} diff --git a/gbl/tests/BUILD b/gbl/tests/BUILD index 833d7a6..3a9fbda 100644 --- a/gbl/tests/BUILD +++ b/gbl/tests/BUILD @@ -26,6 +26,7 @@ test_suite( "@gbl//libsafemath:libsafemath_test", "@gbl//libstorage:libstorage_doc_test", "@gbl//libstorage:libstorage_test", + "@gbl//libstorage:libstorage_testlib_test", "@gbl//third_party/libzbi:libzbi_test", ], ) |