Use composite disk code from crosvm.

Bug: 190503456
Test: atest MicrodroidHostTestCases
Change-Id: Iea978867ad5439b1cdd39d288c99c4ee8ece6f58
diff --git a/virtualizationservice/src/composite.rs b/virtualizationservice/src/composite.rs
index ded0053..40c7e5e 100644
--- a/virtualizationservice/src/composite.rs
+++ b/virtualizationservice/src/composite.rs
@@ -14,279 +14,12 @@
 
 //! Functions for creating a composite disk image.
 
-use crate::gpt::{
-    write_gpt_header, write_protective_mbr, GptPartitionEntry, GPT_BEGINNING_SIZE, GPT_END_SIZE,
-    GPT_HEADER_SIZE, GPT_NUM_PARTITIONS, GPT_PARTITION_ENTRY_SIZE, SECTOR_SIZE,
-};
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::Partition::Partition;
-use anyhow::{anyhow, bail, Context, Error};
-use crc32fast::Hasher;
-use disk::create_disk_file;
-use log::{trace, warn};
-use protobuf::Message;
-use protos::cdisk_spec::{ComponentDisk, CompositeDisk, ReadWriteCapability};
-use std::convert::TryInto;
+use anyhow::{anyhow, Context, Error};
+use disk::{create_composite_disk, create_disk_file, ImagePartitionType, PartitionInfo};
 use std::fs::{File, OpenOptions};
-use std::io::Write;
 use std::os::unix::io::AsRawFd;
 use std::path::{Path, PathBuf};
