Enhances PE image hash algorithm in DxeImageVerificationLib and DxeTpmMeasureBootLib.
authortye1 <tye1@6f19259b-4bc3-4df7-8a09-765794883524>
Sat, 28 Apr 2012 07:48:15 +0000 (07:48 +0000)
committertye1 <tye1@6f19259b-4bc3-4df7-8a09-765794883524>
Sat, 28 Apr 2012 07:48:15 +0000 (07:48 +0000)
Signed-off-by: Ye Ting<ting.ye@intel.com>
Reviewed by: Dong, Eric <yong.dong@intel.com>
Reviewed by: Dong, Guo <guo.dong@intel.com>

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@13228 6f19259b-4bc3-4df7-8a09-765794883524

SecurityPkg/Library/DxeImageVerificationLib/DxeImageVerificationLib.c
SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.c

index 14cd2d7..a591970 100644 (file)
@@ -253,6 +253,8 @@ HashPeImage (
   UINTN                     Pos;\r
   UINTN                     SumOfSectionBytes;\r
   EFI_IMAGE_SECTION_HEADER  *SectionCache;\r
+  UINT32                    CertSize;\r
+  UINT32                    NumberOfRvaAndSizes;\r
 \r
   HashCtx       = NULL;\r
   SectionHeader = NULL;\r
@@ -292,6 +294,7 @@ HashPeImage (
   if (!Status) {\r
     goto Done;\r
   }\r
+\r
   //\r
   // Measuring PE/COFF Image Header;\r
   // But CheckSum field and SECURITY data directory (certificate) are excluded\r
@@ -307,11 +310,13 @@ HashPeImage (
     // Use PE32 offset.\r
     //\r
     HashSize = (UINTN) ((UINT8 *) (&mNtHeader.Pe32->OptionalHeader.CheckSum) - HashBase);\r
+    NumberOfRvaAndSizes = mNtHeader.Pe32->OptionalHeader.NumberOfRvaAndSizes;\r
   } else if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC) {\r
     //\r
     // Use PE32+ offset.\r
     //\r
     HashSize = (UINTN) ((UINT8 *) (&mNtHeader.Pe32Plus->OptionalHeader.CheckSum) - HashBase);\r
+    NumberOfRvaAndSizes = mNtHeader.Pe32Plus->OptionalHeader.NumberOfRvaAndSizes;\r
   } else {\r
     //\r
     // Invalid header magic number.\r
@@ -324,51 +329,86 @@ HashPeImage (
   if (!Status) {\r
     goto Done;\r
   }\r
+\r
   //\r
   // 5.  Skip over the image checksum (it occupies a single ULONG).\r
-  // 6.  Get the address of the beginning of the Cert Directory.\r
-  // 7.  Hash everything from the end of the checksum to the start of the Cert Directory.\r
   //\r
-  if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+  if (NumberOfRvaAndSizes <= EFI_IMAGE_DIRECTORY_ENTRY_SECURITY) {\r
     //\r
-    // Use PE32 offset.\r
+    // 6.  Since there is no Cert Directory in optional header, hash everything\r
+    //     from the end of the checksum to the end of image header.\r
     //\r
-    HashBase = (UINT8 *) &mNtHeader.Pe32->OptionalHeader.CheckSum + sizeof (UINT32);\r
-    HashSize = (UINTN) ((UINT8 *) (&mNtHeader.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY]) - HashBase);\r
+    if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+      //\r
+      // Use PE32 offset.\r
+      //\r
+      HashBase = (UINT8 *) &mNtHeader.Pe32->OptionalHeader.CheckSum + sizeof (UINT32);\r
+      HashSize = mNtHeader.Pe32->OptionalHeader.SizeOfHeaders - (UINTN) (HashBase - mImageBase);\r
+    } else {\r
+      //\r
+      // Use PE32+ offset.\r
+      //\r
+      HashBase = (UINT8 *) &mNtHeader.Pe32Plus->OptionalHeader.CheckSum + sizeof (UINT32);\r
+      HashSize = mNtHeader.Pe32Plus->OptionalHeader.SizeOfHeaders - (UINTN) (HashBase - mImageBase);\r
+    }\r
+\r
+    if (HashSize != 0) {\r
+      Status  = mHash[HashAlg].HashUpdate(HashCtx, HashBase, HashSize);\r
+      if (!Status) {\r
+        goto Done;\r
+      }\r
+    }\r
   } else {\r
     //\r
-    // Use PE32+ offset.\r
+    // 7.  Hash everything from the end of the checksum to the start of the Cert Directory.\r
     //\r
-    HashBase = (UINT8 *) &mNtHeader.Pe32Plus->OptionalHeader.CheckSum + sizeof (UINT32);\r
-    HashSize = (UINTN) ((UINT8 *) (&mNtHeader.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY]) - HashBase);\r
-  }\r
+    if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+      //\r
+      // Use PE32 offset.\r
+      //\r
+      HashBase = (UINT8 *) &mNtHeader.Pe32->OptionalHeader.CheckSum + sizeof (UINT32);\r
+      HashSize = (UINTN) ((UINT8 *) (&mNtHeader.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY]) - HashBase);\r
+    } else {\r
+      //\r
+      // Use PE32+ offset.\r
+      //\r
+      HashBase = (UINT8 *) &mNtHeader.Pe32Plus->OptionalHeader.CheckSum + sizeof (UINT32);\r
+      HashSize = (UINTN) ((UINT8 *) (&mNtHeader.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY]) - HashBase);\r
+    }\r
+\r
+    if (HashSize != 0) {\r
+      Status  = mHash[HashAlg].HashUpdate(HashCtx, HashBase, HashSize);\r
+      if (!Status) {\r
+        goto Done;\r
+      }\r
+    }\r
 \r
