/** @file
  Routines supporting partition discovery

Copyright (c) 2006 - 2017, 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 <PiPei.h>
#include <IndustryStandard/Mbr.h>
#include <Uefi/UefiGpt.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/Crc32Lib.h>
#include <Library/DebugLib.h>
#include <Library/PartitionLib.h>
#include <Library/MediaAccessLib.h>

//
// Extract INT32 from char array
//
#define UNPACK_INT32(a) \
  (INT32) ((((UINT8 *) a)[0] << 0) | (((UINT8 *) a)[1] << 8) | (((UINT8 *) a)[2] << 16) | (((UINT8 *) a)[3] << 24))

//
// Extract UINT32 from char array
//
#define UNPACK_UINT32(a) \
  (UINT32) ((((UINT8 *) a)[0] << 0) | (((UINT8 *) a)[1] << 8) | (((UINT8 *) a)[2] << 16) | (((UINT8 *) a)[3] << 24))


CHAR8 *mPartTypeName[] = {
  "UNKNOWN",
  "MBR",
  "GPT"
};

extern
EFI_STATUS
FindSpiPartitions (
  IN OUT PART_BLOCK_DEVICE  *PartBlockDev
  );

/**
  Convert partition type into readable string.

  @param  Type              Partition type

  @retval                   Ascii string for partition type

**/
CHAR8 *
GetPartitionTypeName (
  IN UINTN    Type
  )
{
  if (Type >= EnumPartTypeMax) {
    Type = EnumPartTypeUnknown;
  }
  return mPartTypeName[Type];
}


/**
  Test to see if the Mbr buffer is a valid MBR.

  @param  Mbr               Parent Handle
  @param  LastLba           Last Lba address on the device.

  @retval TRUE              Mbr is a Valid MBR
  @retval FALSE             Mbr is not a Valid MBR

**/
BOOLEAN
PartitionValidMbr (
  IN  MASTER_BOOT_RECORD      *Mbr,
  IN  EFI_PEI_LBA             LastLba
  )
{
  UINT32  StartingLBA;
  UINT32  EndingLBA;
  UINT32  NewEndingLBA;
  INTN    Index1;
  INTN    Index2;
  BOOLEAN MbrValid;

  if (Mbr->Signature != MBR_SIGNATURE) {
    return FALSE;
  }

  //
  // The BPB also has this signature, so it can not be used alone.
  //
  MbrValid = FALSE;
  for (Index1 = 0; Index1 < MAX_MBR_PARTITIONS; Index1++) {
    if (Mbr->Partition[Index1].OSIndicator == 0x00 || UNPACK_UINT32 (Mbr->Partition[Index1].SizeInLBA) == 0) {
      continue;
    }

    MbrValid    = TRUE;
    StartingLBA = UNPACK_UINT32 (Mbr->Partition[Index1].StartingLBA);
    EndingLBA   = StartingLBA + UNPACK_UINT32 (Mbr->Partition[Index1].SizeInLBA) - 1;
    if (EndingLBA > LastLba) {
      //
      // Compatability Errata:
      //  Some systems try to hide drive space with thier INT 13h driver
      //  This does not hide space from the OS driver. This means the MBR
      //  that gets created from DOS is smaller than the MBR created from
      //  a real OS (NT & Win98). This leads to BlockIo->LastBlock being
      //  wrong on some systems FDISKed by the OS.
      //
      //  return FALSE Because no block devices on a system are implemented
      //  with INT 13h
      //
      return FALSE;
    }

    for (Index2 = Index1 + 1; Index2 < MAX_MBR_PARTITIONS; Index2++) {
      if (Mbr->Partition[Index2].OSIndicator == 0x00 || UNPACK_INT32 (Mbr->Partition[Index2].SizeInLBA) == 0) {
        continue;
      }

      NewEndingLBA = UNPACK_UINT32 (Mbr->Partition[Index2].StartingLBA) + UNPACK_UINT32 (Mbr->Partition[Index2].SizeInLBA) -
                     1;
      if (NewEndingLBA >= StartingLBA && UNPACK_UINT32 (Mbr->Partition[Index2].StartingLBA) <= EndingLBA) {
        //
        // This region overlaps with the Index1'th region
        //
        return FALSE;
      }
    }
  }
  //
  // Non of the regions overlapped so MBR is O.K.
  //
  return MbrValid;
}


