/** @file
  Routines supporting partition discovery and
  logical device reading

Copyright (c) 2019 Intel Corporation. All rights reserved.<BR>

This program and the accompanying materials are licensed and made available
under the terms and conditions of the BSD License which accompanies this
distribution. The full text of the license may be found at
http://opensource.org/licenses/bsd-license.php

THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.

**/

#include <IndustryStandard/Mbr.h>
#include <Uefi/UefiGpt.h>
#include <Library/BaseLib.h>
#include "FatLitePeim.h"

//
// Assumption: 'a' and 'blocksize' are all UINT32 or UINT64.
// If 'a' and 'blocksize' are not the same type, should use DivU64xU32 to calculate.
//
#define EFI_SIZE_TO_BLOCKS(a, blocksize)  (((a) / (blocksize)) + (((a) % (blocksize)) ? 1 : 0))

//
// GPT Partition Entry Status
//
typedef struct {
  BOOLEAN OutOfRange;
  BOOLEAN Overlap;
  BOOLEAN OsSpecific;
} EFI_PARTITION_ENTRY_STATUS;

/**
  Check if the CRC field in the Partition table header is valid.

  @param[in]  PartHeader  Partition table header structure

  @retval TRUE      the CRC is valid
  @retval FALSE     the CRC is invalid

**/
BOOLEAN
PartitionCheckGptHeaderCRC (
  IN  EFI_PARTITION_TABLE_HEADER  *PartHeader
  )
{
  UINT32      GptHdrCrc;
  UINT32      Crc;

  GptHdrCrc = PartHeader->Header.CRC32;

  //
  // Set CRC field to zero when doing calcuation
  //
  PartHeader->Header.CRC32 = 0;

  Crc = CalculateCrc32 (PartHeader, PartHeader->Header.HeaderSize);

  //
  // Restore Header CRC
  //
  PartHeader->Header.CRC32 = GptHdrCrc;

  return (GptHdrCrc == Crc);
}


/**
  Check if the CRC field in the Partition table header is valid
  for Partition entry array.

  @param[in]  PartHeader  Partition table header structure
  @param[in]  PartEntry   The partition entry array

  @retval TRUE      the CRC is valid
  @retval FALSE     the CRC is invalid

**/
BOOLEAN
PartitionCheckGptEntryArrayCRC (
  IN  EFI_PARTITION_TABLE_HEADER *PartHeader,
  IN  EFI_PARTITION_ENTRY        *PartEntry
  )
{
  UINT32      Crc;
  UINTN       Size;

  Size = (UINTN)MultU64x32(PartHeader->NumberOfPartitionEntries, PartHeader->SizeOfPartitionEntry);
  Crc  = CalculateCrc32 (PartEntry, Size);

  return (BOOLEAN) (PartHeader->PartitionEntryArrayCRC32 == Crc);
}