-  Status  = mHash[HashAlg].HashUpdate(HashCtx, HashBase, HashSize);\r
-  if (!Status) {\r
-    goto Done;\r
-  }\r
-  //\r
-  // 8.  Skip over the Cert Directory. (It is sizeof(IMAGE_DATA_DIRECTORY) bytes.)\r
-  // 9.  Hash everything from the end of the Cert Directory to the end of image header.\r
-  //\r
-  if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
-    //\r
-    // Use PE32 offset\r
-    //\r
-    HashBase = (UINT8 *) &mNtHeader.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1];\r
-    HashSize = mNtHeader.Pe32->OptionalHeader.SizeOfHeaders - (UINTN) ((UINT8 *) (&mNtHeader.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1]) - mImageBase);\r
-  } else {\r
     //\r
-    // Use PE32+ offset.\r
+    // 8.  Skip over the Cert Directory. (It is sizeof(IMAGE_DATA_DIRECTORY) bytes.)\r
+    // 9.  Hash everything from the end of the Cert Directory to the end of image header.\r
     //\r
-    HashBase = (UINT8 *) &mNtHeader.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1];\r
-    HashSize = mNtHeader.Pe32Plus->OptionalHeader.SizeOfHeaders - (UINTN) ((UINT8 *) (&mNtHeader.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1]) - mImageBase);\r
-  }\r
+    if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+      //\r
+      // Use PE32 offset\r
+      //\r
+      HashBase = (UINT8 *) &mNtHeader.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1];\r
+      HashSize = mNtHeader.Pe32->OptionalHeader.SizeOfHeaders - (UINTN) (HashBase - mImageBase);\r
+    } else {\r
+      //\r
+      // Use PE32+ offset.\r
+      //\r
+      HashBase = (UINT8 *) &mNtHeader.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1];\r
+      HashSize = mNtHeader.Pe32Plus->OptionalHeader.SizeOfHeaders - (UINTN) (HashBase - mImageBase);\r
+    }\r
 \r
-  Status  = mHash[HashAlg].HashUpdate(HashCtx, HashBase, HashSize);\r
-  if (!Status) {\r
-    goto Done;\r
+    if (HashSize != 0) {\r
+      Status  = mHash[HashAlg].HashUpdate(HashCtx, HashBase, HashSize);\r
+      if (!Status) {\r
+        goto Done;\r
+      }\r
+    }    \r
   }\r