/**
  This function finds Mbr partitions. Main algorithm
  is ported from DXE partition driver.

  @param[in]  PartBlockDev   Parition block device pointer

  @retval EFI_SUCCESS        New partitions are detected and logical block devices
                             are  added to block device array
  @retval EFI_NOT_FOUND      No New partitions are added
  @retval EFI_DEVICE_ERROR   Error occured for device access

**/
EFI_STATUS
FindMbrPartitions (
  IN  PART_BLOCK_DEVICE   *PartBlockDev
  )
{
  EFI_STATUS             Status;
  MASTER_BOOT_RECORD     *Mbr;
  UINTN                  Index;
  LOGICAL_BLOCK_DEVICE  *BlockDev;
  UINTN                  ParentBlockDevNo;
  DEVICE_BLOCK_INFO     *DevBlockInfo;

  DevBlockInfo = &PartBlockDev->BlockInfo;
  ParentBlockDevNo = PartBlockDev->HarewareDevice;
  Mbr    = (MASTER_BOOT_RECORD *) PartBlockDev->BlockData;
  Status = MediaReadBlocks (
             ParentBlockDevNo,
             0,
             DevBlockInfo->BlockSize,
             Mbr
             );
  if (EFI_ERROR (Status)) {
    return EFI_DEVICE_ERROR;
  }

  Status = EFI_NOT_FOUND;
  if (!PartitionValidMbr (Mbr, DevBlockInfo->BlockNum - 1)) {
    return Status;
  }

  //
  // We have a valid mbr - add each partition
  //
  for (Index = 0; Index < MAX_MBR_PARTITIONS; Index++) {
    if (Mbr->Partition[Index].OSIndicator == 0x00 || UNPACK_INT32 (Mbr->Partition[Index].SizeInLBA) == 0) {
      //
      // Don't use null MBR entries
      //
      continue;
    }
    //
    // Register this partition
    //
    if (PartBlockDev->BlockDeviceCount < PART_MAX_BLOCK_DEVICE) {
      Status                      = EFI_SUCCESS;
      BlockDev                    = & (PartBlockDev->BlockDevice[PartBlockDev->BlockDeviceCount]);
      BlockDev->StartBlock        = UNPACK_UINT32 (Mbr->Partition[Index].StartingLBA);
      BlockDev->LastBlock         = BlockDev->StartBlock + UNPACK_INT32 (Mbr->Partition[Index].SizeInLBA) - 1;
      PartBlockDev->BlockDeviceCount++;
    }
  }

  if (!EFI_ERROR (Status)) {
    PartBlockDev->PartitionType = EnumPartTypeMbr;
  }

  PartBlockDev->PartitionChecked = TRUE;

  return Status;
}


/**
  Test to see if the GPT Entry Array CRC is valid.

  @param  GptEntries        GPT Entry Array Buffer
  @param  Gpt               GPT Header info
  @param  ReadSize          Size of the GPT Array Buffer in bytes
  @param  DevBlockInfo      Device block info
  @param  ParentBlockDevNo  Hardware partition of the device to read
                            from

  @retval TRUE              GPT Entry Array has a valid CRC
  @retval FALSE             GPT Entry Array does not have
                            a valid CRC

**/
BOOLEAN
PartitionValidGptEntryArray (
  IN OUT  EFI_PARTITION_ENTRY         *GptEntries,
  IN      EFI_PARTITION_TABLE_HEADER  *Gpt,
  IN      UINTN                        ReadSize,
  IN      DEVICE_BLOCK_INFO           *DevBlockInfo,
  IN      UINTN                        ParentBlockDevNo
  )
{
  UINT32 CrcOut;
  EFI_STATUS Status;
  EFI_LBA GptLba;

  GptLba = Gpt->PartitionEntryLBA;

  Status = MediaReadBlocks (ParentBlockDevNo, GptLba, ReadSize, GptEntries);
  if (EFI_ERROR (Status)) {
    return FALSE;
  }

  Status = CalculateCrc32WithType ((UINT8 *)GptEntries, ReadSize, Crc32TypeDefault, &CrcOut);
  if (EFI_ERROR (Status) || (CrcOut != Gpt->PartitionEntryArrayCRC32)) {
    return FALSE;
  }

  return TRUE;
}