/**
  The function is used for valid GPT table. Both for Primary and Backup GPT header.

  @param[in]  PrivateData       The global memory map
  @param[in]  ParentBlockDevNo  The parent block device
  @param[in]  IsPrimaryHeader   Indicate to which header will be checked.
  @param[in]  PartHdr           Stores the partition table that is read

  @retval TRUE      The partition table is valid
  @retval FALSE     The partition table is not valid

**/
BOOLEAN
PartitionCheckGptHeader (
  IN  PEI_FAT_PRIVATE_DATA        *PrivateData,
  IN  UINTN                       ParentBlockDevNo,
  IN  BOOLEAN                     IsPrimaryHeader,
  IN  EFI_PARTITION_TABLE_HEADER  *PartHdr
  )
{
  PEI_FAT_BLOCK_DEVICE            *ParentBlockDev;
  EFI_PEI_LBA                     Lba;
  EFI_PEI_LBA                     AlternateLba;
  EFI_PEI_LBA                     EntryArrayLastLba;

  UINT64                          PartitionEntryArraySize;
  UINT64                          PartitionEntryBlockNumb;
  UINT32                          EntryArraySizeRemainder;

  ParentBlockDev = &(PrivateData->BlockDevice[ParentBlockDevNo]);

  if (IsPrimaryHeader) {
    Lba          = PRIMARY_PART_HEADER_LBA;
    AlternateLba = ParentBlockDev->LastBlock;
  } else {
    Lba          = ParentBlockDev->LastBlock;
    AlternateLba = PRIMARY_PART_HEADER_LBA;
  }

  if ( (PartHdr->Header.Signature != EFI_PTAB_HEADER_ID) ||
       (PartHdr->Header.Revision != 0x00010000) ||
       (PartHdr->Header.HeaderSize < 92) ||
       (PartHdr->Header.HeaderSize > ParentBlockDev->BlockSize) ||
       (!PartitionCheckGptHeaderCRC (PartHdr)) ||
       (PartHdr->Header.Reserved != 0)
     ) {
    DEBUG ((DEBUG_ERROR, "Invalid efi partition table header\n"));
    return FALSE;
  }

  //
  // |    Block0    |    Block1    |Block2 ~ FirstUsableLBA - 1|FirstUsableLBA, ... ,LastUsableLBA|LastUsableLBA+1 ~ LastBlock-1|  LastBlock  |
  // |Protective MBR|Primary Header|Entry Array(At Least 16384)|             Partition            | Entry Array(At Least 16384) |BackUp Header|
  //
  // 1. Protective MBR is fixed at Block 0.
  // 2. Primary Header is fixed at Block 1.
  // 3. Backup Header is fixed at LastBlock.
  // 4. Must be remain 128*128 bytes for primary entry array.
  // 5. Must be remain 128*128 bytes for backup entry array.
  // 6. SizeOfPartitionEntry must be equals to 128 * 2^n.
  //
  if ( (PartHdr->MyLBA != Lba) ||
       (PartHdr->AlternateLBA != AlternateLba) ||
       (PartHdr->FirstUsableLBA < 2 + EFI_SIZE_TO_BLOCKS (EFI_GPT_PART_ENTRY_MIN_SIZE, ParentBlockDev->BlockSize)) ||
       (PartHdr->LastUsableLBA  > ParentBlockDev->LastBlock - 1 - EFI_SIZE_TO_BLOCKS (EFI_GPT_PART_ENTRY_MIN_SIZE, ParentBlockDev->BlockSize)) ||
       (PartHdr->FirstUsableLBA > PartHdr->LastUsableLBA) ||
       (PartHdr->PartitionEntryLBA < 2) ||
       (PartHdr->PartitionEntryLBA > ParentBlockDev->LastBlock - 1) ||
       (PartHdr->PartitionEntryLBA >= PartHdr->FirstUsableLBA && PartHdr->PartitionEntryLBA <= PartHdr->LastUsableLBA) ||
       (PartHdr->SizeOfPartitionEntry%128 != 0) ||
       (PartHdr->SizeOfPartitionEntry != sizeof (EFI_PARTITION_ENTRY))
     ) {
    DEBUG ((DEBUG_ERROR, "Invalid efi partition table header\n"));
    return FALSE;
  }

  //
  // Ensure the NumberOfPartitionEntries * SizeOfPartitionEntry doesn't overflow.
  //
  if (PartHdr->NumberOfPartitionEntries > DivU64x32 (MAX_UINTN, PartHdr->SizeOfPartitionEntry)) {
    DEBUG ((DEBUG_ERROR, "Memory overflow in GPT Entry Array\n"));
    return FALSE;
  }

  PartitionEntryArraySize = MultU64x32 (PartHdr->NumberOfPartitionEntries, PartHdr->SizeOfPartitionEntry);
  EntryArraySizeRemainder = 0;
  PartitionEntryBlockNumb = DivU64x32Remainder (PartitionEntryArraySize, ParentBlockDev->BlockSize, &EntryArraySizeRemainder);
  if (EntryArraySizeRemainder != 0) {
    PartitionEntryBlockNumb++;
  }

  if (IsPrimaryHeader) {
    EntryArrayLastLba = PartHdr->FirstUsableLBA;
  } else {
    EntryArrayLastLba = ParentBlockDev->LastBlock;
  }

  //
  // Make sure partition entry array not overlaps with partition area or the LastBlock.
  //
  if (PartHdr->PartitionEntryLBA + PartitionEntryBlockNumb > EntryArrayLastLba) {
    DEBUG ((DEBUG_ERROR, "GPT Partition Entry Array Error!\n"));
    DEBUG ((DEBUG_ERROR, "PartitionEntryArraySize = %lu.\n", PartitionEntryArraySize));
    DEBUG ((DEBUG_ERROR, "PartitionEntryLBA = %lu.\n", PartHdr->PartitionEntryLBA));
    DEBUG ((DEBUG_ERROR, "PartitionEntryBlockNumb = %lu.\n", PartitionEntryBlockNumb));
    DEBUG ((DEBUG_ERROR, "EntryArrayLastLba = %lu.\n", EntryArrayLastLba));
    return FALSE;
  }

  return TRUE;
}