+\r
   //\r
   // 10. Set the SUM_OF_BYTES_HASHED to the size of the header.\r
   //\r
@@ -398,15 +438,6 @@ HashPeImage (
     SumOfSectionBytes += SectionCache->SizeOfRawData;\r
   }\r
 \r
-  //\r
-  // Sanity check for file corruption. Sections raw data size should be smaller\r
-  // than Image Size.\r
-  //\r
-  if (SumOfSectionBytes >= mImageSize) {\r
-    Status = FALSE;\r
-    goto Done;\r
-  }\r
-\r
   //\r
   // 11. Build a temporary table of pointers to all the IMAGE_SECTION_HEADER\r
   //     structures in the image. The 'NumberOfSections' field of the image\r
@@ -465,37 +496,36 @@ HashPeImage (
   //\r
   if (mImageSize > SumOfBytesHashed) {\r
     HashBase = mImageBase + SumOfBytesHashed;\r
-    if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
-      if (mImageSize - SumOfBytesHashed < mNtHeader.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size) {\r
-        Status = FALSE;\r
-        goto Done;\r
-      }\r
-      //\r
-      // Use PE32 offset.\r
-      //\r
-      HashSize = (UINTN)(\r
-                 mImageSize -\r
-                 mNtHeader.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size -\r
-                 SumOfBytesHashed);\r
+\r
+    if (NumberOfRvaAndSizes <= EFI_IMAGE_DIRECTORY_ENTRY_SECURITY) {\r
+      CertSize = 0;\r
     } else {\r
-      if (mImageSize - SumOfBytesHashed < mNtHeader.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size) {\r
-        Status = FALSE;\r
-        goto Done;\r
+      if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+        //\r
+        // Use PE32 offset.\r
+        //\r
+        CertSize = mNtHeader.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size;\r
+      } else {\r
+        //\r
+        // Use PE32+ offset.\r
+        //\r
+        CertSize = mNtHeader.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size;\r
       }\r
-      //\r
-      // Use PE32+ offset.\r
-      //\r
-      HashSize = (UINTN)(\r
-                 mImageSize -\r
-                 mNtHeader.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size -\r
-                 SumOfBytesHashed);\r
     }\r
 \r
-    Status  = mHash[HashAlg].HashUpdate(HashCtx, HashBase, HashSize);\r
-    if (!Status) {\r
+    if (mImageSize > CertSize + SumOfBytesHashed) {\r
+      HashSize = (UINTN) (mImageSize - CertSize - SumOfBytesHashed);\r
+\r
+      Status  = mHash[HashAlg].HashUpdate(HashCtx, HashBase, HashSize);\r
+      if (!Status) {\r
+        goto Done;\r
+      }\r
+    } else if (mImageSize < CertSize + SumOfBytesHashed) {\r
+      Status = FALSE;\r
       goto Done;\r
     }\r
   }\r
+\r
   Status  = mHash[HashAlg].HashFinal(HashCtx, mImageDigest);\r
 \r
 Done:\r
@@ -1170,19 +1200,20 @@ DxeImageVerificationHandler (
   IN  UINTN                            FileSize\r
   )\r
 {\r
-  EFI_STATUS                  Status;\r
-  UINT16                      Magic;\r
-  EFI_IMAGE_DOS_HEADER        *DosHdr;\r
-  EFI_STATUS                  VerifyStatus;\r
-  UINT8                       *SetupMode;\r
-  EFI_SIGNATURE_LIST          *SignatureList;\r
-  UINTN                       SignatureListSize;\r
-  EFI_SIGNATURE_DATA          *Signature;\r
-  EFI_IMAGE_EXECUTION_ACTION  Action;\r
-  WIN_CERTIFICATE             *WinCertificate;\r
-  UINT32                      Policy;\r
-  UINT8                       *SecureBootEnable;\r
-  PE_COFF_LOADER_IMAGE_CONTEXT      ImageContext;\r
+  EFI_STATUS                           Status;\r
+  UINT16                               Magic;\r
+  EFI_IMAGE_DOS_HEADER                 *DosHdr;\r
+  EFI_STATUS                           VerifyStatus;\r
+  UINT8                                *SetupMode;\r
+  EFI_SIGNATURE_LIST                   *SignatureList;\r
+  UINTN                                SignatureListSize;\r
+  EFI_SIGNATURE_DATA                   *Signature;\r
+  EFI_IMAGE_EXECUTION_ACTION           Action;\r
+  WIN_CERTIFICATE                      *WinCertificate;\r
+  UINT32                               Policy;\r
+  UINT8                                *SecureBootEnable;\r
+  PE_COFF_LOADER_IMAGE_CONTEXT         ImageContext;\r
+  UINT32                               NumberOfRvaAndSizes;\r
 \r
   if (File == NULL) {\r
     return EFI_INVALID_PARAMETER;\r
@@ -1243,6 +1274,8 @@ DxeImageVerificationHandler (
     return EFI_SUCCESS;\r
   }\r
 \r
+  FreePool (SecureBootEnable);\r
+\r
   SetupMode = GetEfiGlobalVariable (EFI_SETUP_MODE_NAME);\r
 \r
   //\r
@@ -1260,13 +1293,16 @@ DxeImageVerificationHandler (
     FreePool (SetupMode);\r
     return EFI_SUCCESS;\r
   }\r
+\r
+  FreePool (SetupMode);\r
+\r
   //\r
   // Read the Dos header.\r
   //\r
   if (FileBuffer == NULL) {\r
-    FreePool (SetupMode);\r
     return EFI_INVALID_PARAMETER;\r
   }\r
+\r
   mImageBase  = (UINT8 *) FileBuffer;\r
   mImageSize  = FileSize;\r
 \r
@@ -1303,7 +1339,8 @@ DxeImageVerificationHandler (
     //\r
     // It is not a valid Pe/Coff file.\r
     //\r
-    return EFI_ACCESS_DENIED;\r
+    Status = EFI_ACCESS_DENIED;\r
+    goto Done;\r
   }\r
 \r
   Magic = mNtHeader.Pe32->OptionalHeader.Magic;\r
@@ -1311,29 +1348,21 @@ DxeImageVerificationHandler (
     //\r
     // Use PE32 offset.\r
     //\r
-    mSecDataDir = (EFI_IMAGE_DATA_DIRECTORY *) &mNtHeader.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY];\r
-  } else if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC) {\r
-    //\r
-    // Use PE32+ offset.\r
-    //\r
-    mSecDataDir = (EFI_IMAGE_DATA_DIRECTORY *) &mNtHeader.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY];\r
+    NumberOfRvaAndSizes = mNtHeader.Pe32->OptionalHeader.NumberOfRvaAndSizes;\r
+    if (NumberOfRvaAndSizes > EFI_IMAGE_DIRECTORY_ENTRY_SECURITY) {\r
+      mSecDataDir = (EFI_IMAGE_DATA_DIRECTORY *) &mNtHeader.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY];\r
+    }        \r
   } else {\r
     //\r
-    // Invalid header magic number.\r
-    //\r
-    Status       = EFI_INVALID_PARAMETER;\r
-    goto Done;\r
-  }\r
-\r
-  if (mSecDataDir->VirtualAddress >= mImageSize) {\r
-    //\r
-    // Sanity check to see if this file is corrupted.\r
+    // Use PE32+ offset.\r
     //\r
-    Status       = EFI_INVALID_PARAMETER;\r
-    goto Done;\r
+    NumberOfRvaAndSizes = mNtHeader.Pe32Plus->OptionalHeader.NumberOfRvaAndSizes;\r
+    if (NumberOfRvaAndSizes > EFI_IMAGE_DIRECTORY_ENTRY_SECURITY) {\r
+      mSecDataDir = (EFI_IMAGE_DATA_DIRECTORY *) &mNtHeader.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY];\r
+    }\r
   }\r
 \r
