// Copyright 2022 Alberto Ruiz <aruiz@redhat.com>
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

#![cfg_attr(feature = "no_std", no_std)]

use crc::Crc;

#[cfg(feature = "no_std")]
use core::{mem, ops};

#[cfg(not(feature = "no_std"))]
use std::{mem, ops};

use ops::{Index, Range};

use mem::size_of;

mod endian;

use endian::u16le;
use endian::u32le;
use endian::u64le;

#[repr(C, align(1))]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Uuid {
    time_low: u32le,
    time_mid: u16le,
    time_hi_and_version: u16le,
    clock_seq_hi_and_reseved: u8,
    clock_seq_loq: u8,
    node: [u8; 6],
}

impl Uuid {
    pub fn nil() -> Self {
        Self {
            time_low: 0.into(),
            time_mid: 0.into(),
            time_hi_and_version: 0.into(),
            clock_seq_hi_and_reseved: 0,
            clock_seq_loq: 0,
            node: [0, 0, 0, 0, 0, 0],
        }
    }
}

#[repr(C, align(1))]
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub struct GPTPartition {
    part_type: Uuid,
    id: Uuid,
    first_lba: u64le,
    last_lba: u64le,
    attr: [u8; 8],
    name: [u16le; 36],
}

#[repr(C, align(1))]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct GPTHeader {
    /// GPT header magic signature, hardcoded to "EFI PART".
    pub signature: [u8; 8], // Offset  0. "EFI PART", 45h 46h 49h 20h 50h 41h 52h 54h
    /// 00 00 01 00
    pub revision: u32le, // Offset  8
    /// little endian
    pub header_size_le: u32le, // Offset 12
    /// CRC32 of the header with crc32 section zeroed
    pub header_crc32: u32le, // Offset 16
    /// must be 0
    pub reserved: u32le, // Offset 20
    /// For main header, 1
    pub current_lba: u64le, // Offset 24
    /// LBA for backup header
    pub backup_lba: u64le, // Offset 32
    /// First usable LBA for partitions (primary table last LBA + 1)
    pub first_usable: u64le, // Offset 40
    /// Last usable LBA (secondary partition table first LBA - 1)
    pub last_usable: u64le, // Offset 48
    /// UUID of the disk
    pub disk_guid: Uuid, // Offset 56
    /// Starting LBA of partition entries
    pub part_start_lba: u64le, // Offset 72
    /// Number of partition entries
    pub num_parts: u32le, // Offset 80
    /// Size of a partition entry, usually 128
    pub part_size: u32le, // Offset 84
    /// CRC32 of the partition table
    pub part_table_crc32: u32le, // Offset 88
                                 // RESERVED: [u8; 96 - 512],
}

impl GPTPartition {
    pub unsafe fn partition_table_checksum(table: &[GPTPartition]) -> u32 {
        let table: &[u8] = mem::transmute(table);
        let hasher = Crc::<u32>::new(&crc::CRC_32_JAMCRC);
        hasher.checksum(table) ^ 0xffffffff
    }
}

#[repr(C)]
#[derive(Debug, Eq, PartialEq)]
pub struct LBA<T: Sized + LBAIndex>(T);

#[repr(C)]
pub struct LBA512([u8; 512]);

#[repr(C)]
pub struct LBA4K([u8; 4096]);

pub trait LBAIndex {
    fn index<'a>(&'a self, slice: Range<usize>) -> &'a [u8];
}

impl LBAIndex for LBA512 {
    fn index<'a>(&'a self, slice: Range<usize>) -> &'a [u8] {
        &self.0[slice]
    }
}

impl LBAIndex for LBA4K {
    fn index<'a>(&'a self, slice: Range<usize>) -> &'a [u8] {
        &self.0[slice]
    }
}

impl<T: LBAIndex> Index<Range<usize>> for LBA<T> {
    type Output = [u8];
    fn index<'a>(&'a self, slice: Range<usize>) -> &'a [u8] {
        &self.0.index(slice)
    }
}

impl GPTHeader {
    pub fn verify_signature(&self) -> bool {
        const SIGNATURE: [u8; 8] = *b"EFI PART";
        self.signature == SIGNATURE
    }