/**
  This function is used to verify each partition in block device.

  @param[in]  PrivateData       The global memory map
  @param[in]  ParentBlockDevNo  The parent block device
  @param[in]  PartHdr           Stores the partition table that is read

  @retval TRUE      The partition is valid
  @retval FALSE     The partition is not valid

**/
BOOLEAN
PartitionCheckGptEntryArray (
  IN  PEI_FAT_PRIVATE_DATA        *PrivateData,
  IN  UINTN                       ParentBlockDevNo,
  IN  EFI_PARTITION_TABLE_HEADER  *PartHdr
  )
{
  EFI_STATUS                      Status;
  PEI_FAT_BLOCK_DEVICE            *ParentBlockDev;
  PEI_FAT_BLOCK_DEVICE            *BlockDevPtr;

  UINT64                          PartitionEntryArraySize;
  UINT64                          PartitionEntryBlockNumb;
  UINT32                          EntryArraySizeRemainder;

  EFI_PARTITION_ENTRY             *PartitionEntryBuffer;
  EFI_PARTITION_ENTRY_STATUS      *PartitionEntryStatus;

  BOOLEAN                         Found;
  EFI_LBA                         StartingLBA;
  EFI_LBA                         EndingLBA;
  UINTN                           Index;
  UINTN                           Index1;
  UINTN                           Index2;
  EFI_PARTITION_ENTRY             *Entry;

  PartitionEntryBuffer = NULL;
  PartitionEntryStatus = NULL;

  ParentBlockDev  = &(PrivateData->BlockDevice[ParentBlockDevNo]);
  Found           = FALSE;

  PartitionEntryArraySize = MultU64x32 (PartHdr->NumberOfPartitionEntries, PartHdr->SizeOfPartitionEntry);
  EntryArraySizeRemainder = 0;
  PartitionEntryBlockNumb = DivU64x32Remainder (PartitionEntryArraySize, ParentBlockDev->BlockSize, &EntryArraySizeRemainder);
  if (EntryArraySizeRemainder != 0) {
    PartitionEntryBlockNumb++;
  }
  PartitionEntryArraySize = MultU64x32 (PartitionEntryBlockNumb, ParentBlockDev->BlockSize);

  PartitionEntryBuffer = (EFI_PARTITION_ENTRY *) AllocatePages (EFI_SIZE_TO_PAGES ((UINTN)PartitionEntryArraySize));
  if (PartitionEntryBuffer == NULL) {
    DEBUG ((DEBUG_ERROR, "Allocate memory error!\n"));
    goto EXIT;
  }

  PartitionEntryStatus = (EFI_PARTITION_ENTRY_STATUS *) AllocatePages (EFI_SIZE_TO_PAGES (PartHdr->NumberOfPartitionEntries * sizeof (EFI_PARTITION_ENTRY_STATUS)));
  if (PartitionEntryStatus == NULL) {
    DEBUG ((DEBUG_ERROR, "Allocate memory error!\n"));
    goto EXIT;
  }
  ZeroMem (PartitionEntryStatus, PartHdr->NumberOfPartitionEntries * sizeof (EFI_PARTITION_ENTRY_STATUS));

  Status = FatReadBlock (
             PrivateData,
             ParentBlockDevNo,
             PartHdr->PartitionEntryLBA,
             (UINTN)PartitionEntryArraySize,
             PartitionEntryBuffer
             );
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "Read partition entry array error!\n"));
    goto EXIT;
  }

  if (!PartitionCheckGptEntryArrayCRC (PartHdr, PartitionEntryBuffer)) {
    DEBUG ((DEBUG_ERROR, "Partition entries CRC check fail\n"));
    goto EXIT;
  }

  for (Index1 = 0; Index1 < PartHdr->NumberOfPartitionEntries; Index1++) {
    Entry = (EFI_PARTITION_ENTRY *) ((UINT8 *) PartitionEntryBuffer + Index1 * PartHdr->SizeOfPartitionEntry);
    if (CompareGuid (&Entry->PartitionTypeGUID, &gEfiPartTypeUnusedGuid)) {
      continue;
    }

    StartingLBA = Entry->StartingLBA;
    EndingLBA   = Entry->EndingLBA;
    if (StartingLBA > EndingLBA ||
        StartingLBA < PartHdr->FirstUsableLBA ||
        StartingLBA > PartHdr->LastUsableLBA ||
        EndingLBA < PartHdr->FirstUsableLBA ||
        EndingLBA > PartHdr->LastUsableLBA
        ) {
      PartitionEntryStatus[Index1].OutOfRange = TRUE;
      continue;
    }

    if ((Entry->Attributes & BIT1) != 0) {
      //
      // If Bit 1 is set, this indicate that this is an OS specific GUID partition.
      //
      PartitionEntryStatus[Index1].OsSpecific = TRUE;
    }

    for (Index2 = Index1 + 1; Index2 < PartHdr->NumberOfPartitionEntries; Index2++) {
      Entry = (EFI_PARTITION_ENTRY *) ((UINT8 *) PartitionEntryBuffer + Index2 * PartHdr->SizeOfPartitionEntry);
      if (CompareGuid (&Entry->PartitionTypeGUID, &gEfiPartTypeUnusedGuid)) {
        continue;
      }

      if (Entry->EndingLBA >= StartingLBA && Entry->StartingLBA <= EndingLBA) {
        //
        // This region overlaps with the Index1'th region
        //
        PartitionEntryStatus[Index1].Overlap  = TRUE;
        PartitionEntryStatus[Index2].Overlap  = TRUE;
        continue;
      }
    }
  }

  for (Index = 0; Index < PartHdr->NumberOfPartitionEntries; Index++) {
    if (CompareGuid (&PartitionEntryBuffer[Index].PartitionTypeGUID, &gEfiPartTypeUnusedGuid)||
        PartitionEntryStatus[Index].OutOfRange ||
        PartitionEntryStatus[Index].Overlap ||
        PartitionEntryStatus[Index].OsSpecific) {
      //
      // Don't use null EFI Partition Entries, Invalid Partition Entries or OS specific
      // partition Entries
      //
      continue;
    }

    if (PrivateData->BlockDeviceCount >= PEI_FAT_MAX_BLOCK_DEVICE) {
      break;
    }

    Found                         = TRUE;
    BlockDevPtr                   = &(PrivateData->BlockDevice[PrivateData->BlockDeviceCount]);

    BlockDevPtr->BlockSize        = ParentBlockDev->BlockSize;
    BlockDevPtr->LastBlock        = PartitionEntryBuffer[Index].EndingLBA;
    BlockDevPtr->IoAlign          = ParentBlockDev->IoAlign;
    BlockDevPtr->Logical          = TRUE;
    BlockDevPtr->PartitionChecked = FALSE;
    BlockDevPtr->StartingPos      = MultU64x32 (
                                      PartitionEntryBuffer[Index].StartingLBA,
                                      ParentBlockDev->BlockSize
                                      );
    BlockDevPtr->ParentDevNo      = ParentBlockDevNo;

    PrivateData->BlockDeviceCount++;

    DEBUG ((DEBUG_INFO, "Find GPT Partition [0x%lx",  PartitionEntryBuffer[Index].StartingLBA, BlockDevPtr->LastBlock));
    DEBUG ((DEBUG_INFO, ", 0x%lx]\n", BlockDevPtr->LastBlock));
    DEBUG ((DEBUG_INFO, "         BlockSize %x\n",  BlockDevPtr->BlockSize));
  }