-use uuid::Uuid;
-
-/// A magic string placed at the beginning of a composite disk file to identify it.
-const CDISK_MAGIC: &str = "composite_disk\x1d";
-/// The version of the composite disk format supported by this implementation.
-const COMPOSITE_DISK_VERSION: u64 = 1;
-/// The amount of padding needed between the last partition entry and the first partition, to align
-/// the partition appropriately. The two sectors are for the MBR and the GPT header.
-const PARTITION_ALIGNMENT_SIZE: usize = GPT_BEGINNING_SIZE as usize
-    - 2 * SECTOR_SIZE as usize
-    - GPT_NUM_PARTITIONS as usize * GPT_PARTITION_ENTRY_SIZE as usize;
-const HEADER_PADDING_LENGTH: usize = SECTOR_SIZE as usize - GPT_HEADER_SIZE as usize;
-// Keep all partitions 4k aligned for performance.
-const PARTITION_SIZE_SHIFT: u8 = 12;
-// Keep the disk size a multiple of 64k for crosvm's virtio_blk driver.
-const DISK_SIZE_SHIFT: u8 = 16;
-
-const LINUX_FILESYSTEM_GUID: Uuid = Uuid::from_u128(0x0FC63DAF_8483_4772_8E79_3D69D8477DE4);
-const EFI_SYSTEM_PARTITION_GUID: Uuid = Uuid::from_u128(0xC12A7328_F81F_11D2_BA4B_00A0C93EC93B);
-
-/// Information about a partition to create.
-#[derive(Clone, Debug, Eq, PartialEq)]
-pub struct PartitionInfo {
-    label: String,
-    path: PathBuf,
-    partition_type: ImagePartitionType,
-    writable: bool,
-    size: u64,
-}
-
-/// Round `val` up to the next multiple of 2**`align_log`.
-fn align_to_power_of_2(val: u64, align_log: u8) -> u64 {
-    let align = 1 << align_log;
-    ((val + (align - 1)) / align) * align
-}
-
-/// Round `val` to partition size(4K)
-fn align_to_partition_size(val: u64) -> u64 {
-    align_to_power_of_2(val, PARTITION_SIZE_SHIFT)
-}
-
-impl PartitionInfo {
-    fn aligned_size(&self) -> u64 {
-        align_to_partition_size(self.size)
-    }
-}
-
-/// The type of partition.
-#[allow(dead_code)]
-#[derive(Copy, Clone, Debug, Eq, PartialEq)]
-pub enum ImagePartitionType {
-    LinuxFilesystem,
-    EfiSystemPartition,
-}
-
-impl ImagePartitionType {
-    fn guid(self) -> Uuid {
-        match self {
-            Self::LinuxFilesystem => LINUX_FILESYSTEM_GUID,
-            Self::EfiSystemPartition => EFI_SYSTEM_PARTITION_GUID,
-        }
-    }
-}
-
-/// Write protective MBR and primary GPT table.
-fn write_beginning(
-    file: &mut impl Write,
-    disk_guid: Uuid,
-    partitions: &[u8],
-    partition_entries_crc32: u32,
-    secondary_table_offset: u64,
-    disk_size: u64,
-) -> Result<(), Error> {
-    // Write the protective MBR to the first sector.
-    write_protective_mbr(file, disk_size)?;
-
-    // Write the GPT header, and pad out to the end of the sector.
-    write_gpt_header(file, disk_guid, partition_entries_crc32, secondary_table_offset, false)?;
-    file.write_all(&[0; HEADER_PADDING_LENGTH])?;
-
-    // Write partition entries, including unused ones.
-    file.write_all(partitions)?;
-
-    // Write zeroes to align the first partition appropriately.
-    file.write_all(&[0; PARTITION_ALIGNMENT_SIZE])?;
-
-    Ok(())
-}
-
-/// Write secondary GPT table.
-fn write_end(
-    file: &mut impl Write,
-    disk_guid: Uuid,
-    partitions: &[u8],
-    partition_entries_crc32: u32,
-    secondary_table_offset: u64,
-    disk_size: u64,
-) -> Result<(), Error> {
-    // Write partition entries, including unused ones.
-    file.write_all(partitions)?;
-
-    // Write the GPT header, and pad out to the end of the sector.
-    write_gpt_header(file, disk_guid, partition_entries_crc32, secondary_table_offset, true)?;
-    file.write_all(&[0; HEADER_PADDING_LENGTH])?;
-
-    // Pad out to the aligned disk size.
-    let used_disk_size = secondary_table_offset + GPT_END_SIZE;
-    let padding = disk_size - used_disk_size;
-    file.write_all(&vec![0; padding as usize])?;
-
-    Ok(())
-}
-
-/// Create the `GptPartitionEntry` for the given partition.
-fn create_gpt_entry(partition: &PartitionInfo, offset: u64) -> GptPartitionEntry {
-    let mut partition_name: Vec<u16> = partition.label.encode_utf16().collect();
-    partition_name.resize(36, 0);
-
-    GptPartitionEntry {
-        partition_type_guid: partition.partition_type.guid(),
-        unique_partition_guid: Uuid::new_v4(),
-        first_lba: offset / SECTOR_SIZE,
-        last_lba: (offset + partition.aligned_size()) / SECTOR_SIZE - 1,
-        attributes: 0,
-        partition_name: partition_name.try_into().unwrap(),
-    }
-}
-
-/// Create one or more `ComponentDisk` proto messages for the given partition.
-fn create_component_disks(
-    partition: &PartitionInfo,
-    offset: u64,
-    zero_filler_path: &str,
-) -> Result<Vec<ComponentDisk>, Error> {
-    let aligned_size = partition.aligned_size();
-
-    let mut component_disks = vec![ComponentDisk {
-        offset,
-        file_path: partition.path.to_str().context("Invalid partition path")?.to_string(),
-        read_write_capability: if partition.writable {
-            ReadWriteCapability::READ_WRITE
-        } else {
-            ReadWriteCapability::READ_ONLY
-        },
-        ..ComponentDisk::new()
-    }];
-
-    if partition.size != aligned_size {
-        if partition.writable {
-            bail!(
-                "Read-write partition {:?} size is not a multiple of {}.",
-                partition,
-                1 << PARTITION_SIZE_SHIFT
-            );
-        } else {
-            // Fill in the gap by reusing the header file, because we know it is always bigger
-            // than the alignment size (i.e. GPT_BEGINNING_SIZE > 1 << PARTITION_SIZE_SHIFT).
-            warn!(
-                "Read-only partition {:?} size is not a multiple of {}, filling gap.",
-                partition,
-                1 << PARTITION_SIZE_SHIFT
-            );
-            component_disks.push(ComponentDisk {
-                offset: offset + partition.size,
-                file_path: zero_filler_path.to_owned(),
-                read_write_capability: ReadWriteCapability::READ_ONLY,
-                ..ComponentDisk::new()
-            });
-        }
-    }
-
-    Ok(component_disks)
-}
-
-/// Create a new composite disk containing the given partitions, and write it out to the given
-/// files.
-pub fn create_composite_disk(
-    partitions: &[PartitionInfo],
-    zero_filler_path: &Path,
-    header_path: &Path,
-    header_file: &mut File,
-    footer_path: &Path,
-    footer_file: &mut File,
-    output_composite: &mut File,
-) -> Result<(), Error> {
-    let zero_filler_path =
-        zero_filler_path.to_str().context("Invalid zero filler path")?.to_string();
-    let header_path = header_path.to_str().context("Invalid header path")?.to_string();
-    let footer_path = footer_path.to_str().context("Invalid footer path")?.to_string();
-
-    let mut composite_proto = CompositeDisk::new();
-    composite_proto.version = COMPOSITE_DISK_VERSION;
-    composite_proto.component_disks.push(ComponentDisk {
-        file_path: header_path,
-        offset: 0,
-        read_write_capability: ReadWriteCapability::READ_ONLY,
-        ..ComponentDisk::new()
-    });
-
-    // Write partitions to a temporary buffer so that we can calculate the CRC, and construct the
-    // ComponentDisk proto messages at the same time.
-    let mut partitions_buffer =
-        [0u8; GPT_NUM_PARTITIONS as usize * GPT_PARTITION_ENTRY_SIZE as usize];
-    let mut writer: &mut [u8] = &mut partitions_buffer;
-    let mut next_disk_offset = GPT_BEGINNING_SIZE;
-    for partition in partitions {
-        create_gpt_entry(partition, next_disk_offset).write_bytes(&mut writer)?;
-
-        for component_disk in
-            create_component_disks(partition, next_disk_offset, &zero_filler_path)?
-        {
-            composite_proto.component_disks.push(component_disk);
-        }
-
-        next_disk_offset += partition.aligned_size();
-    }
-    let secondary_table_offset = next_disk_offset;
-    let disk_size = align_to_power_of_2(secondary_table_offset + GPT_END_SIZE, DISK_SIZE_SHIFT);
-    trace!("Partitions: {:#?}", partitions);
-    trace!("Secondary table offset: {} disk size: {}", secondary_table_offset, disk_size);
-
-    composite_proto.component_disks.push(ComponentDisk {
-        file_path: footer_path,
-        offset: secondary_table_offset,
-        read_write_capability: ReadWriteCapability::READ_ONLY,
-        ..ComponentDisk::new()
-    });
-
-    // Calculate CRC32 of partition entries.
-    let mut hasher = Hasher::new();
-    hasher.update(&partitions_buffer);
-    let partition_entries_crc32 = hasher.finalize();
-
-    let disk_guid = Uuid::new_v4();
-    write_beginning(
-        header_file,
-        disk_guid,
-        &partitions_buffer,
-        partition_entries_crc32,
-        secondary_table_offset,
-        disk_size,
-    )?;
-    write_end(
-        footer_file,
-        disk_guid,
-        &partitions_buffer,
-        partition_entries_crc32,
-        secondary_table_offset,
-        disk_size,
-    )?;
-
-    composite_proto.length = disk_size;
-    output_composite.write_all(CDISK_MAGIC.as_bytes())?;
-    composite_proto.write_to_writer(output_composite)?;
-
-    Ok(())
-}
 
 /// Constructs a composite disk image for the given list of partitions, and opens it ready to use.
 ///