    pub unsafe fn parse_gpt_header<'a, T: Sized + LBAIndex>(
        gpt_header_lba: &'a LBA<T>,
    ) -> Result<&'a Self, &'static str> {
        if size_of::<LBA<T>>() > isize::MAX as usize {
            return Err("Size of LBA block larger than the 32bit address space");
        }
        if size_of::<LBA<T>>() < 512 {
            return Err("Size of LBA smaller than 512 bytes");
        }

        let header = (gpt_header_lba as *const LBA<T>) as *const GPTHeader;
        let header = &(*header) as &GPTHeader;
        if !header.verify_signature() {
            return Err("Could not verify GPT header signature `EFI PART`");
        }

        let part_header_size = u32::from(header.part_size) as usize;

        // Partition header size has to be 2^x * 128
        if part_header_size % 128 != 0 {
            return Err("Error in GPT table: Partition size not a multiple of 128");
        }
        if (part_header_size / 128).count_ones() != 1 {
            return Err("Error in GPT table: Partition size not a power of 2");
        }

        // Verify header
        assert_eq!(
            Self::header_checksum::<T>(&(*header)),
            u32::from(header.header_crc32),
            "GPT header CRC32 verification failed"
        );

        //FIXME: Check if there are enough LBAs to fit all partitions

        Ok(header) // Figure out if we can make this a reference instead of a copy
    }

    pub fn parse_partitions<F: Fn(&GPTPartition, &mut T), T>(
        part_array: &[GPTPartition],
        size: usize,
        foreach: F,
        data: &mut T,
    ) {
        //FIXME: Verify that part_header_size * num_parts does not overfloy beyond first_usable LBA
        let nil_uuid = Uuid::nil();

        // Size comes from the GPT Header and it is always a multiple of 128 bytes
        // which is the size of GPTPartition, we use it as a stride in our array
        let stride = size / size_of::<GPTPartition>();

        for i in 0..part_array.len() {
            if i % stride != 0 {
                continue;
            }
            if part_array[i].part_type == nil_uuid {
                continue;
            }
            foreach(&part_array[i], data);
        }
    }

    fn header_checksum<T: LBAIndex>(header: &GPTHeader) -> u32 {
        let verify_size = u32::from(header.header_size_le) as usize;
        let mut verify_header: GPTHeader = *header;
        let v_header_block = (&verify_header as *const GPTHeader) as *const LBA<T>;
        verify_header.header_crc32 = 0.into();
        let hasher = Crc::<u32>::new(&crc::CRC_32_JAMCRC);
        (unsafe { hasher.checksum(&(*v_header_block)[0..verify_size]) }) ^ 0xffffffff
    }
}

#[cfg(test)]
mod tests {
    extern crate base64;
    extern crate flate2;
    use crate::{GPTHeader, GPTPartition};
    use crate::{LBAIndex, LBA, LBA4K, LBA512};

    use std::io::Cursor;
    use std::io::Read;
    use std::vec::Vec;

    // This is a base64 encoded gzipped 100K disk image formatted for 512 bytes block size and two partitions
    const GPT_EXAMPLE1: &[u8] = b"H4sIAAAAAAAAA+3cPUgjQRQA4E1AOLAQezFi5d9ZWQgSwSKGEAsJCRwStBTFyoBgIQQMlnZqqVUOQtA+XRolVpJKjlME7YSzsEkjuQib9gqPHCd8Hwxv583beQNT7wYBn1k0+NVutyOdp6sPvJ2rJBZTI8sLmWwQRIJ8J/Pl549v7yuRsKK762gYy2GcfDx6fsilosX+6+HXfDoWDfPFcIyvr7x84ED8Yxfxy4H90nbycDe+cZvce8pOn9xUWom30/O72Xpiqtq999Ue9R8rTNRjreZgLR2cNeYbczsHhWzfZnroeGareZ9Zqq6FdeU/7gIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACfy0X8cmC/tJ083I1v3Cb3nrLTJzeVVuLt9Pxutp6Yqo6Gdas96j9WmKjHWs3BWjo4a8w35nYOCtm+zfTQ8cxW8z6zVF0L68o96g8AAAAAAAAAAAAAAAAAAAD/g8RiamR5IZMNgkiQ78wzX2ul9/xVuB4JY/c/AN3v8Ccfj54fcqlosf96+DWfjn0P88VwjK+vvPT+9Pyt3wAlEswAkAEA";