EXIT:
  if (PartitionEntryBuffer != NULL) {
    FreePages (PartitionEntryBuffer, EFI_SIZE_TO_PAGES ((UINTN)PartitionEntryArraySize));
  }

  if (PartitionEntryStatus != NULL) {
    FreePages (PartitionEntryStatus, EFI_SIZE_TO_PAGES (PartHdr->NumberOfPartitionEntries * sizeof (EFI_PARTITION_ENTRY_STATUS)));
  }

  return Found;
}

/**
  The function is used to check GPT structure, include GPT header and GPT entry array.

  1. Check GPT header.
  2. Check partition entry array.
  3. Check each partitions.

  @param[in]  PrivateData       The global memory map
  @param[in]  ParentBlockDevNo  The parent block device
  @param[in]  IsPrimary         Indicate primary or backup to be check

  @retval TRUE              Primary or backup GPT structure is valid.
  @retval FALSE             Both primary and backup are invalid.

**/
BOOLEAN
PartitionCheckGptStructure (
  IN  PEI_FAT_PRIVATE_DATA      *PrivateData,
  IN  UINTN                     ParentBlockDevNo,
  IN  BOOLEAN                   IsPrimary
  )
{
  EFI_STATUS                    Status;
  PEI_FAT_BLOCK_DEVICE          *ParentBlockDev;
  EFI_PARTITION_TABLE_HEADER    *PartHdr;
  EFI_PEI_LBA                   GptHeaderLBA;

  ParentBlockDev  = &(PrivateData->BlockDevice[ParentBlockDevNo]);
  PartHdr         = (EFI_PARTITION_TABLE_HEADER *) PrivateData->BlockData;

  if (IsPrimary) {
    GptHeaderLBA = PRIMARY_PART_HEADER_LBA;
  } else {
    GptHeaderLBA = ParentBlockDev->LastBlock;
  }

  Status = FatReadBlock (
             PrivateData,
             ParentBlockDevNo,
             GptHeaderLBA,
             ParentBlockDev->BlockSize,
             PartHdr
             );
  if (EFI_ERROR (Status)) {
    return FALSE;
  }

  if (!PartitionCheckGptHeader (PrivateData, ParentBlockDevNo, IsPrimary, PartHdr)) {
    return FALSE;
  }

  if (!PartitionCheckGptEntryArray (PrivateData, ParentBlockDevNo, PartHdr)) {
    return FALSE;
  }

  return TRUE;
}