-  if (mSecDataDir->Size == 0) {\r
+  if ((mSecDataDir == NULL) || ((mSecDataDir != NULL) && (mSecDataDir->Size == 0))) {\r
     //\r
     // This image is not signed.\r
     //\r
@@ -1474,8 +1503,6 @@ Done:
     FreePool (SignatureList);\r
   }\r
 \r
-  FreePool (SetupMode);\r
-\r
   return Status;\r
 }\r
 \r
index 17c41e9..8179422 100644 (file)
@@ -256,7 +256,9 @@ TcgMeasureGptTable (
 \r
   @retval EFI_SUCCESS            Successfully measure image.\r
   @retval EFI_OUT_OF_RESOURCES   No enough resource to measure image.\r
+  @retval EFI_UNSUPPORTED        ImageType is unsupported or PE image is mal-format.  \r
   @retval other error value\r
+\r
 **/\r
 EFI_STATUS\r
 EFIAPI\r
@@ -269,27 +271,31 @@ TcgMeasurePeImage (
   IN  EFI_DEVICE_PATH_PROTOCOL  *FilePath\r
   )\r
 {\r
-  EFI_STATUS                        Status;\r
-  TCG_PCR_EVENT                     *TcgEvent;\r
-  EFI_IMAGE_LOAD_EVENT              *ImageLoad;\r
-  UINT32                            FilePathSize;\r
-  VOID                              *Sha1Ctx;\r
-  UINTN                             CtxSize;\r
-  EFI_IMAGE_DOS_HEADER              *DosHdr;\r
-  UINT32                            PeCoffHeaderOffset;\r
-  EFI_IMAGE_SECTION_HEADER          *Section;\r
-  UINT8                             *HashBase;\r
-  UINTN                             HashSize;\r
-  UINTN                             SumOfBytesHashed;\r
-  EFI_IMAGE_SECTION_HEADER          *SectionHeader;\r
-  UINTN                             Index, Pos;\r
-  UINT16                            Magic;\r
-  UINT32                            EventSize;\r
-  UINT32                            EventNumber;\r
-  EFI_PHYSICAL_ADDRESS              EventLogLastEntry;\r
+  EFI_STATUS                           Status;\r
+  TCG_PCR_EVENT                        *TcgEvent;\r
+  EFI_IMAGE_LOAD_EVENT                 *ImageLoad;\r
+  UINT32                               FilePathSize;\r
+  VOID                                 *Sha1Ctx;\r
+  UINTN                                CtxSize;\r
+  EFI_IMAGE_DOS_HEADER                 *DosHdr;\r
+  UINT32                               PeCoffHeaderOffset;\r
+  EFI_IMAGE_SECTION_HEADER             *Section;\r
+  UINT8                                *HashBase;\r
+  UINTN                                HashSize;\r
+  UINTN                                SumOfBytesHashed;\r
+  EFI_IMAGE_SECTION_HEADER             *SectionHeader;\r
+  UINTN                                Index;\r
+  UINTN                                Pos;\r
+  UINT16                               Magic;\r
+  UINT32                               EventSize;\r
+  UINT32                               EventNumber;\r
+  EFI_PHYSICAL_ADDRESS                 EventLogLastEntry;\r
   EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;\r
+  UINT32                               NumberOfRvaAndSizes;\r
+  BOOLEAN                              HashStatus;\r
+  UINT32                               CertSize;\r
 \r
-  Status        = EFI_SUCCESS;\r
+  Status        = EFI_UNSUPPORTED;\r
   ImageLoad     = NULL;\r
   SectionHeader = NULL;\r
   Sha1Ctx       = NULL;\r
@@ -326,7 +332,6 @@ TcgMeasurePeImage (
         "TcgMeasurePeImage: Unknown subsystem type %d",\r
         ImageType\r
         ));\r
-      Status = EFI_UNSUPPORTED;\r
       goto Finish;\r
   }\r
 \r