/**
  Test to see if the GPT Header CRC is valid.

  @param  Gpt               GPT Header info
  @param  DevBlockInfo      Device block info
  @param  ParentBlockDevNo  Hardware partition of the device to read
                            from
  @param  Primary           Inidicator to check primary or secondary
                            GPT Header

  @retval TRUE              GPT Header has a valid CRC
  @retval FALSE             GPT Header does not have
                            a valid CRC

**/
BOOLEAN
PartitionValidGpt (
  IN OUT  EFI_PARTITION_TABLE_HEADER  *Gpt,
  IN      DEVICE_BLOCK_INFO           *DevBlockInfo,
  IN      UINTN                        ParentBlockDevNo,
  IN      BOOLEAN                      Primary
  )
{
  UINT32 CrcOut;
  EFI_STATUS Status;
  EFI_PARTITION_TABLE_HEADER GptBuff;

  Status = MediaReadBlocks (
             ParentBlockDevNo,
             (Primary) ? 1 : DevBlockInfo->BlockNum - 1,
             DevBlockInfo->BlockSize,
             Gpt
             );
  if (EFI_ERROR (Status)) {
    return FALSE;
  }
  CopyMem (&GptBuff,  Gpt, sizeof (EFI_PARTITION_TABLE_HEADER));
  GptBuff.Header.CRC32 = 0;
  Status = CalculateCrc32WithType ((UINT8 *) (&GptBuff), sizeof (EFI_PARTITION_TABLE_HEADER), Crc32TypeDefault, &CrcOut);
  if (EFI_ERROR (Status) || (CrcOut != Gpt->Header.CRC32)) {
    return FALSE;
  }

  return TRUE;
}