/**
  This function is used to check protective MBR structure before checking GPT.

  @param[in]  PrivateData       The global memory map
  @param[in]  ParentBlockDevNo  The parent block device

  @retval TRUE              Valid protective MBR
  @retval FALSE             Invalid MBR
**/
BOOLEAN
PartitionCheckProtectiveMbr (
  IN  PEI_FAT_PRIVATE_DATA    *PrivateData,
  IN  UINTN                   ParentBlockDevNo
  )
{
  EFI_STATUS                  Status;
  MASTER_BOOT_RECORD          *ProtectiveMbr;
  MBR_PARTITION_RECORD        *MbrPartition;
  PEI_FAT_BLOCK_DEVICE        *ParentBlockDev;
  UINTN                       Index;

  ProtectiveMbr   = (MASTER_BOOT_RECORD *) PrivateData->BlockData;
  ParentBlockDev  = &(PrivateData->BlockDevice[ParentBlockDevNo]);

  //
  // Read Protective MBR
  //
  Status = FatReadBlock (
             PrivateData,
             ParentBlockDevNo,
             0,
             ParentBlockDev->BlockSize,
             ProtectiveMbr
             );
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "GPT Error When Read Protective Mbr From Partition!\n"));
    return FALSE;
  }

  if (ProtectiveMbr->Signature != MBR_SIGNATURE) {
    DEBUG ((DEBUG_ERROR, "Protective Mbr Signature is invalid!\n"));
    return FALSE;
  }

  //
  // The partition define in UEFI Spec Table 17.
  // Boot Code, Unique MBR Disk Signature, Unknown.
  // These parts will not be used by UEFI, so we skip to check them.
  //
  for (Index = 0; Index < MAX_MBR_PARTITIONS; Index++) {
    MbrPartition = (MBR_PARTITION_RECORD *)&ProtectiveMbr->Partition[Index];
    if (MbrPartition->BootIndicator   == 0x00 &&
        MbrPartition->StartSector     == 0x02 &&
        MbrPartition->OSIndicator     == PMBR_GPT_PARTITION &&
        UNPACK_UINT32 (MbrPartition->StartingLBA) == 1
       ) {
      return TRUE;
    }
  }

  DEBUG ((DEBUG_ERROR, "Protective Mbr, All Partition Entry Are Empty!\n"));
  return FALSE;
}