@@ -344,8 +349,9 @@ TcgMeasurePeImage (
   if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {\r
     PeCoffHeaderOffset = DosHdr->e_lfanew;\r
   }\r
-  if (((EFI_TE_IMAGE_HEADER *)((UINT8 *) (UINTN) ImageAddress + PeCoffHeaderOffset))->Signature\r
-       == EFI_TE_IMAGE_HEADER_SIGNATURE) {\r
+\r
+  Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINT8 *) (UINTN) ImageAddress + PeCoffHeaderOffset);\r
+  if (Hdr.Pe32->Signature != EFI_IMAGE_NT_SIGNATURE) {\r
     goto Finish;\r
   }\r
 \r
@@ -367,14 +373,16 @@ TcgMeasurePeImage (
     goto Finish;\r
   }\r
 \r
-  Sha1Init (Sha1Ctx);\r
+  HashStatus = Sha1Init (Sha1Ctx);\r
+  if (!HashStatus) {\r
+    goto Finish;\r
+  }\r
 \r
   //\r
   // Measuring PE/COFF Image Header;\r
   // But CheckSum field and SECURITY data directory (certificate) are excluded\r
   //\r
-  Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINT8 *) (UINTN) ImageAddress + PeCoffHeaderOffset);\r
-  Magic    = Hdr.Pe32->OptionalHeader.Magic;\r
+  Magic = Hdr.Pe32->OptionalHeader.Magic;\r
   \r
   //\r
   // 3.  Calculate the distance from the base of the image header to the image checksum address.\r