/**
  This function finds Gpt partitions. Main algorithm
  is ported from DXE partition driver.

  @param[in]  PartBlockDev   Parition block device pointer

  @retval EFI_SUCCESS        New partitions are detected and logical block devices
                             are  added to block device array
  @retval EFI_NOT_FOUND      No New partitions are added
  @retval EFI_DEVICE_ERROR   Error occured for device access
  @retval EFI_CRC_ERROR      Error occured with primary and secondary GPT CRC
  @retval EFI_UNSUPPORTED    The number of GPT entries exceeds the maximum number allowed

**/
EFI_STATUS
FindGptPartitions (
  IN  PART_BLOCK_DEVICE   *PartBlockDev
  )
{
  EFI_STATUS                   Status;
  UINTN                        Index;
  UINTN                        ReadSize;
  UINTN                        GptPartCnt;
  UINTN                        ParentBlockDevNo;
  LOGICAL_BLOCK_DEVICE        *BlockDev;
  MASTER_BOOT_RECORD          *Mbr;
  UINTN                        ProtectMbrStartLba;
  UINTN                        ProtectMbrLbaSize;
  EFI_PARTITION_TABLE_HEADER  *Gpt;
  UINT64                       MaxBlkNum;
  BOOLEAN                      IsGptValid;
  BOOLEAN                      SecondaryHeader;
  EFI_PARTITION_ENTRY         *GptEntries;
  DEVICE_BLOCK_INFO           *DevBlockInfo;

  DevBlockInfo = &PartBlockDev->BlockInfo;
  ParentBlockDevNo = PartBlockDev->HarewareDevice;
  SecondaryHeader = FALSE;
  GptEntries = NULL;

  if (DevBlockInfo->BlockSize == 0 || DevBlockInfo->BlockNum == 0) {
    Status = EFI_INVALID_PARAMETER;
    goto Done;
  }
  MaxBlkNum     = DevBlockInfo->BlockNum;

  Mbr    = (MASTER_BOOT_RECORD *) PartBlockDev->BlockData;
  Status = MediaReadBlocks (
             ParentBlockDevNo,
             0,
             DevBlockInfo->BlockSize,
             Mbr
             );
  if (EFI_ERROR (Status)) {
    return EFI_DEVICE_ERROR;
  }

  if (!PartitionValidMbr (Mbr, DevBlockInfo->BlockNum - 1)) {
    DEBUG ((DEBUG_WARN, "Protective MBR is not valid format!\n"));
  } else {
    ProtectMbrStartLba = * ((UINT32 *) (Mbr->Partition[0].StartingLBA));
    ProtectMbrLbaSize  = * ((UINT32 *) (Mbr->Partition[0].SizeInLBA));

    if ((ProtectMbrStartLba != 1) || (ProtectMbrLbaSize != (MaxBlkNum - 1))) {
      DEBUG ((DEBUG_WARN, "Protective MBR does not protect the whole drive!\n"));
    }
  }

  Gpt = (EFI_PARTITION_TABLE_HEADER *) PartBlockDev->BlockData;

  IsGptValid = PartitionValidGpt (Gpt, DevBlockInfo, ParentBlockDevNo, TRUE);
  if (!IsGptValid) {
    SecondaryHeader = TRUE;
    DEBUG ((DEBUG_WARN, "Primary GPT Header CalculateCrc32 Error!\n"));
    IsGptValid = PartitionValidGpt (Gpt, DevBlockInfo, ParentBlockDevNo, FALSE);
    if (!IsGptValid) {
      DEBUG ((DEBUG_WARN, "Secondary GPT Header CalculateCrc32 Error!\n"));
      Status = EFI_CRC_ERROR;
      goto Done;
    }
  }

  ReadSize = Gpt->NumberOfPartitionEntries * Gpt->SizeOfPartitionEntry;
  ReadSize = (ReadSize % DevBlockInfo->BlockSize) == 0 ? ReadSize : DevBlockInfo->BlockSize * ((
               ReadSize / DevBlockInfo->BlockSize) + 1);
  GptEntries = (EFI_PARTITION_ENTRY *) AllocatePool (ReadSize);
  if (GptEntries == NULL) {
    Status = EFI_OUT_OF_RESOURCES;
    goto Done;
  }

  IsGptValid = PartitionValidGptEntryArray (GptEntries, Gpt, ReadSize, DevBlockInfo, ParentBlockDevNo);
  if (!IsGptValid) {
    DEBUG ((DEBUG_ERROR, "GPT Array CalculateCrc32 Error!\n"));
    if (SecondaryHeader) {
      DEBUG ((DEBUG_ERROR, "Primary Header and Secondary GPT Array data is corrupt!\n"));
      Status = EFI_CRC_ERROR;
      goto Done;
    }
    IsGptValid = PartitionValidGpt (Gpt, DevBlockInfo, ParentBlockDevNo, FALSE);
    if (!IsGptValid) {
      DEBUG ((DEBUG_ERROR, "Secondary GPT Header CalculateCrc32 Error!\n"));
      Status = EFI_CRC_ERROR;
      goto Done;
    }
    IsGptValid = PartitionValidGptEntryArray (GptEntries, Gpt, ReadSize, DevBlockInfo, ParentBlockDevNo);
    if (!IsGptValid) {
      DEBUG ((DEBUG_ERROR, "Secondary GPT Array CalculateCrc32 Error!\n"));
      Status = EFI_CRC_ERROR;
      goto Done;
    }
  }

  GptPartCnt = Gpt->NumberOfPartitionEntries;

  if (GptPartCnt > GPT_MAX_NUM_PARTITIONS) {
    Status = EFI_UNSUPPORTED;
    goto Done;
  }

  Status = EFI_NOT_FOUND;
  for (Index = 0; Index < GptPartCnt; Index++) {
    if (CompareGuid (&GptEntries[Index].PartitionTypeGUID, &gEfiPartTypeUnusedGuid)) {
      continue;
    }
    if (PartBlockDev->BlockDeviceCount < PART_MAX_BLOCK_DEVICE) {
      Status  = EFI_SUCCESS;
      BlockDev                    = & (PartBlockDev->BlockDevice[PartBlockDev->BlockDeviceCount]);
      BlockDev->StartBlock        = GptEntries[Index].StartingLBA;
      BlockDev->LastBlock         = GptEntries[Index].EndingLBA;
      DEBUG ((DEBUG_INFO, "Part %02d: %12s ", PartBlockDev->BlockDeviceCount, GptEntries[Index].PartitionName));
      DEBUG ((DEBUG_INFO, "0x%08x--0x%08x, LBA count: 0x%x\n", (UINT32)BlockDev->StartBlock, \
              (UINT32)BlockDev->LastBlock, (UINT32) (BlockDev->LastBlock - BlockDev->StartBlock + 1)));

      PartBlockDev->BlockDeviceCount++;
    }
  }

Done:
  if (GptEntries != NULL) {
    FreePool (GptEntries);
  }

  if (!EFI_ERROR (Status)) {
    PartBlockDev->PartitionType = EnumPartTypeGpt;
  }

  PartBlockDev->PartitionChecked = TRUE;

  return Status;
}