    // This is a base64 encoded gzipped 100K disk image formatted for 4Kbytes block size and two partitions
    const GPT_EXAMPLE_LBA4K: &[u8] = b"H4sICFR+fWIAA2xiYTRrLmltZwDt3L9LI0EUAOBdixSJaJErrA57aysJomBiCvEHURCsBRsRGwmiaLS4SrCwCoiV4DVaCEGwE2z8B0RMozYipFYQ1giTVkQu4sH3wTA7b97se+0Uu1HE/6wjaiRJEjefer5wevrvv+4HAACAn2kkX+ydGJoqRVEczTXXS5U/pfd4HPZbt8pUmLNhPuq8WV3pH62dZV6S58b+cEeIb4RxUdvda3/3AAAAwGcc5y67t7aXCzvl3MJ1Ye3hZOBqMb0+/npe3ZypHk7ete796W+qP1jPz9/3jc2edlXq5dTjQSbkZT98CwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAx7nL7q3t5cJOObdwXVh7OBm4Wkyvj7+eVzdnqoeTd6mQl/6m+oP1/Px939jsaVelXk49HmRCXrZN9QEAAAAAAAAAAAAAAAAAAOAnGMkXeyeGpkpRFEdzzfVt8en3e7wn7Mdhbv0HoPUd/lHnzepK/2jtLPOSPDf2h3+F+EYYF7XdvfZ3DwAAAHzGGwzc9x8AkAEA";

    fn test_disk_image<T: LBAIndex>(image: &[u8], uuid_time_low_1: u32, uuid_time_low_2: u32) {
        let mut base_decoded: Vec<u8> = Vec::new();
        let mut cur = Cursor::new(image);
        assert!(
            base64::read::DecoderReader::new(&mut cur, base64::STANDARD)
                .read_to_end(&mut base_decoded)
                .is_ok(),
            "Could not decode base64 blob"
        );
        let mut disk_data: Vec<u8> = Vec::new();
        assert!(
            flate2::read::GzDecoder::new(&base_decoded[..])
                .read_to_end(&mut disk_data)
                .is_ok(),
            "flat2 could not gunzip disk blob"
        );

        // We hardcode LBA512 for the disk, typically you'd get this from the device
        let mbr_lba = disk_data.as_ptr() as *const LBA<T>;

        let gpt_header = unsafe { GPTHeader::parse_gpt_header(&*mbr_lba.add(1)) };
        assert!(
            gpt_header.is_ok(),
            "Could not parse GPT Header: {}",
            gpt_header.unwrap_err()
        );
        let gpt_header = gpt_header.unwrap();

        // This is the offset of LBA blocks starting from the GPT Header LBA
        let partition_lba_offset = u64::from(gpt_header.part_start_lba) as usize;
        let partition_lba_ptr = unsafe { mbr_lba.add(partition_lba_offset) as *const GPTPartition };

        let part_size = u32::from(gpt_header.part_size) as usize;
        let num_parts = u32::from(gpt_header.num_parts) as usize;

        assert!(
            disk_data.len() > (std::mem::size_of::<LBA<T>>() * 2 + (part_size * num_parts)),
            "Partition table is bigger than buffer size"
        );

        let stride = part_size / std::mem::size_of::<GPTPartition>();
        let length = stride * num_parts;
        let part_array_ptr = core::ptr::slice_from_raw_parts(partition_lba_ptr, length);
        let part_array = unsafe { &*part_array_ptr };

        let checksum_array =
            core::ptr::slice_from_raw_parts(partition_lba_ptr, num_parts * part_size);
        let sum = unsafe { GPTPartition::partition_table_checksum(&*checksum_array) };

        assert_eq!(
            sum,
            u32::from(gpt_header.part_table_crc32),
            "Bad CRC for partition array"
        );

        let mut partitions: Vec<GPTPartition> = Vec::new();
        GPTHeader::parse_partitions(
            part_array,
            u32::from(gpt_header.part_size) as usize,
            |part, parts| {
                parts.push(part.clone());
            },
            &mut partitions,
        );

        assert_eq!(partitions.len(), 2, "List of partitions should have been 2");
        let guid_time_low = u32::from(partitions[0].part_type.time_low);
        assert_eq!(
            guid_time_low, uuid_time_low_1,
            "First partition type did not match (time_low) expected value {:x}: {:x}",
            uuid_time_low_1, guid_time_low
        );
        let guid_time_low = u32::from(partitions[1].part_type.time_low);
        assert_eq!(
            guid_time_low, uuid_time_low_2,
            "Second partition type did not match (time_low) expected value {:x}: {:x}",
            uuid_time_low_2, guid_time_low
        );
    }

    #[test]
    fn parse_512_header() {
        test_disk_image::<LBA512>(GPT_EXAMPLE1, 0x0fc63daf, 0xc12a7328);
    }

    #[test]
    fn parse_4k_header() {
        test_disk_image::<LBA4K>(GPT_EXAMPLE_LBA4K, 0x0fc63daf, 0x0fc63daf);
    }
}