// Copyright 2023, The Android Open Source Project // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. use crate::{BlockIo, Disk, Result}; use core::{ array::from_fn, cmp::min, convert::TryFrom, default::Default, fmt::{Debug, Formatter}, mem::size_of, num::NonZeroU64, ops::{Deref, DerefMut}, str::from_utf8, }; use crc32fast::Hasher; use gbl_async::block_on; use liberror::{Error, GptError}; use safemath::SafeNum; use zerocopy::{AsBytes, ByteSlice, FromBytes, FromZeroes, Ref}; /// Number of bytes in GUID. pub const GPT_GUID_LEN: usize = 16; /// The maximum number of UTF-16 characters in a GPT partition name, including termination. pub const GPT_NAME_LEN_U16: usize = 36; const GPT_NAME_LEN_U8: usize = 2 * GPT_GUID_LEN; /// The top-level GPT header. #[repr(C, packed)] #[derive(Debug, Default, Copy, Clone, AsBytes, FromBytes, FromZeroes, PartialEq, Eq)] pub struct GptHeader { /// Magic bytes; must be [GPT_MAGIC]. pub magic: u64, /// Header version. pub revision: u32, /// Header size in bytes. pub size: u32, /// CRC of the first `size` bytes, calculated with this field zeroed. pub crc32: u32, /// Reserved; must be set to 0. pub reserved0: u32, /// The on-disk block location of this header. pub current: u64, /// The on-disk block location of the other header. pub backup: u64, /// First usable block for partition contents. pub first: u64, /// Last usable block for partition contents (inclusive). pub last: u64, /// Disk GUID. pub guid: [u8; GPT_GUID_LEN], /// Starting block for the partition entries array. pub entries: u64, /// Number of partition entries. pub entries_count: u32, /// The size of each partition entry in bytes. pub entries_size: u32, /// CRC of the partition entries array. pub entries_crc: u32, } impl GptHeader { /// Casts a bytes slice into a mutable GptHeader structure. pub fn from_bytes_mut(bytes: &mut [u8]) -> &mut GptHeader { Ref::<_, GptHeader>::new_from_prefix(bytes).unwrap().0.into_mut() } /// Computes the actual crc32 value. fn calculate_header_crc(&self) -> u32 { let mut hasher = Hasher::new(); hasher.update(&self.as_bytes()[..GPT_CRC32_OFFSET]); hasher.update(&[0u8; size_of::()]); hasher.update(&self.as_bytes()[GPT_CRC32_OFFSET + size_of::()..]); hasher.finalize() } /// Update the header crc32 value. pub fn update_crc(&mut self) { self.crc32 = self.calculate_header_crc(); } /// Updates entries and header crc according to the given entries buffer. fn update_entries_crc(&mut self, entries: &[u8]) { let size = SafeNum::from(self.entries_count) * self.entries_size; self.entries_crc = crc32(&entries[..size.try_into().unwrap()]); self.update_crc(); } } /// Computes the number of blocks for the 128 partition entries reserved space in GPT. fn gpt_entries_blk(block_size: u64) -> Result { let size = u64::try_from(GPT_MAX_NUM_ENTRIES_SIZE).unwrap(); match size % block_size { 0 => Ok(size / block_size), _ => Err(Error::InvalidInput), } } /// Checks a header against a block device. /// /// # Args /// /// * `io`: An implementation of [BlockIo], /// * `header`: The GPT header to verify. /// * `is_primary`: If the header is a primary header. fn check_header(io: &mut impl BlockIo, header: &GptHeader, is_primary: bool) -> Result<()> { let num_blks = SafeNum::from(io.info().num_blocks); let blk_sz = io.info().block_size; // GPT spec requires that at least 128 entries worth of space be reserved. let min_reserved_entries_blk = gpt_entries_blk(blk_sz)?; // Minimum space needed: 2 * (header + entries) + MBR. let min_disk_blks: u64 = ((min_reserved_entries_blk + 1) * 2 + 1).try_into().unwrap(); if min_disk_blks > u64::try_from(num_blks).unwrap() { return Err(Error::GptError(GptError::DiskTooSmall)); } if header.magic != GPT_MAGIC { return Err(Error::GptError(GptError::IncorrectMagic(header.magic))); } if header.calculate_header_crc() != header.crc32 { return Err(Error::GptError(GptError::IncorrectHeaderCrc)); } if header.size != size_of::().try_into().unwrap() { return Err(Error::GptError(GptError::UnexpectedHeaderSize { actual: header.size, expect: size_of::(), })); } if header.entries_size != size_of::().try_into().unwrap() { return Err(Error::GptError(GptError::UnexpectedEntrySize { actual: header.entries_size, expect: size_of::(), })); } // Checks first/last usable block. // // Assuming maximum range where partition entries are adjacent to GPT headers. // // Should leave a minimum space for MBR + primary header + primary entries before. let min_first: u64 = (min_reserved_entries_blk + 2).try_into().unwrap(); // Should leave a minimum space for secondary header + secondary entries space after. let max_last: u64 = (num_blks - 1 - min_reserved_entries_blk - 1).try_into().unwrap(); if header.first > header.last + 1 || header.first < min_first || header.last > max_last { return Err(Error::GptError(GptError::InvalidFirstLastUsableBlock { first: header.first, last: header.last, range: (min_first, max_last), })); } // Checks entries starting block. if is_primary { // For primary header, entries must be before first usable block and can hold up to // `GPT_MAX_NUM_ENTRIES` entries let right: u64 = (SafeNum::from(header.first) - min_reserved_entries_blk).try_into().unwrap(); if !(header.entries >= 2 && header.entries <= right) { return Err(Error::GptError(GptError::InvalidPrimaryEntriesStart { value: header.entries, expect_range: (2, right), })); } } else { // For secondary header, entries must be after last usable block and can hold up to // `GPT_MAX_NUM_ENTRIES` entries. if !(header.entries > header.last && header.entries <= max_last + 1) { return Err(Error::GptError(GptError::InvalidSecondaryEntriesStart { value: header.entries, expect_range: (header.last + 1, max_last + 1), })); } } if header.entries_count > GPT_MAX_NUM_ENTRIES.try_into().unwrap() { return Err(Error::GptError(GptError::NumberOfEntriesOverflow { entries: header.entries_count, max_allowed: GPT_MAX_NUM_ENTRIES, })); } Ok(()) } /// Verifies the given entries against a verifed GPT header. /// /// # Args /// /// * `header`: The verified GPT header corresponding to the entries. /// * `entries`: The buffer containing the entries. fn check_entries(header: &GptHeader, entries: &[u8]) -> Result<()> { // Checks entries CRC. assert!(header.entries_count <= GPT_MAX_NUM_ENTRIES.try_into().unwrap()); let entries_size: usize = (SafeNum::from(header.entries_count) * GPT_ENTRY_SIZE).try_into().unwrap(); let entries = entries.get(..entries_size).ok_or(Error::GptError(GptError::EntriesTruncated))?; if header.entries_crc != crc32(entries) { return Err(Error::GptError(GptError::IncorrectEntriesCrc)); } // Checks each entry. let entries = Ref::<_, [GptEntry]>::new_slice(entries) .ok_or(Error::GptError(GptError::EntriesTruncated))? .into_slice(); let entries = &entries[..header.entries_count.try_into().unwrap()]; for (idx, ele) in entries.iter().take_while(|v| !v.is_null()).enumerate() { // Error information uses 1-base partition index. let idx = idx.checked_add(1).unwrap(); let (first, last) = (ele.first, ele.last); if first > last + 1 || last > header.last || first < header.first { return Err(Error::GptError(GptError::InvalidPartitionRange { idx, part_range: (first, last), usable_range: (header.first, header.last), })); } else if ele.part_type == [0u8; GPT_GUID_LEN] { return Err(Error::GptError(GptError::ZeroPartitionTypeGUID { idx })); } else if ele.guid == [0u8; GPT_GUID_LEN] { return Err(Error::GptError(GptError::ZeroPartitionUniqueGUID { idx })); } } // Checks overlap between partition ranges. // Sorts an index array because we don't want to modify input. let mut sorted_indices: [u8; GPT_MAX_NUM_ENTRIES] = from_fn(|i| i.try_into().unwrap()); sorted_indices.sort_unstable_by_key(|v| match entries.get(usize::try_from(*v).unwrap()) { Some(v) if !v.is_null() => v.first, _ => u64::MAX, }); let actual = entries.iter().position(|v| v.is_null()).unwrap_or(entries.len()); if actual > 1 { for i in 0..actual - 1 { let prev: usize = sorted_indices[i].try_into().unwrap(); let next: usize = sorted_indices[i + 1].try_into().unwrap(); if entries[prev].last >= entries[next].first { return Err(Error::GptError(GptError::PartitionRangeOverlap { prev: (prev + 1, entries[prev].first, entries[prev].last), next: (next + 1, entries[next].first, entries[next].last), })); } } } Ok(()) } /// GptEntry is the partition entry data structure in the GPT. #[repr(C, packed)] #[derive(Debug, Copy, Clone, AsBytes, FromBytes, FromZeroes, PartialEq)] pub struct GptEntry { /// Partition type GUID. pub part_type: [u8; GPT_GUID_LEN], /// Unique partition GUID. pub guid: [u8; GPT_GUID_LEN], /// First block. pub first: u64, /// Last block (inclusive). pub last: u64, /// Partition flags. pub flags: u64, /// Partition name in UTF-16. pub name: [u16; GPT_NAME_LEN_U16], } impl GptEntry { /// Return the partition entry size in blocks. pub fn blocks(&self) -> Result { // Must perform "+1" first before subtracting `self.first`. Otherwise if partition size is // zero, where `self.first > self.last`, arithmetic will overflow. u64::try_from(SafeNum::from(self.last) + 1 - self.first).map_err(Into::into) } /// Return whether this is a `NULL` entry. The first null entry marks the end of the partition /// entries. fn is_null(&self) -> bool { self.first == 0 && self.last == 0 } /// Decode the partition name into a string. A length N utf16 string can be at most 2N utf8 /// bytes. Therefore, a safe size of `buffer` is 2*GPT_NAME_LEN_U16 = 72. pub fn name_to_str<'a>(&self, buffer: &'a mut [u8]) -> Result<&'a str> { let mut index = 0; for c in char::decode_utf16(self.name) { match c.unwrap_or(char::REPLACEMENT_CHARACTER) { '\0' => break, c if c.len_utf8() <= buffer[index..].len() => { index += c.encode_utf8(&mut buffer[index..]).len() } _ => return Err(Error::InvalidInput), // Not enough space in `buffer`. } } // SAFETY: // _unchecked should be OK here since we wrote each utf8 byte ourselves, // but it's just an optimization, checked version would be fine also. unsafe { Ok(core::str::from_utf8_unchecked(&buffer[..index])) } } /// Checks if the partition name is the same as the given. pub fn match_name(&self, part: &str) -> Result { Ok(self.name_to_str(&mut [0u8; GPT_NAME_LEN_U16 * 2][..])? == part) } } impl core::fmt::Display for GptEntry { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { // Format: partition name: "abc", [first, last]: [123, 456] let mut name_conversion_buffer = [0u8; GPT_NAME_LEN_U16 * 2]; let name = self.name_to_str(&mut name_conversion_buffer).map_err(|_| core::fmt::Error)?; // Note: The bracket around `{ self.first }` is for forcing a copy of the field because // GptEntry is a packed structure. write!(f, "partition: \"{}\", first: {}, last: {}", name, { self.first }, { self.last }) } } // core::mem::offset_of!(GptHeader, crc32) is unsatble feature and rejected by the compiler in our // settings. We pre-compute the value here. const GPT_CRC32_OFFSET: usize = 16; const GPT_ENTRY_SIZE: usize = size_of::(); const GPT_MAX_NUM_ENTRIES: usize = 128; const GPT_MAX_NUM_ENTRIES_SIZE: usize = GPT_MAX_NUM_ENTRIES * GPT_ENTRY_SIZE; /// GPT header magic bytes ("EFI PART" in ASCII). pub const GPT_MAGIC: u64 = 0x5452415020494645; enum HeaderType { Primary, Secondary, } /// `Partition` contains information about a GPT partition. #[derive(Debug, Copy, Clone, PartialEq)] pub struct Partition { entry: GptEntry, block_size: u64, decoded_name: Option<([u8; GPT_NAME_LEN_U8], usize)>, } impl Partition { /// Creates a new instance. fn new(entry: GptEntry, block_size: u64) -> Self { let mut buf = [0u8; GPT_NAME_LEN_U8]; let decoded_name = match entry.name_to_str(&mut buf[..]).ok().map(|v| v.len()) { Some(len) => Some((buf, len)), _ => None, }; Self { entry, block_size, decoded_name } } /// Gets the decoded partition name. pub fn name(&self) -> Option<&str> { // Correct by construction. `from_utf8` should not fail. self.decoded_name.as_ref().map(|(buf, sz)| from_utf8(&buf[..*sz]).unwrap()) } /// Returns the partition size in bytes. pub fn size(&self) -> Result { u64::try_from(SafeNum::from(self.entry.blocks()?) * self.block_size).map_err(Error::from) } /// Returns the block size of this partition. pub fn block_size(&self) -> u64 { self.block_size } /// Returns the partition entry structure in the GPT header. pub fn gpt_entry(&self) -> &GptEntry { &self.entry } /// Returns the partition's absolute start/end offset in number of bytes. pub fn absolute_range(&self) -> Result<(u64, u64)> { let start = SafeNum::from(self.entry.first) * self.block_size; let end = (SafeNum::from(self.entry.last) + 1) * self.block_size; Ok((start.try_into()?, end.try_into()?)) } /// Checks a given sub range and returns its absolute offset. pub fn check_range(&self, off: u64, size: u64) -> Result { let off = SafeNum::from(off); let end: u64 = (off + size).try_into()?; match end > self.size()? { true => Err(Error::BadIndex(end as usize)), _ => Ok((off + self.absolute_range()?.0).try_into()?), } } } /// `PartitionIterator` iterates all GPT partition entries. pub struct PartitionIterator<'a> { entries: &'a [GptEntry], block_size: u64, idx: usize, } impl Iterator for PartitionIterator<'_> { type Item = Partition; fn next(&mut self) -> Option { let res = self .entries .get(self.idx) .filter(|v| !v.is_null()) .map(|v| Partition::new(*v, self.block_size))?; self.idx += 1; Some(res) } } /// Contains result of GPT syncing/restoration. #[derive(Copy, Clone, PartialEq, Debug, Default)] pub enum GptSyncResult { /// Both primary and secondary GPT are valid. #[default] BothValid, /// Primary GPT is invalid and restored. PrimaryRestored(Error), /// Secondary GPT is invalid and restored. SecondaryRestored(Error), /// Neither primary or secondary GPT is valid. NoValidGpt { /// Primary GPT verify error. primary: Error, /// Secondary GPT verify error. secondary: Error, }, } impl GptSyncResult { /// Combined into a result pub fn res(&self) -> Result<()> { match self { Self::NoValidGpt { primary: e, .. } => Err(*e), _ => Ok(()), } } } impl core::fmt::Display for GptSyncResult { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { match self { Self::BothValid => write!(f, "Found valid GPT."), Self::PrimaryRestored(e) => write!(f, "Primary GPT restored due to {e:?}."), Self::SecondaryRestored(e) => write!(f, "Secondary GPT restored due to {e:?}."), Self::NoValidGpt { primary, secondary } => { write!(f, "No valid GPT. primary: {primary:?}, secondary: {secondary:?}.") } } } } /// A packed wrapper of `Option` #[repr(C, packed)] #[derive(Debug, Copy, Clone, AsBytes, FromBytes, FromZeroes)] struct BlockSize(Option); /// Represents the structure of a load buffer for loading/verifying/syncing up to N GPT entries. #[repr(C, packed)] #[derive(Debug, Copy, Clone, AsBytes, FromBytes, FromZeroes)] pub struct GptLoadBufferN { // GPT doesn't care about block size. But it's easier to have it available for computing offset // and size in bytes for partitions. It's also used as a flag for indicating whether a valid // GPT is loaded. block_size: BlockSize, primary_header: GptHeader, secondary_header: GptHeader, primary_entries: [GptEntry; N], secondary_entries: [GptEntry; N], } impl Deref for GptLoadBufferN { type Target = [u8]; fn deref(&self) -> &Self::Target { self.as_bytes() } } impl DerefMut for GptLoadBufferN { fn deref_mut(&mut self) -> &mut Self::Target { self.as_bytes_mut() } } /// Contains references corresponding to different GPT load entities parsed from a load buffer. /// /// The structure is simply for organizing together the individual references of fields in /// `GptLoadBufferN` parsed from a raw buffer. Note that we can't parse a `Ref` /// directly from a buffer because the number of entries (length of [GptEntry]) in this case needs /// to be computed at run time based on the buffer size. struct LoadBufferRef { block_size: Ref, primary_header: Ref, secondary_header: Ref, primary_entries: Ref, secondary_entries: Ref, } impl LoadBufferRef { fn from(buffer: B) -> Self { let n = min(GPT_MAX_NUM_ENTRIES, max_supported_entries(&buffer[..]).unwrap()); let (block_size, rest) = Ref::new_from_prefix(buffer).unwrap(); let (primary_header, rest) = Ref::new_from_prefix(rest).unwrap(); let (secondary_header, rest) = Ref::new_from_prefix(rest).unwrap(); let (primary_entries, rest) = Ref::new_slice_from_prefix(rest, n).unwrap(); let (secondary_entries, _) = Ref::new_slice_from_prefix(rest, n).unwrap(); Self { block_size, primary_header, secondary_header, primary_entries, secondary_entries } } /// Unpacks into the secondary GPT header/entries fn secondary(self) -> (Ref, Ref) { (self.secondary_header, self.secondary_entries) } } /// The minimum buffer size needed for creating a [Gpt] that can load `entries` number of /// partitions. /// /// # Returns /// /// * Returns Ok(size) on success. /// * Returns Err(Error::InvalidInput) if max_entries is greater than 128. pub fn gpt_buffer_size(entries: usize) -> Result { match entries > GPT_MAX_NUM_ENTRIES { true => Err(Error::InvalidInput), _ => Ok(size_of::>() + entries * GPT_ENTRY_SIZE * 2), } } /// Computes the maximum number of entries that can be loaded if using the given buffer for [Gpt]. fn max_supported_entries(buf: &[u8]) -> Result { match buf.len() < size_of::>() { true => Err(Error::BufferTooSmall(Some(size_of::>()))), _ => Ok((buf.len() - size_of::>()) / 2 / GPT_ENTRY_SIZE), } } /// [Gpt] manages a buffer for loading, verifying and syncing GPT. pub struct Gpt { buffer: B, } impl> Gpt { /// Create an uninitialized Gpt instance from a provided buffer. /// /// The created [Gpt] can then be used in `Disk::sync_gpt()` for loading, verifying and syncing /// GPT on disk. /// /// # Args: /// /// * `buffer`: A buffer to use for loading, verifying and syncing primary and secondary GPT. /// The size of the buffer determines the maximum number of partition entries that can be /// loaded. If actual number of partitions, specified by `entries_count` in the GPT header, /// exceeds it, verification and sync will eventually fail with `Error::BufferTooSmall`. /// `gpt_buffer_size(num_entries)` can be used to compute the required size of buffer for /// loading a specific number of entries. Note that most tools and OS fix the `entries_count` /// value to the max 128 regardless of the actual number of partition entries used. Thus /// unless you have full control of GPT generation in your entire system where you can always /// ensure a smaller bound on it, it is recommended to always provide enough buffer for /// loading 128 entries. /// /// # Returns /// /// * Returns Ok(Self) on success. /// * Returns Err(Error::BufferTooSmall) if buffer is less than the minimum size. pub fn new(mut buffer: B) -> Result { max_supported_entries(&buffer[..])?; LoadBufferRef::from(&mut buffer[..]).block_size.0 = None; Ok(Self { buffer }) } /// Returns the maximum allowed entries. pub fn max_entries(&self) -> usize { max_supported_entries(&self.buffer[..]).unwrap() } /// Creates an instance of `Gpt<&mut [u8]>` that borrows the internal GPT buffer. pub fn as_borrowed(&mut self) -> Gpt<&mut [u8]> { Gpt { buffer: &mut self.buffer[..] } } /// Returns an iterator to GPT partition entries. /// /// If the object does not contain a valid GPT, the method returns Error. pub fn partition_iter(&self) -> Result { let block_size = self.check_valid()?; let entries = LoadBufferRef::from(&self.buffer[..]).primary_entries.into_slice(); Ok(PartitionIterator { entries, idx: 0, block_size }) } /// Checks if a read/write range into a GPT partition overflows and returns the range's absolute /// offset in number of bytes. pub fn check_range(&self, part_name: &str, offset: u64, size: usize) -> Result { self.find_partition(part_name)?.check_range(offset, u64::try_from(size)?) } /// Return the list of GPT entries. /// /// If there is not a valid GPT, the method returns Error. pub fn entries(&self) -> Result<&[GptEntry]> { self.check_valid()?; let entries = LoadBufferRef::from(&self.buffer[..]).primary_entries.into_slice(); let n = entries.iter().position(|v| v.is_null()).unwrap_or(entries.len()); Ok(&entries[..n]) } /// Returns the total number of partitions. pub fn num_partitions(&self) -> Result { Ok(self.entries()?.len()) } /// Gets the `idx`th partition. pub fn get_partition(&self, idx: usize) -> Result { let block_size = self.check_valid()?; let entry = *self.entries()?.get(idx).ok_or(Error::BadIndex(idx))?; Ok(Partition::new(entry, block_size)) } /// Returns the `Partition` for a partition. /// /// # Args /// /// * `part`: Name of the partition. pub fn find_partition(&self, part: &str) -> Result { let block_size = self.check_valid()?; for entry in self.entries()? { let mut name_conversion_buffer = [0u8; GPT_NAME_LEN_U16 * 2]; if entry.name_to_str(&mut name_conversion_buffer)? != part { continue; } return Ok(Partition::new(*entry, block_size)); } Err(Error::NotFound) } /// Checks whether the Gpt has been initialized and returns the block size. fn check_valid(&self) -> Result { Ok(LoadBufferRef::from(&self.buffer[..]).block_size.0.ok_or(Error::InvalidState)?.get()) } /// Helper function for loading and validating GPT header and entries. async fn load_and_validate_gpt( &mut self, disk: &mut Disk>, hdr_type: HeaderType, ) -> Result<()> { let blk_sz = disk.io().info().block_size; let load = LoadBufferRef::from(&mut self.buffer[..]); let (header_start, mut header, mut entries) = match hdr_type { HeaderType::Primary => (blk_sz, load.primary_header, load.primary_entries), HeaderType::Secondary => ( ((SafeNum::from(disk.io().info().num_blocks) - 1) * blk_sz).try_into()?, load.secondary_header, load.secondary_entries, ), }; // Loads the header disk.read(header_start, header.as_bytes_mut()).await?; // Checks header. check_header(disk.io(), &header, matches!(hdr_type, HeaderType::Primary))?; // Loads the entries. let entries_size = SafeNum::from(header.entries_count) * GPT_ENTRY_SIZE; let entries_offset = SafeNum::from(header.entries) * blk_sz; let out = entries.as_bytes_mut().get_mut(..entries_size.try_into().unwrap()).ok_or( Error::BufferTooSmall(Some( gpt_buffer_size(header.entries_count.try_into().unwrap()).unwrap(), )), )?; disk.read(entries_offset.try_into().unwrap(), out).await?; // Checks entries. check_entries(&header, entries.as_bytes()) } /// Loads and syncs GPT from a block device. /// /// * Returns Ok(sync_result) if disk IO is successful, where `sync_result` contains the GPT /// verification and restoration result, /// * Returns Err() if disk IO encounters error. pub(crate) async fn load_and_sync( &mut self, disk: &mut Disk>, ) -> Result { let blk_sz = disk.io().info().block_size; let nonzero_blk_sz = NonZeroU64::new(blk_sz).ok_or(Error::InvalidInput)?; let total_blocks: SafeNum = disk.io().info().num_blocks.into(); let primary_header_blk = 1; let primary_header_pos = blk_sz; let secondary_header_blk = total_blocks - 1; // Entries position for restoring. let primary_entries_blk = 2; let primary_entries_pos = SafeNum::from(primary_entries_blk) * blk_sz; let primary_res = self.load_and_validate_gpt(disk, HeaderType::Primary).await; let secondary_res = self.load_and_validate_gpt(disk, HeaderType::Secondary).await; let LoadBufferRef { mut block_size, mut primary_header, mut secondary_header, mut primary_entries, mut secondary_entries, } = LoadBufferRef::from(&mut self.buffer[..]); block_size.0 = None; let primary_entries = primary_entries.as_bytes_mut(); let secondary_entries = secondary_entries.as_bytes_mut(); let sync_res = match (primary_res, secondary_res) { (Err(primary), Err(secondary)) => GptSyncResult::NoValidGpt { primary, secondary }, (Ok(()), Ok(())) if is_consistent(&primary_header, &secondary_header) => { GptSyncResult::BothValid } (Err(e), Ok(())) => { // Restores to primary primary_header.as_bytes_mut().clone_from_slice(secondary_header.as_bytes()); primary_entries.clone_from_slice(&secondary_entries); primary_header.current = primary_header_blk; primary_header.backup = secondary_header_blk.try_into()?; primary_header.entries = primary_entries_blk; primary_header.update_crc(); disk.write(primary_header_pos, primary_header.as_bytes_mut()).await?; disk.write(primary_entries_pos.try_into()?, primary_entries).await?; GptSyncResult::PrimaryRestored(e) } (Ok(()), v) => { // Restores to secondary let pos = secondary_header_blk * blk_sz; let secondary_entries_pos = pos - GPT_MAX_NUM_ENTRIES_SIZE; let secondary_entries_blk = secondary_entries_pos / blk_sz; secondary_header.as_bytes_mut().clone_from_slice(primary_header.as_bytes()); secondary_entries.clone_from_slice(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(); disk.write(pos.try_into()?, secondary_header.as_bytes_mut()).await?; disk.write(secondary_entries_pos.try_into()?, secondary_entries).await?; GptSyncResult::SecondaryRestored(match v { Err(e) => e, _ => Error::GptError(GptError::DifferentFromPrimary), }) } }; block_size.0 = Some(nonzero_blk_sz); Ok(sync_res) } } /// Checks whether primary and secondary header fn is_consistent(primary: &GptHeader, secondary: &GptHeader) -> bool { let mut expected_secondary = *primary; expected_secondary.crc32 = secondary.crc32; expected_secondary.current = secondary.current; expected_secondary.backup = 1; expected_secondary.entries = secondary.entries; &expected_secondary == secondary } /// A [Gpt] that owns a `GptLoadBufferN` and can load up to N partition entries. /// /// Note: The size of this type increases with N and can be expensive to store on stack. It is /// typically intended for resource abundant environment such as test. pub type GptN = Gpt>; /// Creates an instance of GptN. pub fn new_gpt_n() -> GptN { Gpt::new(GptLoadBufferN::::new_zeroed()).unwrap() } /// A [Gpt] that owns a `GptLoadBufferN<128>` and can load the maximum 128 partition entries. /// /// Note: The size of this type is approximately 34K and can be expensive to store on stack. It /// is typically intended for resource abundant environment such as test. pub type GptMax = GptN; /// Creates an instance of GptMax. pub fn new_gpt_max() -> GptMax { new_gpt_n::() } /// Updates GPT on a block device. /// /// # Args /// /// * `io`: An implementation of [BlockIo] /// * `scratch`: Scratch buffer for unaligned read write. /// * `mbr_primary`: A buffer containing the MBR block, primary GPT header and entries. /// * `resize`: If set to true, the method updates the last partition to cover the rest of the /// storage. /// * `gpt`: The output [Gpt] to update. pub(crate) async fn update_gpt( disk: &mut Disk>, mbr_primary: &mut [u8], resize: bool, gpt: &mut Gpt>, ) -> Result<()> { let blk_sz: usize = disk.io().info().block_size.try_into()?; let (header, remain) = mbr_primary .get_mut(blk_sz..) .map(|v| v.split_at_mut_checked(blk_sz)) .flatten() .ok_or(Error::BufferTooSmall(Some(blk_sz * 2)))?; let header = Ref::<_, GptHeader>::new_from_prefix(&mut header[..]).unwrap().0.into_mut(); // Adjusts last usable block according to this device in case the GPT was generated for a // different disk size. If this results in some partition being out of range, it will be // caught during `check_header()`. let entries_blk = SafeNum::from(GPT_MAX_NUM_ENTRIES_SIZE) / blk_sz; // Reserves only secondary GPT header and entries. let num_blks = SafeNum::from(disk.io().info().num_blocks); header.last = (num_blks - entries_blk - 2).try_into().unwrap(); header.backup = (num_blks - 1).try_into().unwrap(); header.update_crc(); check_header(disk.io(), &header, true)?; // Computes entries offset in bytes relative to `remain` let entries_off: usize = ((SafeNum::from(header.entries) - 2) * blk_sz).try_into().unwrap(); let entries_size: usize = (SafeNum::from(header.entries_count) * header.entries_size).try_into().unwrap(); let entries = remain .get_mut(entries_off..) .map(|v| v.get_mut(..entries_size)) .flatten() .ok_or(Error::BufferTooSmall(Some(2 * blk_sz + entries_off + entries_size)))?; check_entries(&header, entries)?; if resize { // Updates the last entry to cover the rest of the storage. let gpt_entries = Ref::<_, [GptEntry]>::new_slice(&mut entries[..]).unwrap().into_mut_slice(); gpt_entries.iter_mut().filter(|e| !e.is_null()).last().map(|v| v.last = header.last); header.update_entries_crc(entries); // Re-verifies everything. check_header(disk.io(), &header, true).unwrap(); check_entries(&header, entries).unwrap(); } disk.write(0, mbr_primary).await?; disk.sync_gpt(gpt).await?.res() } /// Erases GPT if there is one on the device. pub(crate) async fn erase_gpt( disk: &mut Disk>, gpt: &mut Gpt>, ) -> Result<()> { match disk.sync_gpt(gpt).await?.res() { Err(_) => Ok(()), // No valid GPT. Nothing to erase. _ => { let blk_sz = disk.block_info().block_size; let mut load = LoadBufferRef::from(&mut gpt.buffer[..]); let entries_size = SafeNum::from(load.primary_header.entries_count) * GPT_ENTRY_SIZE; let scratch = load.primary_entries.as_bytes_mut(); // Invalidate GPT first. load.block_size.0 = None; // Erases primary header/entries. let header = load.primary_header.current; let entries = load.primary_header.entries; disk.fill(header * blk_sz, blk_sz, 0, scratch).await?; disk.fill(entries * blk_sz, entries_size.try_into().unwrap(), 0, scratch).await?; // Erases secondary header/entries. let header = load.secondary_header.current; let entries = load.secondary_header.entries; disk.fill(header * blk_sz, blk_sz, 0, scratch).await?; disk.fill(entries * blk_sz, entries_size.try_into().unwrap(), 0, scratch).await?; Ok(()) } } } /// Computes the minimum blocks needed for creating a GPT. fn min_required_blocks(block_size: u64) -> Result { // MBR + primary/secondary GPT header block + primary/secondary entries blocks. Ok(1 + (1 + gpt_entries_blk(block_size)?) * 2) } /// `GptBuilder` provides API for modifying/creating GPT partition table on a disk. pub struct GptBuilder { disk: D, gpt: G, } impl Debug for GptBuilder { fn fmt(&self, f: &mut Formatter<'_>) -> core::result::Result<(), core::fmt::Error> { write!(f, "GptBuilder {{ disk: {:?}, gpt: {:?} }}", self.disk, self.gpt) } } // Generic parameters: // // * T: The type that implement BlockIo. // * S: The type for the scratch buffer in `Self::disk`. // * B: The type for the GPT buffer in `Self::gpt`. // * D: The type for `Self::disk` which can dereference to a Disk. // * G: The type for `Self::gpt` which can dereference to a Gpt. impl<'a, T, S, B, D, G> GptBuilder where T: BlockIo, S: DerefMut, B: DerefMut, D: DerefMut>, G: DerefMut>, { /// Creates a new instance. /// /// The method always re-syncs the GPT. If `disk` does not contain a valid GPT, a new GPT is /// started from scratch. /// /// The partition entries will always be sorted when writing back to disk by `Self::persist()`. /// /// # Returns /// /// * Returns Ok((Self, true)) if an instance is created and the disk has a valid GPT. /// * Returns Ok((Self, false)) if an instance is created but disk does not have a valid GPT. /// * Returns Err() otherwise. pub fn new(mut disk: D, mut gpt: G) -> Result<(Self, bool)> { if disk.block_info().num_blocks < min_required_blocks(disk.block_info().block_size)? { return Err(Error::GptError(GptError::DiskTooSmall)); } let has_valid_gpt = block_on(disk.sync_gpt(&mut gpt))?.res().is_ok(); // Uses the buffer for secondary GPT header/entries as construction buffer, as it is not // used by Gpt once loaded and synced. let (mut header, mut entries) = LoadBufferRef::from(&mut gpt.buffer[..]).secondary(); if !has_valid_gpt { header.as_bytes_mut().fill(0); entries.as_bytes_mut().fill(0); let entries_blk = gpt_entries_blk(disk.block_info().block_size).unwrap(); // Initializes a secondary header. let num_blks = SafeNum::from(disk.block_info().num_blocks); header.magic = GPT_MAGIC; header.current = (num_blks - 1).try_into().unwrap(); header.backup = 1; header.size = size_of::().try_into().unwrap(); header.first = 1 + 1 + entries_blk; // MBR + GPT header blocks + entries block header.last = (num_blks - 1 - entries_blk - 1).try_into().unwrap(); header.entries = (num_blks - 1 - entries_blk).try_into().unwrap(); header.entries_count = 0; header.entries_size = size_of::().try_into().unwrap(); } // Normalizes `entries_count` to actual valid entries. Some GPT disk fixes `entry_count` to // 128. header.entries_count = entries.iter().position(|v| v.is_null()).unwrap_or(entries.len()).try_into().unwrap(); entries.sort_unstable_by_key(|v| match v.is_null() { true => u64::MAX, _ => v.first, }); Ok((Self { disk, gpt }, has_valid_gpt)) } /// Removes a partition. /// /// # Returns /// /// * Returns Ok(true) if found and removed. /// * Returns Ok(false) if not found. /// * Returns Err() otherwise. pub fn remove(&mut self, part: &str) -> Result { let (mut header, mut entries) = LoadBufferRef::from(&mut self.gpt.buffer[..]).secondary(); let entries = &mut entries[..header.entries_count.try_into().unwrap()]; match entries.iter().position(|v| v.match_name(part).unwrap_or(false)) { Some(n) => { // Shift the elements behind forward. entries[n..].rotate_left(1); // Zeroizes the last element. entries.last_mut().unwrap().as_bytes_mut().fill(0); header.entries_count -= 1; Ok(true) } _ => Ok(false), } } /// Inserts a new partition before a partition. /// /// # Args /// /// * `idx`: Index of the partition to insert before. If index is out of range of valid entries, /// the partition will be inserted at the last. /// * `name`: Name of the partition. /// * `part_type`: Type GUID. /// * `unique_guid`: Unique GUID. /// * `flags`: Partition flag. /// * `size`: If Some(_), specifies the size in number of bytes for the partition. The method /// will round it up to multiple of disk block size and check that there is enough space for /// the partition. If None, the method will insert the partition and consumes all the /// available space in between. fn insert_before( &mut self, idx: usize, name: &str, part_type: [u8; GPT_GUID_LEN], unique_guid: [u8; GPT_GUID_LEN], flags: u64, size: Option, ) -> Result<()> { let (mut header, mut entries) = LoadBufferRef::from(&mut self.gpt.buffer[..]).secondary(); // Gets position to the first NULL entry. let n = entries.iter().position(|v| v.is_null()).ok_or(Error::OutOfResources)?; let entries = &mut entries[..n + 1]; // Caps `idx` to no more than the first NULL entry. let idx = min(n, idx); // Comptues the ending block index (non-inclusive) of the previous partition entry. // Entries are guaranteed sorted in `Self::new()`. let prev_end = match idx { 0 => header.first, _ => entries[idx - 1].last + 1, }; // Comptues the starting block index (inclusive) of the next partition entry. let next_start = match idx == n { true => header.last + 1, _ => entries[idx].first, }; // Computes the size in number of blocks let blk_sz = self.disk.block_info().block_size; let blocks: u64 = match size { Some(v) => (SafeNum::from(v).round_up(blk_sz) / blk_sz).try_into()?, _ => next_start - prev_end, // If not given, uses up all the gap space }; // Checks if there is enough space. if next_start - prev_end < blocks { return Err(Error::OutOfResources); } // Inserts the new entry. entries[idx..].rotate_right(1); let entry = &mut entries[idx]; assert!(entry.is_null()); entry.part_type = part_type; entry.guid = unique_guid; entry.flags = flags; entry.first = prev_end; entry.last = prev_end + blocks - 1; for (idx, ele) in name.encode_utf16().enumerate() { match idx < GPT_NAME_LEN_U16 { true => entry.name[idx] = ele, _ => break, } } header.entries_count += 1; Ok(()) } /// Adds a partition. /// /// # Args /// /// * `name`: Name of the partition. /// * `part_type`: Type GUID. /// * `unique_guid`: Unique GUID. /// * `flags`: Partition flag. /// * `size`: If Some(_), specifies the size in number of bytes for the partition. The method /// will round it up to multiple of disk block size and search for the first large enough /// space in the unused spae for putting the partition. If None, the method will add the /// partition at the last and have it consume all remaining usable disk space. pub fn add( &mut self, name: &str, part_type: [u8; GPT_GUID_LEN], unique_guid: [u8; GPT_GUID_LEN], flags: u64, size: Option, ) -> Result<()> { let (header, _) = LoadBufferRef::from(&mut self.gpt.buffer[..]).secondary(); let entry_count = usize::try_from(header.entries_count).unwrap(); let search_start = size.is_some().then_some(0).unwrap_or(entry_count); for i in search_start..entry_count + 1 { if self.insert_before(i, name, part_type, unique_guid, flags, size).is_ok() { return Ok(()); } } Err(Error::OutOfResources) } /// Persists the constructed GPT table to the disk and syncs. The builder is consumed. pub async fn persist(mut self) -> Result<()> { let (mut header, mut entries) = LoadBufferRef::from(&mut self.gpt.buffer[..]).secondary(); header.update_entries_crc(entries.as_bytes()); // Check validity. Should not fail if implementation is correct. check_header(self.disk.io(), &header, false).unwrap(); check_entries(&header, entries.as_bytes()).unwrap(); let blk_sz = self.disk.block_info().block_size; // Writes to secondary header/ entries self.disk.write(header.current * blk_sz, header.as_bytes_mut()).await?; self.disk.write(header.entries * blk_sz, entries.as_bytes_mut()).await?; // Clears primary header magic self.disk.write(blk_sz, &mut 0u64.to_be_bytes()).await?; // Re-syncs GPT self.disk.sync_gpt(&mut self.gpt).await?.res() } } /// Helper for calculcating the Crc32. fn crc32(data: &[u8]) -> u32 { let mut hasher = Hasher::new(); hasher.update(data); hasher.finalize() } #[cfg(test)] pub(crate) mod test { use super::*; use crate::test::TestDisk; use gbl_async::block_on; /// A helper for creating a [TestDisk] from given data. fn test_disk(data: impl AsRef<[u8]>) -> TestDisk { // All tests cases use pre-generated GPT disk of 512 block size. TestDisk::new_ram_alloc(512, 512, data.as_ref().to_vec()).unwrap() } /// A helper for creating a [TestDisk] from given data and a [Gpt] for 128 entries. fn test_disk_and_gpt(data: impl AsRef<[u8]>) -> (TestDisk, GptMax) { (test_disk(data), new_gpt_max()) } #[test] fn test_load_and_sync() { let (mut dev, mut gpt) = test_disk_and_gpt(include_bytes!("../test/gpt_test_1.bin")); block_on(dev.sync_gpt(&mut gpt)).unwrap(); assert_eq!(gpt.partition_iter().unwrap().count(), 2); gpt.find_partition("boot_a").unwrap(); gpt.find_partition("boot_b").unwrap(); assert!(gpt.find_partition("boot_c").is_err()); // Creating a new [Gpt] using the same buffer should reset the valid state. let gpt = Gpt::new(gpt.buffer).unwrap(); assert!(gpt.partition_iter().is_err()); assert!(gpt.find_partition("boot_a").is_err()); assert!(gpt.find_partition("boot_b").is_err()); } #[test] fn test_load_with_unaligned_buffer() { #[repr(align(8))] struct AlignedBuffer([u8; 34 * 1024]); let mut buffer = AlignedBuffer([0u8; 34 * 1024]); let buffer = &mut buffer.0[1..]; assert_ne!(buffer.as_ptr() as usize % 2, 0); let mut disk = test_disk(include_bytes!("../test/gpt_test_1.bin")); let mut gpt = Gpt::new(buffer).unwrap(); block_on(disk.sync_gpt(&mut gpt)).unwrap(); } #[test] fn test_gpt_buffer_too_small() { assert!(Gpt::new(vec![0u8; size_of::>() - 1]).is_err()); } #[test] fn test_gpt_buffer_not_enough_for_all_entries() { let mut dev = test_disk(include_bytes!("../test/gpt_test_1.bin")); let mut gpt = new_gpt_n::<127>(); assert_eq!(gpt.max_entries(), 127); // Actual entries_count is 128 in the GPT. assert!(block_on(dev.sync_gpt(&mut gpt)).unwrap().res().is_err()); } #[test] fn test_good_gpt_no_repair_write() { let (mut dev, mut gpt) = test_disk_and_gpt(include_bytes!("../test/gpt_test_1.bin")); assert_eq!(block_on(dev.sync_gpt(&mut gpt)).unwrap(), GptSyncResult::BothValid); } /// A helper for testing restoration of invalid primary/secondary header modified by caller. fn test_gpt_sync_restore<'a>( modify_primary: impl FnOnce(&mut GptHeader, Ref<&mut [u8], [GptEntry]>), modify_secondary: impl FnOnce(&mut GptHeader, Ref<&mut [u8], [GptEntry]>), expect_primary_err: Error, expect_secondary_err: Error, ) { let disk_orig = include_bytes!("../test/gpt_test_1.bin"); // Restores from secondary to primary. let mut disk = disk_orig.to_vec(); let (header, entries) = (&mut disk[512..]).split_at_mut(512); let mut header = GptHeader::from_bytes_mut(header); modify_primary(&mut header, Ref::<_, [GptEntry]>::new_slice(entries).unwrap()); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); assert_ne!(dev.io().storage(), disk_orig); let sync_res = block_on(dev.sync_gpt(&mut gpt)).unwrap(); assert_eq!(sync_res, GptSyncResult::PrimaryRestored(expect_primary_err)); assert_eq!(dev.io().storage(), disk_orig); // Restores from primary to secondary. let mut disk = disk_orig.to_vec(); let (entries, header) = (&mut disk[512..]).split_last_chunk_mut::<512>().unwrap(); let (_, entries) = entries.split_last_chunk_mut::<{ 512 * 32 }>().unwrap(); let mut header = GptHeader::from_bytes_mut(&mut header[..]); modify_secondary(&mut header, Ref::<_, [GptEntry]>::new_slice(&mut entries[..]).unwrap()); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); assert_ne!(dev.io().storage(), disk_orig); let sync_res = block_on(dev.sync_gpt(&mut gpt)).unwrap(); assert_eq!(sync_res, GptSyncResult::SecondaryRestored(expect_secondary_err)); assert_eq!(dev.io().storage(), disk_orig); } #[test] fn test_sync_gpt_incorrect_magic() { fn modify(hdr: &mut GptHeader, _: Ref<&mut [u8], [GptEntry]>) { hdr.magic = 0x123456; hdr.update_crc(); } let err = Error::GptError(GptError::IncorrectMagic(0x123456)); test_gpt_sync_restore(modify, modify, err, err); } #[test] fn test_sync_gpt_incorrect_crc() { fn modify(hdr: &mut GptHeader, _: Ref<&mut [u8], [GptEntry]>) { hdr.crc32 = !hdr.crc32; } let err = Error::GptError(GptError::IncorrectHeaderCrc); test_gpt_sync_restore(modify, modify, err, err); } #[test] fn test_sync_gpt_unexpected_header_size() { fn modify(hdr: &mut GptHeader, _: Ref<&mut [u8], [GptEntry]>) { hdr.size += 1; hdr.update_crc(); } let err = Error::GptError(GptError::UnexpectedHeaderSize { actual: 93, expect: 92 }); test_gpt_sync_restore(modify, modify, err, err); } #[test] fn test_sync_gpt_unexpected_entry_size() { fn modify(hdr: &mut GptHeader, _: Ref<&mut [u8], [GptEntry]>) { hdr.entries_size += 1; hdr.update_crc(); } let err = Error::GptError(GptError::UnexpectedEntrySize { actual: 129, expect: 128 }); test_gpt_sync_restore(modify, modify, err, err); } #[test] fn test_sync_gpt_first_usable_gt_last() { fn modify(hdr: &mut GptHeader, _: Ref<&mut [u8], [GptEntry]>) { hdr.first = hdr.last; hdr.last = hdr.first - 2; hdr.update_crc(); } let err = Error::GptError(GptError::InvalidFirstLastUsableBlock { first: 94, last: 92, range: (34, 94), }); test_gpt_sync_restore(modify, modify, err, err); } #[test] fn test_sync_gpt_first_usable_out_of_range() { fn modify(hdr: &mut GptHeader, _: Ref<&mut [u8], [GptEntry]>) { hdr.first = 33; hdr.update_crc(); } let err = Error::GptError(GptError::InvalidFirstLastUsableBlock { first: 33, last: 94, range: (34, 94), }); test_gpt_sync_restore(modify, modify, err, err); } #[test] fn test_sync_gpt_last_usable_out_of_range() { fn modify(hdr: &mut GptHeader, _: Ref<&mut [u8], [GptEntry]>) { hdr.last += 1; hdr.update_crc(); } let err = Error::GptError(GptError::InvalidFirstLastUsableBlock { first: 34, last: 95, range: (34, 94), }); test_gpt_sync_restore(modify, modify, err, err); } #[test] fn test_sync_gpt_primary_entries_out_of_range() { test_gpt_sync_restore( |hdr, _| { hdr.entries = 1; hdr.update_crc(); }, |hdr, _| { hdr.entries = hdr.last; hdr.update_crc(); }, Error::GptError(GptError::InvalidPrimaryEntriesStart { value: 1, expect_range: (2, 2), }), Error::GptError(GptError::InvalidSecondaryEntriesStart { value: 94, expect_range: (95, 95), }), ); } #[test] fn test_sync_gpt_incorrect_entry_crc() { fn modify(hdr: &mut GptHeader, _: Ref<&mut [u8], [GptEntry]>) { hdr.entries_crc = !hdr.entries_crc; hdr.update_crc(); } let err = Error::GptError(GptError::IncorrectEntriesCrc); test_gpt_sync_restore(modify, modify, err, err); } #[test] fn test_sync_gpt_partition_range_overflow() { fn modify(hdr: &mut GptHeader, mut entries: Ref<&mut [u8], [GptEntry]>) { entries[1].last = hdr.last + 1; hdr.update_entries_crc(entries.as_bytes()); } let err = Error::GptError(GptError::InvalidPartitionRange { idx: 2, part_range: (50, 95), usable_range: (34, 94), }); test_gpt_sync_restore(modify, modify, err, err); } #[test] fn test_sync_gpt_invalid_partition_range() { fn modify(hdr: &mut GptHeader, mut entries: Ref<&mut [u8], [GptEntry]>) { entries[1].first = entries[1].last; entries[1].last = entries[1].first - 2; hdr.update_entries_crc(entries.as_bytes()); } let err = Error::GptError(GptError::InvalidPartitionRange { idx: 2, part_range: (73, 71), usable_range: (34, 94), }); test_gpt_sync_restore(modify, modify, err, err); } #[test] fn test_sync_gpt_partition_overlap() { fn modify(hdr: &mut GptHeader, mut entries: Ref<&mut [u8], [GptEntry]>) { entries[0].last = entries[1].first; entries.swap(0, 1); hdr.update_entries_crc(entries.as_bytes()); } let err = Error::GptError(GptError::PartitionRangeOverlap { prev: (2, 34, 50), next: (1, 50, 73), }); test_gpt_sync_restore(modify, modify, err, err); } #[test] fn test_sync_gpt_zero_partition_type_guid() { fn modify(hdr: &mut GptHeader, mut entries: Ref<&mut [u8], [GptEntry]>) { entries[1].part_type = [0u8; GPT_GUID_LEN]; hdr.update_entries_crc(entries.as_bytes()); } let err = Error::GptError(GptError::ZeroPartitionTypeGUID { idx: 2 }); test_gpt_sync_restore(modify, modify, err, err); } #[test] fn test_sync_gpt_zero_partition_unique_guid() { fn modify(hdr: &mut GptHeader, mut entries: Ref<&mut [u8], [GptEntry]>) { entries[1].guid = [0u8; GPT_GUID_LEN]; hdr.update_entries_crc(entries.as_bytes()); } let err = Error::GptError(GptError::ZeroPartitionUniqueGUID { idx: 2 }); test_gpt_sync_restore(modify, modify, err, err); } #[test] fn test_load_gpt_disk_primary_override_secondary() { let mut disk = include_bytes!("../test/gpt_test_1.bin").to_vec(); // Modifies secondary header. let secondary_hdr = GptHeader::from_bytes_mut(disk.last_chunk_mut::<512>().unwrap()); secondary_hdr.revision = !secondary_hdr.revision; secondary_hdr.update_crc(); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); assert_eq!( block_on(dev.sync_gpt(&mut gpt)).unwrap(), GptSyncResult::SecondaryRestored(Error::GptError(GptError::DifferentFromPrimary)), ); } #[test] fn test_load_gpt_disk_too_small() { let disk_orig = include_bytes!("../test/gpt_test_1.bin"); let mut disk = disk_orig.to_vec(); // Resizes so that it's not enough to hold a full 128 maximum entries. // MBR + (header + entries) * 2 - 1 disk.resize((1 + (32 + 1) * 2 - 1) * 512, 0); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); let sync_res = block_on(dev.sync_gpt(&mut gpt)).unwrap(); let err = Error::GptError(GptError::DiskTooSmall); assert_eq!(sync_res, GptSyncResult::NoValidGpt { primary: err, secondary: err }); } #[test] fn test_uninitialized_gpt() { let disk = include_bytes!("../test/gpt_test_1.bin"); // Load a good GPT first. let (mut dev, mut gpt) = test_disk_and_gpt(&disk); assert_eq!(block_on(dev.sync_gpt(&mut gpt)).unwrap(), GptSyncResult::BothValid); gpt.find_partition("boot_a").unwrap(); // Corrupt GPT. block_on(dev.write(0, &mut vec![0u8; disk.len()])).unwrap(); assert!(block_on(dev.sync_gpt(&mut gpt)).unwrap().res().is_err()); assert!(gpt.find_partition("").is_err()); } #[test] fn test_update_gpt() { let disk_orig = include_bytes!("../test/gpt_test_1.bin"); let mut disk = disk_orig.to_vec(); // Erases all GPT headers. disk[512..][..512].fill(0); disk.last_chunk_mut::<512>().unwrap().fill(0); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); assert_ne!(dev.io().storage(), disk_orig); let mut mbr_primary = disk_orig[..34 * 512].to_vec(); block_on(dev.update_gpt(&mut mbr_primary, false, &mut gpt)).unwrap(); assert_eq!(dev.io().storage(), disk_orig); } #[test] fn test_update_gpt_has_existing_valid_secondary() { let disk_orig = include_bytes!("../test/gpt_test_1.bin"); let mut disk = disk_orig.to_vec(); // Erases all GPT headers. disk[512..][..512].fill(0); // Leaves a valid but different secondary GPT. let secondary_hdr = GptHeader::from_bytes_mut(disk.last_chunk_mut::<512>().unwrap()); secondary_hdr.revision = !secondary_hdr.revision; secondary_hdr.update_crc(); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); assert_ne!(dev.io().storage(), disk_orig); let mut mbr_primary = disk_orig[..34 * 512].to_vec(); block_on(dev.update_gpt(&mut mbr_primary, false, &mut gpt)).unwrap(); assert_eq!(dev.io().storage(), disk_orig); } #[test] fn test_update_gpt_last_usable_adjusted() { let disk_orig = include_bytes!("../test/gpt_test_1.bin"); let mut disk = disk_orig.to_vec(); // Erases all GPT headers. disk[512..][..512].fill(0); disk.last_chunk_mut::<512>().unwrap().fill(0); // Doubles the disk size. disk.resize(disk_orig.len() * 2, 0); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); assert_ne!(dev.io().storage, disk_orig); let mut mbr_primary = disk_orig[..34 * 512].to_vec(); block_on(dev.update_gpt(&mut mbr_primary, true, &mut gpt)).unwrap(); let expected_last = (disk.len() - GPT_MAX_NUM_ENTRIES_SIZE - 512) / 512 - 1; let (primary, secondary) = dev.io().storage().split_last_chunk_mut::<512>().unwrap(); let primary_hdr = GptHeader::from_bytes_mut(&mut primary[512..]); let secondary_hdr = GptHeader::from_bytes_mut(secondary); // Header's last usable block is updated. assert_eq!({ primary_hdr.last }, expected_last.try_into().unwrap()); assert_eq!({ primary_hdr.backup }, (disk.len() / 512 - 1).try_into().unwrap()); assert_eq!({ secondary_hdr.last }, expected_last.try_into().unwrap()); } #[test] fn test_update_gpt_resize() { let disk_orig = include_bytes!("../test/gpt_test_1.bin"); let mut disk = disk_orig.to_vec(); // Erases all GPT headers. disk[512..][..512].fill(0); disk.last_chunk_mut::<512>().unwrap().fill(0); // Doubles the disk size. disk.resize(disk_orig.len() * 2, 0); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); assert_ne!(dev.io().storage, disk_orig); let mut mbr_primary = disk_orig[..34 * 512].to_vec(); block_on(dev.update_gpt(&mut mbr_primary, true, &mut gpt)).unwrap(); // Last entry is extended. let expected_last = (disk.len() - GPT_MAX_NUM_ENTRIES_SIZE - 512) / 512 - 1; assert_eq!({ gpt.entries().unwrap()[1].last }, expected_last.try_into().unwrap()); } #[test] fn test_update_gpt_new_partition_out_of_range() { // `gpt_test_1.bin` has a 8k "boot_a" and a 12k "boot_b". Thus partitions space is 40 // blocks (512 bytes block size) and in total the GPT disk needs (40 + 1 + (33) * 2) = 107 // blocks. let (mut dev, mut gpt) = test_disk_and_gpt(&vec![0u8; 106 * 512]); let mut mbr_primary = include_bytes!("../test/gpt_test_1.bin")[..34 * 512].to_vec(); assert!(block_on(dev.update_gpt(&mut mbr_primary, true, &mut gpt)).is_err()); } #[test] fn test_update_gpt_buffer_truncated() { let mut disk = include_bytes!("../test/gpt_test_1.bin").to_vec(); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); // Less than 1 MBR block. assert_eq!( block_on(dev.update_gpt(&mut disk[..511], false, &mut gpt)), Err(Error::BufferTooSmall(Some(1024))) ); // Less than MBR + GPT header. assert_eq!( block_on(dev.update_gpt(&mut disk[..1023], false, &mut gpt)), Err(Error::BufferTooSmall(Some(1024))) ); // Less than MBR + GPT header + entries. assert_eq!( block_on(dev.update_gpt(&mut disk[..34 * 512 - 1], false, &mut gpt)), Err(Error::BufferTooSmall(Some(34 * 512))) ); } #[test] fn test_update_gpt_check_header_fail() { let disk = include_bytes!("../test/gpt_test_1.bin"); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); let mut mbr_primary = disk[..34 * 512].to_vec(); // Corrupts the first byte of the GPT header. mbr_primary[512] = !mbr_primary[512]; assert_eq!( block_on(dev.update_gpt(&mut mbr_primary, false, &mut gpt)), Err(Error::GptError(GptError::IncorrectMagic(0x54524150204946BA))) ); } #[test] fn test_update_gpt_check_entries_fail() { let disk = include_bytes!("../test/gpt_test_1.bin"); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); let mut mbr_primary = disk[..34 * 512].to_vec(); // Corrupts the first byte of the entries. mbr_primary[1024] = !mbr_primary[1024]; assert_eq!( block_on(dev.update_gpt(&mut mbr_primary, false, &mut gpt)), Err(Error::GptError(GptError::IncorrectEntriesCrc)) ); } #[test] fn test_erase_gpt_no_gpt() { let (mut dev, mut gpt) = test_disk_and_gpt(&[0u8; 1024 * 1024]); block_on(dev.erase_gpt(&mut gpt)).unwrap(); } #[test] fn test_erase_gpt() { let (mut dev, mut gpt) = test_disk_and_gpt(include_bytes!("../test/gpt_test_1.bin")); block_on(dev.erase_gpt(&mut gpt)).unwrap(); const GPT_SECTOR: usize = 33 * 512; assert_eq!(dev.io().storage[512..][..GPT_SECTOR], vec![0u8; GPT_SECTOR]); assert_eq!(*dev.io().storage.last_chunk::().unwrap(), *vec![0u8; GPT_SECTOR]); assert!(matches!( block_on(dev.sync_gpt(&mut gpt)).unwrap(), GptSyncResult::NoValidGpt { .. } )); } #[test] fn test_zero_partition_size() { let disk = include_bytes!("../test/gpt_test_1.bin").to_vec(); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); let (mut builder, _) = GptBuilder::new(&mut dev, &mut gpt).unwrap(); assert_eq!(builder.remove("boot_a"), Ok(true)); assert_eq!(builder.remove("boot_b"), Ok(true)); builder.add("boot_b", [1u8; GPT_GUID_LEN], [1u8; GPT_GUID_LEN], 0, Some(0)).unwrap(); block_on(builder.persist()).unwrap(); assert_eq!(gpt.partition_iter().unwrap().next().unwrap().size().unwrap(), 0); } #[test] fn test_sync_gpt_non_sorted_entries() { let mut disk = include_bytes!("../test/gpt_test_1.bin").to_vec(); let (header, entries) = disk[512..].split_at_mut(512); let header = GptHeader::from_bytes_mut(header); let mut entries = Ref::<_, [GptEntry]>::new_slice(entries).unwrap(); // Makes partition non-sorted. entries.swap(0, 1); header.update_entries_crc(entries.as_bytes()); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); block_on(dev.sync_gpt(&mut gpt)).unwrap().res().unwrap(); } #[test] fn test_gpt_builder_initialize_gpt_if_no_valid_gpt() { let (mut dev, mut gpt) = test_disk_and_gpt(vec![0u8; 1024 * 1024]); let (builder, valid) = GptBuilder::new(&mut dev, &mut gpt).unwrap(); assert!(!valid); block_on(builder.persist()).unwrap(); // A new GPT is created. block_on(dev.sync_gpt(&mut gpt)).unwrap().res().unwrap(); assert!(gpt.partition_iter().unwrap().next().is_none()); } #[test] fn test_gpt_builder_remove_partition() { let (mut dev, mut gpt) = test_disk_and_gpt(include_bytes!("../test/gpt_test_1.bin")); let (mut builder, valid) = GptBuilder::new(&mut dev, &mut gpt).unwrap(); assert!(valid); assert_eq!(builder.remove("boot_b"), Ok(true)); assert_eq!(builder.remove("non-existent"), Ok(false)); block_on(builder.persist()).unwrap(); block_on(dev.sync_gpt(&mut gpt)).unwrap().res().unwrap(); let part_iter = gpt.partition_iter().unwrap(); assert_eq!( part_iter.map(|v| v.name().unwrap().into()).collect::>(), ["boot_a"] ); } #[test] fn test_gpt_builder_add_partition_find_first() { let (mut dev, mut gpt) = test_disk_and_gpt(include_bytes!("../test/gpt_test_1.bin")); let (mut builder, _) = GptBuilder::new(&mut dev, &mut gpt).unwrap(); assert!(builder.remove("boot_a").unwrap()); // Adds at the beginning. builder.add("new_0", [1u8; GPT_GUID_LEN], [1u8; GPT_GUID_LEN], 0, Some(1024)).unwrap(); // Adds following "new_0" builder.add("new_1", [1u8; GPT_GUID_LEN], [1u8; GPT_GUID_LEN], 0, Some(1)).unwrap(); block_on(builder.persist()).unwrap(); block_on(dev.sync_gpt(&mut gpt)).unwrap().res().unwrap(); assert_eq!(gpt.find_partition("new_0").unwrap().absolute_range().unwrap(), (17408, 18432)); assert_eq!(gpt.find_partition("new_1").unwrap().absolute_range().unwrap(), (18432, 18944)); assert_eq!(gpt.find_partition("boot_b").unwrap().absolute_range().unwrap(), (25600, 37888)); } #[test] fn test_gpt_builder_non_sorted_add_partition() { let mut disk = include_bytes!("../test/gpt_test_1.bin").to_vec(); let (mut dev, mut gpt) = test_disk_and_gpt(&disk); let (header, entries) = disk[512..].split_at_mut(512); let header = GptHeader::from_bytes_mut(header); let mut entries = Ref::<_, [GptEntry]>::new_slice(entries).unwrap(); // Makes partition non-sorted. entries.swap(0, 1); header.update_entries_crc(entries.as_bytes()); let (mut builder, _) = GptBuilder::new(&mut dev, &mut gpt).unwrap(); // Adds following boot_b. builder.add("new", [1u8; GPT_GUID_LEN], [1u8; GPT_GUID_LEN], 0, Some(1024)).unwrap(); block_on(builder.persist()).unwrap(); assert_eq!(gpt.find_partition("boot_a").unwrap().absolute_range().unwrap(), (17408, 25600)); assert_eq!(gpt.find_partition("boot_b").unwrap().absolute_range().unwrap(), (25600, 37888)); assert_eq!(gpt.find_partition("new").unwrap().absolute_range().unwrap(), (37888, 38912)); } #[test] fn test_gpt_builder_add_partition_append() { let (mut dev, mut gpt) = test_disk_and_gpt(include_bytes!("../test/gpt_test_1.bin")); let (mut builder, _) = GptBuilder::new(&mut dev, &mut gpt).unwrap(); assert!(builder.remove("boot_b").unwrap()); // Adds following "boot_a". builder.add("new_0", [1u8; GPT_GUID_LEN], [1u8; GPT_GUID_LEN], 0, Some(1024)).unwrap(); // Consumes the rest of the space. builder.add("new_1", [1u8; GPT_GUID_LEN], [1u8; GPT_GUID_LEN], 0, None).unwrap(); block_on(builder.persist()).unwrap(); block_on(dev.sync_gpt(&mut gpt)).unwrap().res().unwrap(); assert_eq!(gpt.find_partition("boot_a").unwrap().absolute_range().unwrap(), (17408, 25600)); assert_eq!(gpt.find_partition("new_0").unwrap().absolute_range().unwrap(), (25600, 26624)); assert_eq!(gpt.find_partition("new_1").unwrap().absolute_range().unwrap(), (26624, 48640)); } #[test] fn test_gpt_builder_not_enough_resource() { // Create a Gpt that can only load 1 entry. let mut gpt = new_gpt_n::<1>(); let mut dev = test_disk(vec![0u8; 64 * 1024]); let (mut builder, _) = GptBuilder::new(&mut dev, &mut gpt).unwrap(); builder.add("new_0", [1u8; GPT_GUID_LEN], [1u8; GPT_GUID_LEN], 0, Some(1024)).unwrap(); assert!(builder.add("new_1", [1u8; GPT_GUID_LEN], [1u8; GPT_GUID_LEN], 0, None).is_err()); } }