/**
Get logical partition information.

This function will fill in logical partition information.

  @param[in]     SwPart           Software partition index.
  @param[in]     PartHandle       Partition handle.
  @param[in,out] LogicBlkDev      Pointer of the logic block device LOGICAL_BLOCK_DEVICE
                                  structure to receive block start and end information.
  @retval EFI_SUCCESS             The logic partition information was filled successfully.
  @retval EFI_INVALID_PARAMETER   Invalid parameters.

**/
EFI_STATUS
EFIAPI
GetLogicalPartitionInfo (
  IN     UINT32                 SwPart,
  IN     EFI_HANDLE             PartHandle,
  IN OUT LOGICAL_BLOCK_DEVICE  *LogicBlkDev
  )
{
  PART_BLOCK_DEVICE   *PartBlockDev;

  // Validate parameters
  PartBlockDev = (PART_BLOCK_DEVICE *)PartHandle;
  if ( (LogicBlkDev == NULL) || (PartBlockDev == NULL) || \
       (PartBlockDev->Signature != PART_INFO_SIGNATURE) ) {
    return EFI_INVALID_PARAMETER;
  }

  if (SwPart >= PartBlockDev->BlockDeviceCount) {
    return EFI_INVALID_PARAMETER;
  }

  CopyMem (LogicBlkDev, &PartBlockDev->BlockDevice[SwPart], sizeof (LOGICAL_BLOCK_DEVICE));

  return EFI_SUCCESS;
}

/**
Find partitions from OS boot medium

This function will check hardware partition for MBR, GPT or NONE parition.

  @param[in]   HwDevice      The hardware device index.
  @param[out]  PartHandle    The pointer to return parition handle

  @retval EFI_SUCCESS        The partition was found successfully.
  @retval EFI_NOT_FOUND      No partition was found.
  @retval EFI_DEVICE_ERROR   Error occured for device access.
  @retval EFI_INVALID_PARAMETER   Invalid parameters.
  @retval EFI_OUT_OF_RESOURCES    Insufficant memory resource pool.

**/
EFI_STATUS
EFIAPI
FindPartitions (
  IN   UINT32               HwDevice,
  OUT  EFI_HANDLE          *PartHandle
  )
{
  EFI_STATUS                   Status;
  DEVICE_BLOCK_INFO            DevBlockInfo;
  UINT8                       *Buffer;
  LOGICAL_BLOCK_DEVICE        *BlockDev;
  EFI_PARTITION_TABLE_HEADER  *Gpt;
  PART_BLOCK_DEVICE           *PartBlockDev;

  if (PartHandle == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  Status = MediaGetMediaInfo (HwDevice, &DevBlockInfo);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  PartBlockDev = (PART_BLOCK_DEVICE *) AllocateZeroPool (sizeof (PART_BLOCK_DEVICE));
  if (PartBlockDev == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  PartBlockDev->Signature      = PART_INFO_SIGNATURE;
  PartBlockDev->HarewareDevice = HwDevice;

  CopyMem (&PartBlockDev->BlockInfo, &DevBlockInfo, sizeof (DEVICE_BLOCK_INFO));

  if (DevBlockInfo.BlockSize == 1) {
    // It is memory mapped SPI block device.
    Status = FindSpiPartitions (PartBlockDev);
    if (!EFI_ERROR (Status)) {
      *PartHandle = (EFI_HANDLE)PartBlockDev;
    }
    return Status;
  }

  // Read GPT Partition first
  Buffer = (UINT8 *) PartBlockDev->BlockData;
  Status = MediaReadBlocks (HwDevice, 1, DevBlockInfo.BlockSize, Buffer);
  if (!EFI_ERROR (Status)) {
    // Check GPT partition
    Gpt = (EFI_PARTITION_TABLE_HEADER *)Buffer;
    if ((UINT32)Gpt->Header.Signature == 0x20494645) {
      Status = FindGptPartitions (PartBlockDev);
    } else {
      Status = FindMbrPartitions (PartBlockDev);
    }

    // Check result
    if (EFI_ERROR (Status)) {
      // Could not find any partition, so assume no partitions
      BlockDev              = & (PartBlockDev->BlockDevice[0]);
      BlockDev->StartBlock  = 0;
      BlockDev->LastBlock   = DevBlockInfo.BlockNum - 1;
      PartBlockDev->PartitionType    = EnumPartTypeUnknown;
      PartBlockDev->PartitionChecked = TRUE;
      PartBlockDev->BlockDeviceCount = 1;
      Status = EFI_SUCCESS;
    }

    DEBUG ((DEBUG_INFO, "Partition type: %a  (%d logical partitions)\n", \
            GetPartitionTypeName (PartBlockDev->PartitionType), \
            PartBlockDev->BlockDeviceCount));
  } else {
    Status = EFI_DEVICE_ERROR;
  }

  if (EFI_ERROR (Status)) {
    FreePool (PartBlockDev);
  } else {
    *PartHandle = (EFI_HANDLE)PartBlockDev;
  }

  return Status;
}