/**
  This function is used for finding GPT partition on block device.
  As follow UEFI spec we should check protective MBR first and then
  try to check both primary/backup GPT structures.

  @param[in]  PrivateData       The global memory map
  @param[in]  ParentBlockDevNo  The parent block device

  @retval TRUE              New partitions are detected and logical block devices
                            are added to block device array
  @retval FALSE             No new partitions are added

**/
BOOLEAN
FatFindGptPartitions (
  IN  PEI_FAT_PRIVATE_DATA *PrivateData,
  IN  UINTN                ParentBlockDevNo
  )
{
  BOOLEAN                      Found;
  PEI_FAT_BLOCK_DEVICE         *ParentBlockDev;

  if (ParentBlockDevNo > PEI_FAT_MAX_BLOCK_DEVICE - 1) {
    return FALSE;
  }

  ParentBlockDev  = &(PrivateData->BlockDevice[ParentBlockDevNo]);
  if (ParentBlockDev->BlockSize > PEI_FAT_MAX_BLOCK_SIZE) {
    DEBUG ((DEBUG_ERROR, "Device BlockSize %x exceed FAT_MAX_BLOCK_SIZE\n", ParentBlockDev->BlockSize));
    return FALSE;
  }

  if (!PartitionCheckProtectiveMbr (PrivateData, ParentBlockDevNo)) {
    return FALSE;
  }

  Found = PartitionCheckGptStructure (PrivateData, ParentBlockDevNo, TRUE);
  if (!Found) {
    DEBUG ((DEBUG_ERROR, "Primary GPT Header Error, Try to Check Backup GPT Header!\n"));
    Found = PartitionCheckGptStructure (PrivateData, ParentBlockDevNo, FALSE);
  }

  if (Found) {
    ParentBlockDev->PartitionChecked = TRUE;
  }

  return Found;
}