// 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);
    }
}