X-Git-Url: https://git.proxmox.com/?p=mirror_edk2.git;a=blobdiff_plain;f=MdeModulePkg%2FUniversal%2FDisk%2FPartitionDxe%2FGpt.c;h=35860515c13b9da3284b2a5f0d5c57cad2032016;hp=7dcb9b45b2f1fadb6068cf92fc7e84dc38d21b82;hb=493d8e3a5e86f78ef975e18dfa5754adce9e82de;hpb=8a7d75b0625cffee0c67b85afe56763f93d86481

diff --git a/MdeModulePkg/Universal/Disk/PartitionDxe/Gpt.c b/MdeModulePkg/Universal/Disk/PartitionDxe/Gpt.c
index 7dcb9b45b2..35860515c1 100644
--- a/MdeModulePkg/Universal/Disk/PartitionDxe/Gpt.c
+++ b/MdeModulePkg/Universal/Disk/PartitionDxe/Gpt.c
@@ -2,21 +2,47 @@
   Decode a hard disk partitioned with the GPT scheme in the UEFI 2.0
   specification.
 
-  Copyright (c) 2006 - 2007, Intel Corporation                                              
-  All rights reserved. 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                                            
+  Caution: This file requires additional review when modified.
+  This driver will have external input - disk partition.
+  This external input must be validated carefully to avoid security issue like
+  buffer overflow, integer overflow.
 
-  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,                     
-  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.             
+  PartitionInstallGptChildHandles() routine will read disk partition content and
+  do basic validation before PartitionInstallChildHandle().
+
+  PartitionValidGptTable(), PartitionCheckGptEntry() routine will accept disk
+  partition content and validate the GPT table and GPT entry.
+
+Copyright (c) 2006 - 2013, 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 "Partition.h"
 
