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