Create composite disk image in VirtualizationService.

This is simpler than spawning mk_cdisk, and will also be useful for
making the payload image.

Bug: 190503456
Test: Ran microdroid, compared log output
Change-Id: Id67d6280696c4221b675eec99c65ea44e1c549ab
diff --git a/virtualizationservice/src/gpt.rs b/virtualizationservice/src/gpt.rs
new file mode 100644
index 0000000..346a40a
--- /dev/null
+++ b/virtualizationservice/src/gpt.rs
@@ -0,0 +1,240 @@
+// Copyright 2021, 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.
+
+//! Functions for writing GUID Partition Tables for use in a composite disk image.
+
+use anyhow::Error;
+use crc32fast::Hasher;
+use std::convert::TryInto;
+use std::io::Write;
+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;
+
+/// 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])?;
+
+    // Partition status
+    file.write_all(&[0x00])?;
+    // Begin CHS
+    file.write_all(&[0; 3])?;
+    // Partition type
+    file.write_all(&[0xEE])?;
+    // End CHS
+    file.write_all(&[0; 3])?;
+    let first_lba: u32 = 1;
+    file.write_all(&first_lba.to_le_bytes())?;
+    let number_of_sectors: u32 = (disk_size / SECTOR_SIZE).try_into()?;
+    file.write_all(&number_of_sectors.to_le_bytes())?;
+
+    // Three more empty partitions
+    file.write_all(&[0; MBR_PARTITION_ENTRY_SIZE * 3])?;
+
+    // Boot signature
+    file.write_all(&[0x55, 0xAA])?;
+
+    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)?;
+        out.write_all(&self.revision)?;
+        out.write_all(&self.header_size.to_le_bytes())?;
+        out.write_all(&self.header_crc32.to_le_bytes())?;
+        // Reserved
+        out.write_all(&[0; 4])?;
+        out.write_all(&self.current_lba.to_le_bytes())?;
+        out.write_all(&self.backup_lba.to_le_bytes())?;
+        out.write_all(&self.first_usable_lba.to_le_bytes())?;
+        out.write_all(&self.last_usable_lba.to_le_bytes())?;
+
+        // GUID is mixed-endian for some reason, so we can't just use `Uuid::as_bytes()`.
+        write_guid(out, self.disk_guid)?;
+
+        out.write_all(&self.partition_entries_lba.to_le_bytes())?;
+        out.write_all(&self.num_partition_entries.to_le_bytes())?;
+        out.write_all(&self.partition_entry_size.to_le_bytes())?;
+        out.write_all(&self.partition_entries_crc32.to_le_bytes())?;
+        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],
+}
+
+// TODO: Derive this once arrays of more than 32 elements have default values.
+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)?;
+        write_guid(out, self.unique_partition_guid)?;
+        out.write_all(&self.first_lba.to_le_bytes())?;
+        out.write_all(&self.last_lba.to_le_bytes())?;
+        out.write_all(&self.attributes.to_le_bytes())?;
+        for code_unit in &self.partition_name {
+            out.write_all(&code_unit.to_le_bytes())?;
+        }
+        Ok(())
+    }
+}
+
+/// Write a UUID in the mixed-endian format which GPT uses for GUIDs.
+fn write_guid(out: &mut impl Write, guid: Uuid) -> Result<(), 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);
+    }
+}