+/**
+  Install child handles if the Handle supports GPT partition structure.
+
+  Caution: This function may receive untrusted input.
+  The GPT partition table header is external input, so this routine
+  will do basic validation for GPT partition table header before return.
+
+  @param[in]  BlockIo     Parent BlockIo interface.
+  @param[in]  DiskIo      Disk Io protocol.
+  @param[in]  Lba         The starting Lba of the Partition Table
+  @param[out] PartHeader  Stores the partition table that is read
 
+  @retval TRUE      The partition table is valid
+  @retval FALSE     The partition table is not valid
+
+**/
 BOOLEAN
 PartitionValidGptTable (
   IN  EFI_BLOCK_IO_PROTOCOL       *BlockIo,
@@ -25,7 +51,18 @@ PartitionValidGptTable (
   OUT EFI_PARTITION_TABLE_HEADER  *PartHeader
   );
 
+/**
+  Check if the CRC field in the Partition table header is valid
+  for Partition entry array.
+
+  @param[in]  BlockIo     Parent BlockIo interface
+  @param[in]  DiskIo      Disk Io Protocol.
+  @param[in]  PartHeader  Partition table header structure
+
+  @retval TRUE      the CRC is valid
+  @retval FALSE     the CRC is invalid
 
+**/
 BOOLEAN
 PartitionCheckGptEntryArrayCRC (
   IN  EFI_BLOCK_IO_PROTOCOL       *BlockIo,
@@ -34,6 +71,18 @@ PartitionCheckGptEntryArrayCRC (
   );
 
 
+/**
+  Restore Partition Table to its alternate place
+  (Primary -> Backup or Backup -> Primary).
+
+  @param[in]  BlockIo     Parent BlockIo interface.
+  @param[in]  DiskIo      Disk Io Protocol.
+  @param[in]  PartHeader  Partition table header structure.
+
+  @retval TRUE      Restoring succeeds
+  @retval FALSE     Restoring failed
+
+**/
 BOOLEAN
 PartitionRestoreGptTable (
   IN  EFI_BLOCK_IO_PROTOCOL       *BlockIo,
@@ -42,6 +91,19 @@ PartitionRestoreGptTable (
   );
 
 
+/**
+  This routine will check GPT partition entry and return entry status.
+
+  Caution: This function may receive untrusted input.
+  The GPT partition entry is external input, so this routine
+  will do basic validation for GPT partition entry and report status.
+
+  @param[in]    PartHeader    Partition table header structure
+  @param[in]    PartEntry     The partition entry array
+  @param[out]   PEntryStatus  the partition entry status array 
+                              recording the status of each partition
+
+**/
 VOID
 PartitionCheckGptEntry (
   IN  EFI_PARTITION_TABLE_HEADER  *PartHeader,
@@ -50,6 +112,17 @@ PartitionCheckGptEntry (
   );
 
 
+/**
+  Checks the CRC32 value in the table header.
+
+  @param  MaxSize   Max Size limit
+  @param  Size      The size of the table
+  @param  Hdr       Table to check
+
+  @return TRUE    CRC Valid
+  @return FALSE   CRC Invalid
+
+**/
 BOOLEAN
 PartitionCheckCrcAltSize (
   IN UINTN                 MaxSize,
@@ -58,6 +131,16 @@ PartitionCheckCrcAltSize (
   );
 
 
+/**
+  Checks the CRC32 value in the table header.
+
+  @param  MaxSize   Max Size limit
+  @param  Hdr       Table to check
+
+  @return TRUE      CRC Valid
+  @return FALSE     CRC Invalid
+
+**/
 BOOLEAN
 PartitionCheckCrc (
   IN UINTN                 MaxSize,
@@ -65,6 +148,13 @@ PartitionCheckCrc (
   );
 
 
+/**
+  Updates the CRC32 value in the table header.
+
+  @param  Size   The size of the table
+  @param  Hdr    Table to update
+
+**/
 VOID
 PartitionSetCrcAltSize (
   IN UINTN                 Size,
@@ -72,6 +162,12 @@ PartitionSetCrcAltSize (
   );
 
 
+/**
+  Updates the CRC32 value in the table header.
+
+  @param  Hdr    Table to update
+
+**/
 VOID
 PartitionSetCrc (
   IN OUT EFI_TABLE_HEADER *Hdr
@@ -80,15 +176,22 @@ PartitionSetCrc (
 /**
   Install child handles if the Handle supports GPT partition structure.
 
-  @param[in]  This       - Calling context.
-  @param[in]  Handle     - Parent Handle
-  @param[in]  DiskIo     - Parent DiskIo interface
-  @param[in]  BlockIo    - Parent BlockIo interface
-  @param[in]  DevicePath - Parent Device Path
+  Caution: This function may receive untrusted input.
+  The GPT partition table is external input, so this routine
+  will do basic validation for GPT partition table before install
+  child handle for each GPT partition.
 
-  @retval EFI_SUCCESS         Valid GPT disk
-  @retval EFI_MEDIA_CHANGED   Media changed Detected
-  @retval other               Not a valid GPT disk
+  @param[in]  This       Calling context.
+  @param[in]  Handle     Parent Handle.
+  @param[in]  DiskIo     Parent DiskIo interface.
+  @param[in]  DiskIo2    Parent DiskIo2 interface.
+  @param[in]  BlockIo    Parent BlockIo interface.
+  @param[in]  BlockIo2   Parent BlockIo2 interface.
+  @param[in]  DevicePath Parent Device Path.
+
+  @retval EFI_SUCCESS           Valid GPT disk.
+  @retval EFI_MEDIA_CHANGED     Media changed Detected.
+  @retval other                 Not a valid GPT disk.
 
 **/
 EFI_STATUS
@@ -96,7 +199,9 @@ PartitionInstallGptChildHandles (
   IN  EFI_DRIVER_BINDING_PROTOCOL  *This,
   IN  EFI_HANDLE                   Handle,
   IN  EFI_DISK_IO_PROTOCOL         *DiskIo,
+  IN  EFI_DISK_IO2_PROTOCOL        *DiskIo2,
   IN  EFI_BLOCK_IO_PROTOCOL        *BlockIo,
+  IN  EFI_BLOCK_IO2_PROTOCOL       *BlockIo2,
   IN  EFI_DEVICE_PATH_PROTOCOL     *DevicePath
   )
 {
@@ -107,10 +212,12 @@ PartitionInstallGptChildHandles (
   EFI_PARTITION_TABLE_HEADER  *PrimaryHeader;
   EFI_PARTITION_TABLE_HEADER  *BackupHeader;
   EFI_PARTITION_ENTRY         *PartEntry;
+  EFI_PARTITION_ENTRY         *Entry;
   EFI_PARTITION_ENTRY_STATUS  *PEntryStatus;
   UINTN                       Index;
-  EFI_STATUS                  GptValid;
+  EFI_STATUS                  GptValidStatus;
   HARDDRIVE_DEVICE_PATH       HdDev;
+  UINT32                      MediaId;
 
   ProtectiveMbr = NULL;
   PrimaryHeader = NULL;
@@ -120,11 +227,12 @@ PartitionInstallGptChildHandles (
 
   BlockSize     = BlockIo->Media->BlockSize;
   LastBlock     = BlockIo->Media->LastBlock;
+  MediaId       = BlockIo->Media->MediaId;
 
   DEBUG ((EFI_D_INFO, " BlockSize : %d \n", BlockSize));
-  DEBUG ((EFI_D_INFO, " LastBlock : %x \n", LastBlock));
+  DEBUG ((EFI_D_INFO, " LastBlock : %lx \n", LastBlock));
 
-  GptValid = EFI_NOT_FOUND;
+  GptValidStatus = EFI_NOT_FOUND;
 
   //
   // Allocate a buffer for the Protective MBR
@@ -137,24 +245,30 @@ PartitionInstallGptChildHandles (
   //
   // Read the Protective MBR from LBA #0
   //
-  Status = BlockIo->ReadBlocks (
-                      BlockIo,
-                      BlockIo->Media->MediaId,
-                      0,
-                      BlockIo->Media->BlockSize,
-                      ProtectiveMbr
-                      );
+  Status = DiskIo->ReadDisk (
+                     DiskIo,
+                     MediaId,
+                     0,
+                     BlockSize,
+                     ProtectiveMbr
+                     );
   if (EFI_ERROR (Status)) {
-    GptValid = Status;
+    GptValidStatus = Status;
     goto Done;
   }
+
   //
   // Verify that the Protective MBR is valid
   //
-  if (ProtectiveMbr->Partition[0].BootIndicator != 0x00 ||
-      ProtectiveMbr->Partition[0].OSIndicator != PMBR_GPT_PARTITION ||
-      UNPACK_UINT32 (ProtectiveMbr->Partition[0].StartingLBA) != 1
-      ) {
+  for (Index = 0; Index < MAX_MBR_PARTITIONS; Index++) {
+    if (ProtectiveMbr->Partition[Index].BootIndicator == 0x00 &&
+        ProtectiveMbr->Partition[Index].OSIndicator == PMBR_GPT_PARTITION &&
+        UNPACK_UINT32 (ProtectiveMbr->Partition[Index].StartingLBA) == 1
+        ) {
+      break;
+    }
+  }
+  if (Index == MAX_MBR_PARTITIONS) {
     goto Done;
   }
 
@@ -167,7 +281,6 @@ PartitionInstallGptChildHandles (
   }
 
   BackupHeader = AllocateZeroPool (sizeof (EFI_PARTITION_TABLE_HEADER));
-
   if (BackupHeader == NULL) {
     goto Done;
   }
@@ -196,7 +309,7 @@ PartitionInstallGptChildHandles (
     DEBUG ((EFI_D_INFO, " Valid primary and !Valid backup partition table\n"));
     DEBUG ((EFI_D_INFO, " Restore backup partition table by the primary\n"));
     if (!PartitionRestoreGptTable (BlockIo, DiskIo, PrimaryHeader)) {
-      DEBUG ((EFI_D_INFO, " Restore  backup partition table error\n"));
+      DEBUG ((EFI_D_INFO, " Restore backup partition table error\n"));
     }
 
     if (PartitionValidGptTable (BlockIo, DiskIo, PrimaryHeader->AlternateLBA, BackupHeader)) {
@@ -210,22 +323,22 @@ PartitionInstallGptChildHandles (
   //
   // Read the EFI Partition Entries
   //
-  PartEntry = AllocatePool (PrimaryHeader->NumberOfPartitionEntries * sizeof (EFI_PARTITION_ENTRY));
+  PartEntry = AllocatePool (PrimaryHeader->NumberOfPartitionEntries * PrimaryHeader->SizeOfPartitionEntry);
   if (PartEntry == NULL) {
     DEBUG ((EFI_D_ERROR, "Allocate pool error\n"));
     goto Done;
   }
 
   Status = DiskIo->ReadDisk (
-                    DiskIo,
-                    BlockIo->Media->MediaId,
-                    MultU64x32(PrimaryHeader->PartitionEntryLBA, BlockSize),
-                    PrimaryHeader->NumberOfPartitionEntries * (PrimaryHeader->SizeOfPartitionEntry),
-                    PartEntry
-                    );
+                     DiskIo,
+                     MediaId,
+                     MultU64x32(PrimaryHeader->PartitionEntryLBA, BlockSize),
+                     PrimaryHeader->NumberOfPartitionEntries * (PrimaryHeader->SizeOfPartitionEntry),
+                     PartEntry
+                     );
   if (EFI_ERROR (Status)) {
-    GptValid = Status;
-    DEBUG ((EFI_D_INFO, " Partition Entry ReadBlocks error\n"));
+    GptValidStatus = Status;
+    DEBUG ((EFI_D_ERROR, " Partition Entry ReadDisk error\n"));
     goto Done;
   }
 
@@ -247,18 +360,21 @@ PartitionInstallGptChildHandles (
   //
   // If we got this far the GPT layout of the disk is valid and we should return true
   //
-  GptValid = EFI_SUCCESS;
+  GptValidStatus = EFI_SUCCESS;
 
   //
   // Create child device handles
   //
   for (Index = 0; Index < PrimaryHeader->NumberOfPartitionEntries; Index++) {
-    if (CompareGuid (&PartEntry[Index].PartitionTypeGUID, &gEfiPartTypeUnusedGuid) ||
+    Entry = (EFI_PARTITION_ENTRY *) ((UINT8 *) PartEntry + Index * PrimaryHeader->SizeOfPartitionEntry);
+    if (CompareGuid (&Entry->PartitionTypeGUID, &gEfiPartTypeUnusedGuid) ||
         PEntryStatus[Index].OutOfRange ||
-        PEntryStatus[Index].Overlap
+        PEntryStatus[Index].Overlap ||
+        PEntryStatus[Index].OsSpecific
         ) {
       //
-      // Don't use null EFI Partition Entries or Invalid Partition Entries
+      // Don't use null EFI Partition Entries, Invalid Partition Entries or OS specific
+      // partition Entries
       //
       continue;
     }
@@ -271,29 +387,31 @@ PartitionInstallGptChildHandles (
     HdDev.PartitionNumber = (UINT32) Index + 1;
     HdDev.MBRType         = MBR_TYPE_EFI_PARTITION_TABLE_HEADER;
     HdDev.SignatureType   = SIGNATURE_TYPE_GUID;
-    HdDev.PartitionStart  = PartEntry[Index].StartingLBA;
-    HdDev.PartitionSize   = PartEntry[Index].EndingLBA - PartEntry[Index].StartingLBA + 1;
-    CopyMem (HdDev.Signature, &PartEntry[Index].UniquePartitionGUID, sizeof (EFI_GUID));
+    HdDev.PartitionStart  = Entry->StartingLBA;
+    HdDev.PartitionSize   = Entry->EndingLBA - Entry->StartingLBA + 1;
+    CopyMem (HdDev.Signature, &Entry->UniquePartitionGUID, sizeof (EFI_GUID));
 
-    DEBUG ((EFI_D_INFO, " Index : %d\n", Index));
-    DEBUG ((EFI_D_INFO, " Start LBA : %x\n", HdDev.PartitionStart));
-    DEBUG ((EFI_D_INFO, " End LBA : %x\n", PartEntry[Index].EndingLBA));
-    DEBUG ((EFI_D_INFO, " Partition size: %x\n", HdDev.PartitionSize));
-    DEBUG ((EFI_D_INFO, " Start : %x", MultU64x32 (PartEntry[Index].StartingLBA, BlockSize)));
-    DEBUG ((EFI_D_INFO, " End : %x\n", MultU64x32 (PartEntry[Index].EndingLBA, BlockSize)));
+    DEBUG ((EFI_D_INFO, " Index : %d\n", (UINT32) Index));
+    DEBUG ((EFI_D_INFO, " Start LBA : %lx\n", (UINT64) HdDev.PartitionStart));
+    DEBUG ((EFI_D_INFO, " End LBA : %lx\n", (UINT64) Entry->EndingLBA));
+    DEBUG ((EFI_D_INFO, " Partition size: %lx\n", (UINT64) HdDev.PartitionSize));
+    DEBUG ((EFI_D_INFO, " Start : %lx", MultU64x32 (Entry->StartingLBA, BlockSize)));
+    DEBUG ((EFI_D_INFO, " End : %lx\n", MultU64x32 (Entry->EndingLBA, BlockSize)));
 
     Status = PartitionInstallChildHandle (
-              This,
-              Handle,
-              DiskIo,
-              BlockIo,
-              DevicePath,
-              (EFI_DEVICE_PATH_PROTOCOL *) &HdDev,
-              PartEntry[Index].StartingLBA,
-              PartEntry[Index].EndingLBA,
-              BlockSize,
-              CompareGuid(&PartEntry[Index].PartitionTypeGUID, &gEfiPartTypeSystemPartGuid)
-              );
+               This,
+               Handle,
+               DiskIo,
+               DiskIo2,
+               BlockIo,
+               BlockIo2,
+               DevicePath,
+               (EFI_DEVICE_PATH_PROTOCOL *) &HdDev,
+               Entry->StartingLBA,
+               Entry->EndingLBA,
+               BlockSize,
+               CompareGuid(&Entry->PartitionTypeGUID, &gEfiPartTypeSystemPartGuid)
+               );
   }
 
   DEBUG ((EFI_D_INFO, "Prepare to Free Pool\n"));
@@ -315,17 +433,20 @@ Done:
     FreePool (PEntryStatus);
   }
 
-  return GptValid;
+  return GptValidStatus;
 }
 
-
 /**
-  Install child handles if the Handle supports GPT partition structure.
+  This routine will read GPT partition table header and return it.
 
-  @param[in]  BlockIo     Parent BlockIo interface
+  Caution: This function may receive untrusted input.
+  The GPT partition table header is external input, so this routine
+  will do basic validation for GPT partition table header before return.
+
+  @param[in]  BlockIo     Parent BlockIo interface.
   @param[in]  DiskIo      Disk Io protocol.
   @param[in]  Lba         The starting Lba of the Partition Table
-  @param[in]  PartHeader  Stores the partition table that is read
+  @param[out] PartHeader  Stores the partition table that is read
 
   @retval TRUE      The partition table is valid
   @retval FALSE     The partition table is not valid
@@ -342,9 +463,10 @@ PartitionValidGptTable (
   EFI_STATUS                  Status;
   UINT32                      BlockSize;
   EFI_PARTITION_TABLE_HEADER  *PartHdr;
+  UINT32                      MediaId;
 
   BlockSize = BlockIo->Media->BlockSize;
-
+  MediaId   = BlockIo->Media->MediaId;
   PartHdr   = AllocateZeroPool (BlockSize);
 
   if (PartHdr == NULL) {
@@ -354,13 +476,13 @@ PartitionValidGptTable (
   //
   // Read the EFI Partition Table Header
   //
-  Status = BlockIo->ReadBlocks (
-                      BlockIo,
-                      BlockIo->Media->MediaId,
-                      Lba,
-                      BlockSize,
-                      PartHdr
-                      );
+  Status = DiskIo->ReadDisk (
+                     DiskIo,
+                     MediaId,
+                     MultU64x32 (Lba, BlockSize),
+                     BlockSize,
+                     PartHdr
+                     );
   if (EFI_ERROR (Status)) {
     FreePool (PartHdr);
     return FALSE;
@@ -368,9 +490,18 @@ PartitionValidGptTable (
 
   if ((PartHdr->Header.Signature != EFI_PTAB_HEADER_ID) ||
       !PartitionCheckCrc (BlockSize, &PartHdr->Header) ||
-      PartHdr->MyLBA != Lba
+      PartHdr->MyLBA != Lba ||
+      (PartHdr->SizeOfPartitionEntry < sizeof (EFI_PARTITION_ENTRY))
       ) {
-    DEBUG ((EFI_D_INFO, " !Valid efi partition table header\n"));
+    DEBUG ((EFI_D_INFO, "Invalid efi partition table header\n"));
+    FreePool (PartHdr);
+    return FALSE;
+  }
+
+  //
+  // Ensure the NumberOfPartitionEntries * SizeOfPartitionEntry doesn't overflow.
+  //
+  if (PartHdr->NumberOfPartitionEntries > DivU64x32 (MAX_UINTN, PartHdr->SizeOfPartitionEntry)) {
     FreePool (PartHdr);
     return FALSE;
   }
@@ -386,7 +517,6 @@ PartitionValidGptTable (
   return TRUE;
 }
 
-
 /**
   Check if the CRC field in the Partition table header is valid
   for Partition entry array.
@@ -449,11 +579,11 @@ PartitionCheckGptEntryArrayCRC (
 
 /**
   Restore Partition Table to its alternate place
-  (Primary -> Backup or Backup -> Primary)
+  (Primary -> Backup or Backup -> Primary).
 
-  @param[in]  BlockIo     Parent BlockIo interface
+  @param[in]  BlockIo     Parent BlockIo interface.
   @param[in]  DiskIo      Disk Io Protocol.
-  @param[in]  PartHeader  Partition table header structure
+  @param[in]  PartHeader  Partition table header structure.
 
   @retval TRUE      Restoring succeeds
   @retval FALSE     Restoring failed
@@ -471,11 +601,13 @@ PartitionRestoreGptTable (
   EFI_PARTITION_TABLE_HEADER  *PartHdr;
   EFI_LBA                     PEntryLBA;
   UINT8                       *Ptr;
+  UINT32                      MediaId;
 
   PartHdr   = NULL;
   Ptr       = NULL;
 
   BlockSize = BlockIo->Media->BlockSize;
+  MediaId   = BlockIo->Media->MediaId;
 
   PartHdr   = AllocateZeroPool (BlockSize);
 
@@ -495,22 +627,28 @@ PartitionRestoreGptTable (
   PartHdr->PartitionEntryLBA  = PEntryLBA;
   PartitionSetCrc ((EFI_TABLE_HEADER *) PartHdr);
 
-  Status = BlockIo->WriteBlocks (BlockIo, BlockIo->Media->MediaId, PartHdr->MyLBA, BlockSize, PartHdr);
+  Status = DiskIo->WriteDisk (
+                     DiskIo,
+                     MediaId,
+                     MultU64x32 (PartHdr->MyLBA, (UINT32) BlockSize),
+                     BlockSize,
+                     PartHdr
+                     );
   if (EFI_ERROR (Status)) {
     goto Done;
   }
 
   Ptr = AllocatePool (PartHeader->NumberOfPartitionEntries * PartHeader->SizeOfPartitionEntry);
   if (Ptr == NULL) {
-    DEBUG ((EFI_D_ERROR, " Allocate pool effor\n"));
+    DEBUG ((EFI_D_ERROR, " Allocate pool error\n"));
     Status = EFI_OUT_OF_RESOURCES;
     goto Done;
   }
 
   Status = DiskIo->ReadDisk (
                     DiskIo,
-                    BlockIo->Media->MediaId,
-                    MultU64x32(PartHeader->PartitionEntryLBA, BlockIo->Media->BlockSize),
+                    MediaId,
+                    MultU64x32(PartHeader->PartitionEntryLBA, (UINT32) BlockSize),
                     PartHeader->NumberOfPartitionEntries * PartHeader->SizeOfPartitionEntry,
                     Ptr
                     );
@@ -520,15 +658,18 @@ PartitionRestoreGptTable (
 
   Status = DiskIo->WriteDisk (
                     DiskIo,
-                    BlockIo->Media->MediaId,
-                    MultU64x32(PEntryLBA, BlockIo->Media->BlockSize),
+                    MediaId,
+                    MultU64x32(PEntryLBA, (UINT32) BlockSize),
                     PartHeader->NumberOfPartitionEntries * PartHeader->SizeOfPartitionEntry,
                     Ptr
                     );
 
 Done:
   FreePool (PartHdr);
-  FreePool (Ptr);
+
+  if (Ptr != NULL) {
+    FreePool (Ptr);
+  }
 
   if (EFI_ERROR (Status)) {
     return FALSE;
@@ -537,15 +678,18 @@ Done:
   return TRUE;
 }
 
-
 /**
-  Restore Partition Table to its alternate place
-  (Primary -> Backup or Backup -> Primary)
+  This routine will check GPT partition entry and return entry status.
+
+  Caution: This function may receive untrusted input.
+  The GPT partition entry is external input, so this routine
+  will do basic validation for GPT partition entry and report status.
 
   @param[in]    PartHeader    Partition table header structure
   @param[in]    PartEntry     The partition entry array
   @param[out]   PEntryStatus  the partition entry status array 
                               recording the status of each partition
+
 **/
 VOID
 PartitionCheckGptEntry (
@@ -554,19 +698,21 @@ PartitionCheckGptEntry (
   OUT EFI_PARTITION_ENTRY_STATUS  *PEntryStatus
   )
 {
-  EFI_LBA StartingLBA;
-  EFI_LBA EndingLBA;
-  UINTN   Index1;
-  UINTN   Index2;
+  EFI_LBA              StartingLBA;
+  EFI_LBA              EndingLBA;
+  EFI_PARTITION_ENTRY  *Entry;
+  UINTN                Index1;
+  UINTN                Index2;
 
   DEBUG ((EFI_D_INFO, " start check partition entries\n"));
   for (Index1 = 0; Index1 < PartHeader->NumberOfPartitionEntries; Index1++) {
-    if (CompareGuid (&PartEntry[Index1].PartitionTypeGUID, &gEfiPartTypeUnusedGuid)) {
+    Entry = (EFI_PARTITION_ENTRY *) ((UINT8 *) PartEntry + Index1 * PartHeader->SizeOfPartitionEntry);
+    if (CompareGuid (&Entry->PartitionTypeGUID, &gEfiPartTypeUnusedGuid)) {
       continue;
     }
 
-    StartingLBA = PartEntry[Index1].StartingLBA;
-    EndingLBA   = PartEntry[Index1].EndingLBA;
+    StartingLBA = Entry->StartingLBA;
+    EndingLBA   = Entry->EndingLBA;
     if (StartingLBA > EndingLBA ||
         StartingLBA < PartHeader->FirstUsableLBA ||
         StartingLBA > PartHeader->LastUsableLBA ||
@@ -577,20 +723,26 @@ PartitionCheckGptEntry (
       continue;
     }
 
-    for (Index2 = Index1 + 1; Index2 < PartHeader->NumberOfPartitionEntries; Index2++) {
+    if ((Entry->Attributes & BIT1) != 0) {
+      //
+      // If Bit 1 is set, this indicate that this is an OS specific GUID partition. 
+      //
+      PEntryStatus[Index1].OsSpecific = TRUE;
+    }
 
-      if (CompareGuid (&PartEntry[Index2].PartitionTypeGUID, &gEfiPartTypeUnusedGuid)) {
+    for (Index2 = Index1 + 1; Index2 < PartHeader->NumberOfPartitionEntries; Index2++) {
+      Entry = (EFI_PARTITION_ENTRY *) ((UINT8 *) PartEntry + Index2 * PartHeader->SizeOfPartitionEntry);
+      if (CompareGuid (&Entry->PartitionTypeGUID, &gEfiPartTypeUnusedGuid)) {
         continue;
       }
 
-      if (PartEntry[Index2].EndingLBA >= StartingLBA && PartEntry[Index2].StartingLBA <= EndingLBA) {
+      if (Entry->EndingLBA >= StartingLBA && Entry->StartingLBA <= EndingLBA) {
         //
         // This region overlaps with the Index1'th region
         //
         PEntryStatus[Index1].Overlap  = TRUE;
         PEntryStatus[Index2].Overlap  = TRUE;
         continue;
-
       }
     }
   }
@@ -600,9 +752,9 @@ PartitionCheckGptEntry (
 
 
 /**
-  Updates the CRC32 value in the table header
+  Updates the CRC32 value in the table header.
 
-  @param[in,out]  Hdr    Table to update
+  @param  Hdr    Table to update
 
 **/
 VOID
@@ -615,10 +767,10 @@ PartitionSetCrc (
 
 
 /**
-  Updates the CRC32 value in the table header
+  Updates the CRC32 value in the table header.
 
-  @param[in]      Size   The size of the table
-  @param[in,out]  Hdr    Table to update
+  @param  Size   The size of the table
+  @param  Hdr    Table to update
 
 **/
 VOID
@@ -626,7 +778,6 @@ PartitionSetCrcAltSize (
   IN UINTN                 Size,
   IN OUT EFI_TABLE_HEADER  *Hdr
   )
-
 {
   UINT32  Crc;
 
@@ -637,13 +788,13 @@ PartitionSetCrcAltSize (
 
 
 /**
-  Checks the CRC32 value in the table header
+  Checks the CRC32 value in the table header.
 
-  @param[in]      MaxSize   Max Size limit
-  @param[in,out]  Hdr       Table to check
+  @param  MaxSize   Max Size limit
+  @param  Hdr       Table to check
 
-  @return TRUE    CRC Valid
-  @return FALSE   CRC Invalid
+  @return TRUE      CRC Valid
+  @return FALSE     CRC Invalid
 
 **/
 BOOLEAN
@@ -657,11 +808,11 @@ PartitionCheckCrc (
 
 
 /**
-  Checks the CRC32 value in the table header
+  Checks the CRC32 value in the table header.
 
-  @param[in]      MaxSize   Max Size limit
-  @param[in]      Size      The size of the table
-  @param[in,out]  Hdr       Table to check
+  @param  MaxSize   Max Size limit
+  @param  Size      The size of the table
+  @param  Hdr       Table to check
 
   @return TRUE    CRC Valid
   @return FALSE   CRC Invalid
@@ -687,7 +838,7 @@ PartitionCheckCrcAltSize (
     return FALSE;
   }
 
-  if (MaxSize && Size > MaxSize) {
+  if ((MaxSize != 0) && (Size > MaxSize)) {
     DEBUG ((EFI_D_ERROR, "CheckCrc32: Size > MaxSize\n"));
     return FALSE;
   }