@@ -390,63 +123,3 @@
         .map_err(|e| anyhow!("Failed to open partition image: {}", e))?
         .get_len()?)
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn beginning_size() {
-        let mut buffer = vec![];
-        let partitions = [0u8; GPT_NUM_PARTITIONS as usize * GPT_PARTITION_ENTRY_SIZE as usize];
-        let disk_size = 1000 * SECTOR_SIZE;
-        write_beginning(
-            &mut buffer,
-            Uuid::from_u128(0x12345678_1234_5678_abcd_12345678abcd),
-            &partitions,
-            42,
-            disk_size - GPT_END_SIZE,
-            disk_size,
-        )
-        .unwrap();
-
-        assert_eq!(buffer.len(), GPT_BEGINNING_SIZE as usize);
-    }
-
-    #[test]
-    fn end_size() {
-        let mut buffer = vec![];
-        let partitions = [0u8; GPT_NUM_PARTITIONS as usize * GPT_PARTITION_ENTRY_SIZE as usize];
-        let disk_size = 1000 * SECTOR_SIZE;
-        write_end(
-            &mut buffer,
-            Uuid::from_u128(0x12345678_1234_5678_abcd_12345678abcd),
-            &partitions,
-            42,
-            disk_size - GPT_END_SIZE,
-            disk_size,
-        )
-        .unwrap();
-
-        assert_eq!(buffer.len(), GPT_END_SIZE as usize);
-    }
-
-    #[test]
-    fn end_size_with_padding() {
-        let mut buffer = vec![];
-        let partitions = [0u8; GPT_NUM_PARTITIONS as usize * GPT_PARTITION_ENTRY_SIZE as usize];
-        let disk_size = 1000 * SECTOR_SIZE;
-        let padding = 3 * SECTOR_SIZE;
-        write_end(
-            &mut buffer,
-            Uuid::from_u128(0x12345678_1234_5678_abcd_12345678abcd),
-            &partitions,
-            42,
-            disk_size - GPT_END_SIZE - padding,
-            disk_size,
-        )
-        .unwrap();
-
-        assert_eq!(buffer.len(), GPT_END_SIZE as usize + padding as usize);
-    }
-}