@@ -385,58 +393,99 @@ TcgMeasurePeImage (
     //\r
     // Use PE32 offset\r
     //\r
+    NumberOfRvaAndSizes = Hdr.Pe32->OptionalHeader.NumberOfRvaAndSizes;\r
     HashSize = (UINTN) ((UINT8 *)(&Hdr.Pe32->OptionalHeader.CheckSum) - HashBase);\r
   } else {\r
     //\r
     // Use PE32+ offset\r
     //\r
+    NumberOfRvaAndSizes = Hdr.Pe32Plus->OptionalHeader.NumberOfRvaAndSizes;\r
     HashSize = (UINTN) ((UINT8 *)(&Hdr.Pe32Plus->OptionalHeader.CheckSum) - HashBase);\r
   }\r
 \r
-  Sha1Update (Sha1Ctx, HashBase, HashSize);\r
+  HashStatus = Sha1Update (Sha1Ctx, HashBase, HashSize);\r
+  if (!HashStatus) {\r
+    goto Finish;\r
+  }  \r
 \r
   //\r
   // 5.  Skip over the image checksum (it occupies a single ULONG).\r
-  // 6.  Get the address of the beginning of the Cert Directory.\r
-  // 7.  Hash everything from the end of the checksum to the start of the Cert Directory.\r
   //\r
