disk/src/gpt.rs - platform/external/crosvm - Git at Google

 // Copyright 2021 The ChromiumOS Authors
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.

 //! Functions for writing GUID Partition Tables for use in a composite disk image.

 use std::convert::TryInto;
 use std::io;
 use std::io::Write;
 use std::num::TryFromIntError;

 use crc32fast::Hasher;
 use remain::sorted;
 use thiserror::Error as ThisError;
 use uuid::Uuid;

 /// The size in bytes of a disk sector (also called a block).
 pub const SECTOR_SIZE: u64 = 1 << 9;
 /// The size in bytes on an MBR partition entry.
 const MBR_PARTITION_ENTRY_SIZE: usize = 16;
 /// The size in bytes of a GPT header.
 pub const GPT_HEADER_SIZE: u32 = 92;
 /// The number of partition entries in the GPT, which is the maximum number of partitions which are
 /// supported.
 pub const GPT_NUM_PARTITIONS: u32 = 128;
 /// The size in bytes of a single GPT partition entry.
 pub const GPT_PARTITION_ENTRY_SIZE: u32 = 128;
 /// The size in bytes of everything before the first partition: i.e. the MBR, GPT header and GPT
 /// partition entries.
 pub const GPT_BEGINNING_SIZE: u64 = SECTOR_SIZE * 40;
 /// The size in bytes of everything after the last partition: i.e. the GPT partition entries and GPT
 /// footer.
 pub const GPT_END_SIZE: u64 = SECTOR_SIZE * 33;

 #[sorted]
 #[derive(ThisError, Debug)]
 pub enum Error {
     /// The disk size was invalid (too large).
     #[error("invalid disk size: {0}")]
     InvalidDiskSize(TryFromIntError),
     /// There was an error writing data to one of the image files.
     #[error("failed to write data: {0}")]
     WritingData(io::Error),
 }

 /// Write a protective MBR for a disk of the given total size (in bytes).
 ///
 /// This should be written at the start of the disk, before the GPT header. It is one `SECTOR_SIZE`
 /// long.
 pub fn write_protective_mbr(file: &mut impl Write, disk_size: u64) -> Result<(), Error> {
     // Bootstrap code
     file.write_all(&[0; 446]).map_err(Error::WritingData)?;

     // Partition status
     file.write_all(&[0x00]).map_err(Error::WritingData)?;
     // Begin CHS
     file.write_all(&[0; 3]).map_err(Error::WritingData)?;
     // Partition type
     file.write_all(&[0xEE]).map_err(Error::WritingData)?;
     // End CHS
     file.write_all(&[0; 3]).map_err(Error::WritingData)?;
     let first_lba: u32 = 1;
     file.write_all(&first_lba.to_le_bytes())
         .map_err(Error::WritingData)?;
     let number_of_sectors: u32 = (disk_size / SECTOR_SIZE)
         .try_into()
         .map_err(Error::InvalidDiskSize)?;
     file.write_all(&number_of_sectors.to_le_bytes())
         .map_err(Error::WritingData)?;

     // Three more empty partitions
     file.write_all(&[0; MBR_PARTITION_ENTRY_SIZE * 3])
         .map_err(Error::WritingData)?;

     // Boot signature
     file.write_all(&[0x55, 0xAA]).map_err(Error::WritingData)?;

     Ok(())
 }

 #[derive(Clone, Debug, Default, Eq, PartialEq)]
 struct GptHeader {
     signature: [u8; 8],
     revision: [u8; 4],
     header_size: u32,
     header_crc32: u32,
     current_lba: u64,
     backup_lba: u64,
     first_usable_lba: u64,
     last_usable_lba: u64,
     disk_guid: Uuid,
     partition_entries_lba: u64,
     num_partition_entries: u32,
     partition_entry_size: u32,
     partition_entries_crc32: u32,
 }

 impl GptHeader {
     fn write_bytes(&self, out: &mut impl Write) -> Result<(), Error> {
         out.write_all(&self.signature).map_err(Error::WritingData)?;
         out.write_all(&self.revision).map_err(Error::WritingData)?;
         out.write_all(&self.header_size.to_le_bytes())
             .map_err(Error::WritingData)?;
         out.write_all(&self.header_crc32.to_le_bytes())
             .map_err(Error::WritingData)?;
         // Reserved
         out.write_all(&[0; 4]).map_err(Error::WritingData)?;
         out.write_all(&self.current_lba.to_le_bytes())
             .map_err(Error::WritingData)?;
         out.write_all(&self.backup_lba.to_le_bytes())
             .map_err(Error::WritingData)?;
         out.write_all(&self.first_usable_lba.to_le_bytes())
             .map_err(Error::WritingData)?;
         out.write_all(&self.last_usable_lba.to_le_bytes())
             .map_err(Error::WritingData)?;

         // GUID is mixed-endian for some reason, so we can't just use `Uuid::as_bytes()`.
         write_guid(out, self.disk_guid).map_err(Error::WritingData)?;

         out.write_all(&self.partition_entries_lba.to_le_bytes())
             .map_err(Error::WritingData)?;
         out.write_all(&self.num_partition_entries.to_le_bytes())
             .map_err(Error::WritingData)?;
         out.write_all(&self.partition_entry_size.to_le_bytes())
             .map_err(Error::WritingData)?;
         out.write_all(&self.partition_entries_crc32.to_le_bytes())
             .map_err(Error::WritingData)?;
         Ok(())
     }
 }

 /// Write a GPT header for the disk.
 ///
 /// It may either be a primary header (which should go at LBA 1) or a secondary header (which should
 /// go at the end of the disk).
 pub fn write_gpt_header(
     out: &mut impl Write,
     disk_guid: Uuid,
     partition_entries_crc32: u32,
     secondary_table_offset: u64,
     secondary: bool,
 ) -> Result<(), Error> {
     let primary_header_lba = 1;
     let secondary_header_lba = (secondary_table_offset + GPT_END_SIZE) / SECTOR_SIZE - 1;
     let mut gpt_header = GptHeader {
         signature: *b"EFI PART",
         revision: [0, 0, 1, 0],
         header_size: GPT_HEADER_SIZE,
         current_lba: if secondary {
             secondary_header_lba
         } else {
             primary_header_lba
         },
         backup_lba: if secondary {
             primary_header_lba
         } else {
             secondary_header_lba
         },
         first_usable_lba: GPT_BEGINNING_SIZE / SECTOR_SIZE,
         last_usable_lba: secondary_table_offset / SECTOR_SIZE - 1,
         disk_guid,
         partition_entries_lba: 2,
         num_partition_entries: GPT_NUM_PARTITIONS,
         partition_entry_size: GPT_PARTITION_ENTRY_SIZE,
         partition_entries_crc32,
         header_crc32: 0,
     };

     // Write once to a temporary buffer to calculate the CRC.
     let mut header_without_crc = [0u8; GPT_HEADER_SIZE as usize];
     gpt_header.write_bytes(&mut &mut header_without_crc[..])?;
     let mut hasher = Hasher::new();
     hasher.update(&header_without_crc);
     gpt_header.header_crc32 = hasher.finalize();

     gpt_header.write_bytes(out)?;

     Ok(())
 }

 /// A GPT entry for a particular partition.
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct GptPartitionEntry {
     pub partition_type_guid: Uuid,
     pub unique_partition_guid: Uuid,
     pub first_lba: u64,
     pub last_lba: u64,
     pub attributes: u64,
     /// UTF-16LE
     pub partition_name: [u16; 36],
 }

 // This is implemented manually because `Default` isn't implemented in the standard library for
 // arrays of more than 32 elements. If that gets implemented (now than const generics are in) then
 // we can derive this instead.
 impl Default for GptPartitionEntry {
     fn default() -> Self {
         Self {
             partition_type_guid: Default::default(),
             unique_partition_guid: Default::default(),
             first_lba: 0,
             last_lba: 0,
             attributes: 0,
             partition_name: [0; 36],
         }
     }
 }

 impl GptPartitionEntry {
     /// Write out the partition table entry. It will take
     /// `GPT_PARTITION_ENTRY_SIZE` bytes.
     pub fn write_bytes(&self, out: &mut impl Write) -> Result<(), Error> {
         write_guid(out, self.partition_type_guid).map_err(Error::WritingData)?;
         write_guid(out, self.unique_partition_guid).map_err(Error::WritingData)?;
         out.write_all(&self.first_lba.to_le_bytes())
             .map_err(Error::WritingData)?;
         out.write_all(&self.last_lba.to_le_bytes())
             .map_err(Error::WritingData)?;
         out.write_all(&self.attributes.to_le_bytes())
             .map_err(Error::WritingData)?;
         for code_unit in &self.partition_name {
             out.write_all(&code_unit.to_le_bytes())
                 .map_err(Error::WritingData)?;
         }
         Ok(())
     }
 }

 /// Write a UUID in the mixed-endian format which GPT uses for GUIDs.
 fn write_guid(out: &mut impl Write, guid: Uuid) -> Result<(), io::Error> {
     let guid_fields = guid.as_fields();
     out.write_all(&guid_fields.0.to_le_bytes())?;
     out.write_all(&guid_fields.1.to_le_bytes())?;
     out.write_all(&guid_fields.2.to_le_bytes())?;
     out.write_all(guid_fields.3)?;

     Ok(())
 }

 #[cfg(test)]
 mod tests {
     use super::*;

     #[test]
     fn protective_mbr_size() {
         let mut buffer = vec![];
         write_protective_mbr(&mut buffer, 1000 * SECTOR_SIZE).unwrap();

         assert_eq!(buffer.len(), SECTOR_SIZE as usize);
     }

     #[test]
     fn header_size() {
         let mut buffer = vec![];
         write_gpt_header(
             &mut buffer,
             Uuid::from_u128(0x12345678_1234_5678_abcd_12345678abcd),
             42,
             1000 * SECTOR_SIZE,
             false,
         )
         .unwrap();

         assert_eq!(buffer.len(), GPT_HEADER_SIZE as usize);
     }

     #[test]
     fn partition_entry_size() {
         let mut buffer = vec![];
         GptPartitionEntry::default()
             .write_bytes(&mut buffer)
             .unwrap();

         assert_eq!(buffer.len(), GPT_PARTITION_ENTRY_SIZE as usize);
     }
 }
	// Copyright 2021 The ChromiumOS Authors
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	//! Functions for writing GUID Partition Tables for use in a composite disk image.

	use std::convert::TryInto;
	use std::io;
	use std::io::Write;
	use std::num::TryFromIntError;

	use crc32fast::Hasher;
	use remain::sorted;
	use thiserror::Error as ThisError;
	use uuid::Uuid;

	/// The size in bytes of a disk sector (also called a block).
	pub const SECTOR_SIZE: u64 = 1 << 9;
	/// The size in bytes on an MBR partition entry.
	const MBR_PARTITION_ENTRY_SIZE: usize = 16;
	/// The size in bytes of a GPT header.
	pub const GPT_HEADER_SIZE: u32 = 92;
	/// The number of partition entries in the GPT, which is the maximum number of partitions which are
	/// supported.
	pub const GPT_NUM_PARTITIONS: u32 = 128;
	/// The size in bytes of a single GPT partition entry.
	pub const GPT_PARTITION_ENTRY_SIZE: u32 = 128;
	/// The size in bytes of everything before the first partition: i.e. the MBR, GPT header and GPT
	/// partition entries.
	pub const GPT_BEGINNING_SIZE: u64 = SECTOR_SIZE * 40;
	/// The size in bytes of everything after the last partition: i.e. the GPT partition entries and GPT
	/// footer.
	pub const GPT_END_SIZE: u64 = SECTOR_SIZE * 33;

	#[sorted]
	#[derive(ThisError, Debug)]
	pub enum Error {
	/// The disk size was invalid (too large).
	#[error("invalid disk size: {0}")]
	InvalidDiskSize(TryFromIntError),
	/// There was an error writing data to one of the image files.
	#[error("failed to write data: {0}")]
	WritingData(io::Error),
	}

	/// Write a protective MBR for a disk of the given total size (in bytes).
	///
	/// This should be written at the start of the disk, before the GPT header. It is one `SECTOR_SIZE`
	/// long.
	pub fn write_protective_mbr(file: &mut impl Write, disk_size: u64) -> Result<(), Error> {
	// Bootstrap code
	file.write_all(&[0; 446]).map_err(Error::WritingData)?;

	// Partition status
	file.write_all(&[0x00]).map_err(Error::WritingData)?;
	// Begin CHS
	file.write_all(&[0; 3]).map_err(Error::WritingData)?;
	// Partition type
	file.write_all(&[0xEE]).map_err(Error::WritingData)?;
	// End CHS
	file.write_all(&[0; 3]).map_err(Error::WritingData)?;
	let first_lba: u32 = 1;
	file.write_all(&first_lba.to_le_bytes())
	.map_err(Error::WritingData)?;
	let number_of_sectors: u32 = (disk_size / SECTOR_SIZE)
	.try_into()
	.map_err(Error::InvalidDiskSize)?;
	file.write_all(&number_of_sectors.to_le_bytes())
	.map_err(Error::WritingData)?;

	// Three more empty partitions
	file.write_all(&[0; MBR_PARTITION_ENTRY_SIZE * 3])
	.map_err(Error::WritingData)?;

	// Boot signature
	file.write_all(&[0x55, 0xAA]).map_err(Error::WritingData)?;

	Ok(())
	}

	#[derive(Clone, Debug, Default, Eq, PartialEq)]
	struct GptHeader {
	signature: [u8; 8],
	revision: [u8; 4],
	header_size: u32,
	header_crc32: u32,
	current_lba: u64,
	backup_lba: u64,
	first_usable_lba: u64,
	last_usable_lba: u64,
	disk_guid: Uuid,
	partition_entries_lba: u64,
	num_partition_entries: u32,
	partition_entry_size: u32,
	partition_entries_crc32: u32,
	}

	impl GptHeader {
	fn write_bytes(&self, out: &mut impl Write) -> Result<(), Error> {
	out.write_all(&self.signature).map_err(Error::WritingData)?;
	out.write_all(&self.revision).map_err(Error::WritingData)?;
	out.write_all(&self.header_size.to_le_bytes())
	.map_err(Error::WritingData)?;
	out.write_all(&self.header_crc32.to_le_bytes())
	.map_err(Error::WritingData)?;
	// Reserved
	out.write_all(&[0; 4]).map_err(Error::WritingData)?;
	out.write_all(&self.current_lba.to_le_bytes())
	.map_err(Error::WritingData)?;
	out.write_all(&self.backup_lba.to_le_bytes())
	.map_err(Error::WritingData)?;
	out.write_all(&self.first_usable_lba.to_le_bytes())
	.map_err(Error::WritingData)?;
	out.write_all(&self.last_usable_lba.to_le_bytes())
	.map_err(Error::WritingData)?;

	// GUID is mixed-endian for some reason, so we can't just use `Uuid::as_bytes()`.
	write_guid(out, self.disk_guid).map_err(Error::WritingData)?;

	out.write_all(&self.partition_entries_lba.to_le_bytes())
	.map_err(Error::WritingData)?;
	out.write_all(&self.num_partition_entries.to_le_bytes())
	.map_err(Error::WritingData)?;
	out.write_all(&self.partition_entry_size.to_le_bytes())
	.map_err(Error::WritingData)?;
	out.write_all(&self.partition_entries_crc32.to_le_bytes())
	.map_err(Error::WritingData)?;
	Ok(())
	}
	}

	/// Write a GPT header for the disk.
	///
	/// It may either be a primary header (which should go at LBA 1) or a secondary header (which should
	/// go at the end of the disk).
	pub fn write_gpt_header(
	out: &mut impl Write,
	disk_guid: Uuid,
	partition_entries_crc32: u32,
	secondary_table_offset: u64,
	secondary: bool,
	) -> Result<(), Error> {
	let primary_header_lba = 1;
	let secondary_header_lba = (secondary_table_offset + GPT_END_SIZE) / SECTOR_SIZE - 1;
	let mut gpt_header = GptHeader {
	signature: *b"EFI PART",
	revision: [0, 0, 1, 0],
	header_size: GPT_HEADER_SIZE,
	current_lba: if secondary {
	secondary_header_lba
	} else {
	primary_header_lba
	},
	backup_lba: if secondary {
	primary_header_lba
	} else {
	secondary_header_lba
	},
	first_usable_lba: GPT_BEGINNING_SIZE / SECTOR_SIZE,
	last_usable_lba: secondary_table_offset / SECTOR_SIZE - 1,
	disk_guid,
	partition_entries_lba: 2,
	num_partition_entries: GPT_NUM_PARTITIONS,
	partition_entry_size: GPT_PARTITION_ENTRY_SIZE,
	partition_entries_crc32,
	header_crc32: 0,
	};

	// Write once to a temporary buffer to calculate the CRC.
	let mut header_without_crc = [0u8; GPT_HEADER_SIZE as usize];
	gpt_header.write_bytes(&mut &mut header_without_crc[..])?;
	let mut hasher = Hasher::new();
	hasher.update(&header_without_crc);
	gpt_header.header_crc32 = hasher.finalize();

	gpt_header.write_bytes(out)?;

	Ok(())
	}

	/// A GPT entry for a particular partition.
	#[derive(Clone, Debug, Eq, PartialEq)]
	pub struct GptPartitionEntry {
	pub partition_type_guid: Uuid,
	pub unique_partition_guid: Uuid,
	pub first_lba: u64,
	pub last_lba: u64,
	pub attributes: u64,
	/// UTF-16LE
	pub partition_name: [u16; 36],
	}

	// This is implemented manually because `Default` isn't implemented in the standard library for
	// arrays of more than 32 elements. If that gets implemented (now than const generics are in) then
	// we can derive this instead.
	impl Default for GptPartitionEntry {
	fn default() -> Self {
	Self {
	partition_type_guid: Default::default(),
	unique_partition_guid: Default::default(),
	first_lba: 0,
	last_lba: 0,
	attributes: 0,
	partition_name: [0; 36],
	}
	}
	}

	impl GptPartitionEntry {
	/// Write out the partition table entry. It will take
	/// `GPT_PARTITION_ENTRY_SIZE` bytes.
	pub fn write_bytes(&self, out: &mut impl Write) -> Result<(), Error> {
	write_guid(out, self.partition_type_guid).map_err(Error::WritingData)?;
	write_guid(out, self.unique_partition_guid).map_err(Error::WritingData)?;
	out.write_all(&self.first_lba.to_le_bytes())
	.map_err(Error::WritingData)?;
	out.write_all(&self.last_lba.to_le_bytes())
	.map_err(Error::WritingData)?;
	out.write_all(&self.attributes.to_le_bytes())
	.map_err(Error::WritingData)?;
	for code_unit in &self.partition_name {
	out.write_all(&code_unit.to_le_bytes())
	.map_err(Error::WritingData)?;
	}
	Ok(())
	}
	}

	/// Write a UUID in the mixed-endian format which GPT uses for GUIDs.
	fn write_guid(out: &mut impl Write, guid: Uuid) -> Result<(), io::Error> {
	let guid_fields = guid.as_fields();
	out.write_all(&guid_fields.0.to_le_bytes())?;
	out.write_all(&guid_fields.1.to_le_bytes())?;
	out.write_all(&guid_fields.2.to_le_bytes())?;
	out.write_all(guid_fields.3)?;

	Ok(())
	}

	#[cfg(test)]
	mod tests {
	use super::*;

	#[test]
	fn protective_mbr_size() {
	let mut buffer = vec![];
	write_protective_mbr(&mut buffer, 1000 * SECTOR_SIZE).unwrap();

	assert_eq!(buffer.len(), SECTOR_SIZE as usize);
	}

	#[test]
	fn header_size() {
	let mut buffer = vec![];
	write_gpt_header(
	&mut buffer,
	Uuid::from_u128(0x12345678_1234_5678_abcd_12345678abcd),
	42,
	1000 * SECTOR_SIZE,
	false,
	)
	.unwrap();

	assert_eq!(buffer.len(), GPT_HEADER_SIZE as usize);
	}

	#[test]
	fn partition_entry_size() {
	let mut buffer = vec![];
	GptPartitionEntry::default()
	.write_bytes(&mut buffer)
	.unwrap();

	assert_eq!(buffer.len(), GPT_PARTITION_ENTRY_SIZE as usize);
	}
	}