-  if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+  if (NumberOfRvaAndSizes <= EFI_IMAGE_DIRECTORY_ENTRY_SECURITY) {\r
     //\r
-    // Use PE32 offset\r
+    // 6.  Since there is no Cert Directory in optional header, hash everything\r
+    //     from the end of the checksum to the end of image header.\r
     //\r
-    HashBase = (UINT8 *) &Hdr.Pe32->OptionalHeader.CheckSum + sizeof (UINT32);\r
-    HashSize = (UINTN) ((UINT8 *)(&Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY]) - HashBase);\r
+    if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+      //\r
+      // Use PE32 offset.\r
+      //\r
+      HashBase = (UINT8 *) &Hdr.Pe32->OptionalHeader.CheckSum + sizeof (UINT32);\r
+      HashSize = Hdr.Pe32->OptionalHeader.SizeOfHeaders - (UINTN) (HashBase - ImageAddress);\r
+    } else {\r
+      //\r
+      // Use PE32+ offset.\r
+      //\r
+      HashBase = (UINT8 *) &Hdr.Pe32Plus->OptionalHeader.CheckSum + sizeof (UINT32);\r
+      HashSize = Hdr.Pe32Plus->OptionalHeader.SizeOfHeaders - (UINTN) (HashBase - ImageAddress);\r
+    }\r
+\r
+    if (HashSize != 0) {\r
+      HashStatus  = Sha1Update (Sha1Ctx, HashBase, HashSize);\r
+      if (!HashStatus) {\r
+        goto Finish;\r
+      }\r
+    }    \r
   } else {\r
     //\r
-    // Use PE32+ offset\r
-    //    \r
-    HashBase = (UINT8 *) &Hdr.Pe32Plus->OptionalHeader.CheckSum + sizeof (UINT32);\r
-    HashSize = (UINTN) ((UINT8 *)(&Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY]) - HashBase);\r
-  }\r
+    // 7.  Hash everything from the end of the checksum to the start of the Cert Directory.\r
+    //\r
+    if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+      //\r
+      // Use PE32 offset\r
+      //\r
+      HashBase = (UINT8 *) &Hdr.Pe32->OptionalHeader.CheckSum + sizeof (UINT32);\r
+      HashSize = (UINTN) ((UINT8 *)(&Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY]) - HashBase);\r
+    } else {\r
+      //\r
+      // Use PE32+ offset\r
+      //    \r
+      HashBase = (UINT8 *) &Hdr.Pe32Plus->OptionalHeader.CheckSum + sizeof (UINT32);\r
+      HashSize = (UINTN) ((UINT8 *)(&Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY]) - HashBase);\r
+    }\r
 \r
-  Sha1Update (Sha1Ctx, HashBase, HashSize);\r
+    if (HashSize != 0) {\r
+      HashStatus  = Sha1Update (Sha1Ctx, HashBase, HashSize);\r
+      if (!HashStatus) {\r
+        goto Finish;\r
+      }\r
+    }\r
 \r
-  //\r
-  // 8.  Skip over the Cert Directory. (It is sizeof(IMAGE_DATA_DIRECTORY) bytes.)\r
-  // 9.  Hash everything from the end of the Cert Directory to the end of image header.\r
-  //\r
-  if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
     //\r
-    // Use PE32 offset\r
-    //\r
-    HashBase = (UINT8 *) &Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1];\r
-    HashSize = Hdr.Pe32->OptionalHeader.SizeOfHeaders -\r
-             (UINTN) ((UINT8 *)(&Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1]) - (UINT8 *) (UINTN) ImageAddress);\r
-  } else {\r
+    // 8.  Skip over the Cert Directory. (It is sizeof(IMAGE_DATA_DIRECTORY) bytes.)\r
+    // 9.  Hash everything from the end of the Cert Directory to the end of image header.\r
     //\r
-    // Use PE32+ offset\r
-    //\r
-    HashBase = (UINT8 *) &Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1];\r
-    HashSize = Hdr.Pe32Plus->OptionalHeader.SizeOfHeaders -\r
-             (UINTN) ((UINT8 *)(&Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1]) - (UINT8 *) (UINTN) ImageAddress);\r
+    if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+      //\r
+      // Use PE32 offset\r
+      //\r
+      HashBase = (UINT8 *) &Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1];\r
+      HashSize = Hdr.Pe32->OptionalHeader.SizeOfHeaders - (UINTN) (HashBase - ImageAddress);\r
+    } else {\r
+      //\r
+      // Use PE32+ offset\r
+      //\r
+      HashBase = (UINT8 *) &Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1];\r
+      HashSize = Hdr.Pe32Plus->OptionalHeader.SizeOfHeaders - (UINTN) (HashBase - ImageAddress);\r
+    }\r
+    \r
+    if (HashSize != 0) {\r
+      HashStatus  = Sha1Update (Sha1Ctx, HashBase, HashSize);\r
+      if (!HashStatus) {\r
+        goto Finish;\r
+      }\r
+    }\r
   }\r
-  \r
-  Sha1Update (Sha1Ctx, HashBase, HashSize);\r
 \r
   //\r
   // 10. Set the SUM_OF_BYTES_HASHED to the size of the header\r
@@ -459,7 +508,7 @@ TcgMeasurePeImage (
   //     header indicates how big the table should be. Do not include any\r
   //     IMAGE_SECTION_HEADERs in the table whose 'SizeOfRawData' field is zero.\r
   //\r
-  SectionHeader = (EFI_IMAGE_SECTION_HEADER *)AllocateZeroPool (sizeof (EFI_IMAGE_SECTION_HEADER) * Hdr.Pe32->FileHeader.NumberOfSections);\r
+  SectionHeader = (EFI_IMAGE_SECTION_HEADER *) AllocateZeroPool (sizeof (EFI_IMAGE_SECTION_HEADER) * Hdr.Pe32->FileHeader.NumberOfSections);\r
   if (SectionHeader == NULL) {\r
     Status = EFI_OUT_OF_RESOURCES;\r
     goto Finish;\r
@@ -503,7 +552,10 @@ TcgMeasurePeImage (
     HashBase = (UINT8 *) (UINTN) ImageAddress + Section->PointerToRawData;\r
     HashSize = (UINTN) Section->SizeOfRawData;\r
 \r
-    Sha1Update (Sha1Ctx, HashBase, HashSize);\r
+    HashStatus = Sha1Update (Sha1Ctx, HashBase, HashSize);\r
+    if (!HashStatus) {\r
+      goto Finish;\r
+    }\r
 \r
     SumOfBytesHashed += HashSize;\r
   }\r
@@ -516,37 +568,42 @@ TcgMeasurePeImage (
   //\r
   if (ImageSize > SumOfBytesHashed) {\r
     HashBase = (UINT8 *) (UINTN) ImageAddress + SumOfBytesHashed;\r
-    if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
-      if (ImageSize - SumOfBytesHashed < Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size) {\r
-        Status = EFI_INVALID_PARAMETER;\r
-        goto Finish;\r
-      }\r
-      //\r
-      // Use PE32 offset\r
-      //\r
-      HashSize = (UINTN)(ImageSize -\r
-                 Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size -\r
-                 SumOfBytesHashed);\r
+\r
+    if (NumberOfRvaAndSizes <= EFI_IMAGE_DIRECTORY_ENTRY_SECURITY) {\r
+      CertSize = 0;\r
     } else {\r
-      if (ImageSize - SumOfBytesHashed < Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size) {\r
-        Status = EFI_INVALID_PARAMETER;\r
-        goto Finish;\r
+      if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+        //\r
+        // Use PE32 offset.\r
+        //\r
+        CertSize = Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size;\r
+      } else {\r
+        //\r
+        // Use PE32+ offset.\r
+        //\r
+        CertSize = Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size;\r
       }\r
-      //\r
-      // Use PE32+ offset\r
-      //\r
-      HashSize = (UINTN)(ImageSize -\r
-                 Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size -\r
-                 SumOfBytesHashed);      \r
     }\r
 \r
-    Sha1Update (Sha1Ctx, HashBase, HashSize);\r
+    if (ImageSize > CertSize + SumOfBytesHashed) {\r
+      HashSize = (UINTN) (ImageSize - CertSize - SumOfBytesHashed);\r
+\r
+      HashStatus = Sha1Update (Sha1Ctx, HashBase, HashSize);\r
+      if (!HashStatus) {\r
+        goto Finish;\r
+      }\r
+    } else if (ImageSize < CertSize + SumOfBytesHashed) {\r
+      goto Finish;\r
+    }\r
   }\r
 \r
   //\r
   // 17.  Finalize the SHA hash.\r
   //\r
-  Sha1Final (Sha1Ctx, (UINT8 *)&TcgEvent->Digest);\r
+  HashStatus = Sha1Final (Sha1Ctx, (UINT8 *) &TcgEvent->Digest);\r
+  if (!HashStatus) {\r
+    goto Finish;\r
+  }\r
 \r
   //\r
   // Log the PE data\r