Validate some fields in PE image to make sure not access violation for later code.

Signed-off-by: Eric Dong <eric.dong@intel.com>
Reviewed-by: Liming Gao <liming.gao@intel.com>
git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@13211 6f19259b-4bc3-4df7-8a09-765794883524

diff --git a/MdeModulePkg/Core/Dxe/Image/Image.c b/MdeModulePkg/Core/Dxe/Image/Image.c
index abafa222e7e63d3010c6f9471bf38b46219249fc..e51a9fe1743b8391e2d7bbf5c0341f84ffec7e4a 100644 (file)
--- a/MdeModulePkg/Core/Dxe/Image/Image.c
+++ b/MdeModulePkg/Core/Dxe/Image/Image.c
@@ -1,7 +1,7 @@
 /** @file\r
   Core image handling services to load and unload PeImage.\r
 \r
-Copyright (c) 2006 - 2011, Intel Corporation. All rights reserved.<BR>\r
+Copyright (c) 2006 - 2012, 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
@@ -232,6 +232,14 @@ CoreReadImageFile (
   UINTN               EndPosition;\r
   IMAGE_FILE_HANDLE  *FHand;\r
 \r
+  if (UserHandle == NULL || ReadSize == NULL || Buffer == NULL) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  if (MAX_ADDRESS - Offset < *ReadSize) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
   FHand = (IMAGE_FILE_HANDLE  *)UserHandle;\r
   ASSERT (FHand->Signature == IMAGE_FILE_HANDLE_SIGNATURE);\r
 \r

diff --git a/MdePkg/Library/BasePeCoffLib/BasePeCoff.c b/MdePkg/Library/BasePeCoffLib/BasePeCoff.c
index 8b2a78e448dda99a2ce839c4397296325e8f09a3..6342a665df329875aa2d1f983fec681974f7e826 100644 (file)
--- a/MdePkg/Library/BasePeCoffLib/BasePeCoff.c
+++ b/MdePkg/Library/BasePeCoffLib/BasePeCoff.c
@@ -47,7 +47,9 @@ PeCoffLoaderGetPeHeaderMagicValue (
 \r
 \r
 /**\r
-  Retrieves the PE or TE Header from a PE/COFF or TE image.\r
+  Retrieves the PE or TE Header from a PE/COFF or TE image. \r
+  Also done many checks in PE image to make sure PE image DosHeader, PeOptionHeader, \r
+  SizeOfHeader, Section Data Region and Security Data Region be in PE image range. \r
 \r
   @param  ImageContext    The context of the image being loaded.\r
   @param  Hdr             The buffer in which to return the PE32, PE32+, or TE header.\r
@@ -66,6 +68,10 @@ PeCoffLoaderGetPeHeader (
   EFI_IMAGE_DOS_HEADER  DosHdr;\r
   UINTN                 Size;\r
   UINT16                Magic;\r
+  UINT32                SectionHeaderOffset;\r
+  UINT32                Index;\r
+  CHAR8                 BufferData;\r
+  EFI_IMAGE_SECTION_HEADER  SectionHeader;\r
 \r
   //\r
   // Read the DOS image header to check for its existence\r
@@ -131,6 +137,74 @@ PeCoffLoaderGetPeHeader (
     Magic = PeCoffLoaderGetPeHeaderMagicValue (Hdr);\r
 \r
     if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+      //\r
+      // 1. Check FileHeader.SizeOfOptionalHeader filed.\r
+      //\r
+      if (EFI_IMAGE_NUMBER_OF_DIRECTORY_ENTRIES < Hdr.Pe32->OptionalHeader.NumberOfRvaAndSizes) {\r
+        return RETURN_UNSUPPORTED;\r
+      }\r
+\r
+      if (Hdr.Pe32->FileHeader.SizeOfOptionalHeader != sizeof (EFI_IMAGE_OPTIONAL_HEADER32) - (EFI_IMAGE_NUMBER_OF_DIRECTORY_ENTRIES - Hdr.Pe32->OptionalHeader.NumberOfRvaAndSizes) * sizeof (EFI_IMAGE_DATA_DIRECTORY)) {\r
+        return RETURN_UNSUPPORTED;\r
+      }\r
+\r
+      //\r
+      // 2. Check the OptionalHeader.SizeOfHeaders field.\r
+      // This field will be use like the following mode, so just compare the result.\r
+      // The DataDirectory array begin with 1, not 0, so here use < to compare not <=.\r
+      //\r
+      if (EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1 < Hdr.Pe32->OptionalHeader.NumberOfRvaAndSizes) {\r
+        if (Hdr.Pe32->OptionalHeader.SizeOfHeaders < (UINT32)((UINT8 *)(&Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1]) - (UINT8 *) &Hdr)) {\r
+          return RETURN_UNSUPPORTED;\r
+        }\r
+      }\r
+\r
+      //\r
+      // Read Hdr.Pe32.OptionalHeader.SizeOfHeaders data from file\r
+      //\r
+      Size = 1;\r
+      Status = ImageContext->ImageRead (\r
+                               ImageContext->Handle,\r
+                               Hdr.Pe32->OptionalHeader.SizeOfHeaders - 1,\r
+                               &Size,\r
+                               &BufferData\r
+                               );\r
+      if (RETURN_ERROR (Status)) {\r
+        return Status;\r
+      }\r
+\r
+      //\r
+      // Check the EFI_IMAGE_DIRECTORY_ENTRY_SECURITY data.\r
+      // Read the last byte to make sure the data is in the image region.\r
+      // The DataDirectory array begin with 1, not 0, so here use < to compare not <=.\r
+      //\r
+      if (EFI_IMAGE_DIRECTORY_ENTRY_SECURITY < Hdr.Pe32->OptionalHeader.NumberOfRvaAndSizes) {\r
+        if (Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size != 0) {\r
+          //\r
+          // Check the member data to avoid overflow.\r
+          //\r
+          if ((UINT32) (~0) - Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].VirtualAddress <\r
+              Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size) {\r
+            return RETURN_INVALID_PARAMETER;\r
+          }\r
+\r
+          //\r
+          // Read section header from file\r
+          //\r
+          Size = 1;\r
+          Status = ImageContext->ImageRead (\r
+                                   ImageContext->Handle,\r
+                                   Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].VirtualAddress +\r
+                                    Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size - 1,\r
+                                   &Size,\r
+                                   &BufferData\r
+                                   );\r
+          if (RETURN_ERROR (Status)) {\r
+            return Status;\r
+          }\r
+        }\r
+      }\r
+\r
       //\r
       // Use PE32 offset\r
       //\r
@@ -140,6 +214,74 @@ PeCoffLoaderGetPeHeader (
       ImageContext->SizeOfHeaders     = Hdr.Pe32->OptionalHeader.SizeOfHeaders;\r
 \r
     } else if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC) {\r
+      //\r
+      // 1. Check FileHeader.SizeOfOptionalHeader filed.\r
+      //\r
+      if (EFI_IMAGE_NUMBER_OF_DIRECTORY_ENTRIES < Hdr.Pe32Plus->OptionalHeader.NumberOfRvaAndSizes) {\r
+        return RETURN_UNSUPPORTED;\r
+      }\r
+\r
+      if (Hdr.Pe32Plus->FileHeader.SizeOfOptionalHeader != sizeof (EFI_IMAGE_OPTIONAL_HEADER32) - (EFI_IMAGE_NUMBER_OF_DIRECTORY_ENTRIES - Hdr.Pe32Plus->OptionalHeader.NumberOfRvaAndSizes) * sizeof (EFI_IMAGE_DATA_DIRECTORY)) {\r
+        return RETURN_UNSUPPORTED;\r
+      }\r
+\r
+      //\r
+      // 2. Check the OptionalHeader.SizeOfHeaders field.\r
+      // This field will be use like the following mode, so just compare the result.\r
+      // The DataDirectory array begin with 1, not 0, so here use < to compare not <=.\r
+      //\r
+      if (EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1 < Hdr.Pe32Plus->OptionalHeader.NumberOfRvaAndSizes) {\r
+        if (Hdr.Pe32Plus->OptionalHeader.SizeOfHeaders < (UINT32)((UINT8 *)(&Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY + 1]) - (UINT8 *) &Hdr)) {\r
+          return RETURN_UNSUPPORTED;\r
+        }\r
+      }\r
+\r
+      //\r
+      // Read Hdr.Pe32.OptionalHeader.SizeOfHeaders data from file\r
+      //\r
+      Size = 1;\r
+      Status = ImageContext->ImageRead (\r
+                               ImageContext->Handle,\r
+                               Hdr.Pe32Plus->OptionalHeader.SizeOfHeaders - 1,\r
+                               &Size,\r
+                               &BufferData\r
+                               );\r
+      if (RETURN_ERROR (Status)) {\r
+        return Status;\r
+      }\r
+\r
+      //\r
+      // Check the EFI_IMAGE_DIRECTORY_ENTRY_SECURITY data.\r
+      // Read the last byte to make sure the data is in the image region.\r
+      // The DataDirectory array begin with 1, not 0, so here use < to compare not <=.\r
+      //\r
+      if (EFI_IMAGE_DIRECTORY_ENTRY_SECURITY < Hdr.Pe32Plus->OptionalHeader.NumberOfRvaAndSizes) {\r
+        if (Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size != 0) {\r
+          //\r
+          // Check the member data to avoid overflow.\r
+          //\r
+          if ((UINT32) (~0) - Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].VirtualAddress <\r
+              Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size) {\r
+            return RETURN_INVALID_PARAMETER;\r
+          }\r
+\r
+          //\r
+          // Read section header from file\r
+          //\r
+          Size = 1;\r
+          Status = ImageContext->ImageRead (\r
+                                   ImageContext->Handle,\r
+                                   Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].VirtualAddress +\r
+                                    Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size - 1,\r
+                                   &Size,\r
+                                   &BufferData\r
+                                   );\r
+          if (RETURN_ERROR (Status)) {\r
+            return Status;\r
+          }\r
+        }\r
+      }\r
+\r
       //\r
       // Use PE32+ offset\r
       //\r
@@ -166,6 +308,55 @@ PeCoffLoaderGetPeHeader (
     return RETURN_UNSUPPORTED;\r
   }\r
 \r
+  //\r
+  // Check each section field.\r
+  //\r
+  SectionHeaderOffset = ImageContext->PeCoffHeaderOffset + sizeof (UINT32) + sizeof (EFI_IMAGE_FILE_HEADER) + Hdr.Pe32->FileHeader.SizeOfOptionalHeader;\r
+  for (Index = 0; Index < Hdr.Pe32->FileHeader.NumberOfSections; Index++) {\r
+    //\r
+    // Read section header from file\r
+    //\r
+    Size = sizeof (EFI_IMAGE_SECTION_HEADER);\r
+    Status = ImageContext->ImageRead (\r
+                             ImageContext->Handle,\r
+                             SectionHeaderOffset,\r
+                             &Size,\r
+                             &SectionHeader\r
+                             );\r
+    if (RETURN_ERROR (Status)) {\r
+      return Status;\r
+    }\r
+\r
+    if (SectionHeader.SizeOfRawData > 0) {\r
+      //\r
+      // Check the member data to avoid overflow.\r
+      //\r
+      if ((UINT32) (~0) - SectionHeader.PointerToRawData < SectionHeader.SizeOfRawData) {\r
+        return RETURN_INVALID_PARAMETER;\r
+      }\r
+\r
+      //\r
+      // Base on the ImageRead function to check the section data field.\r
+      // Read the last byte to make sure the data is in the image region.\r
+      //\r
+      Size = 1;\r
+      Status = ImageContext->ImageRead (\r
+                               ImageContext->Handle,\r
+                               SectionHeader.PointerToRawData + SectionHeader.SizeOfRawData - 1,\r
+                               &Size,\r
+                               &BufferData\r
+                               );\r
+      if (RETURN_ERROR (Status)) {\r
+        return Status;\r
+      }\r
+    }\r
+\r
+    //\r
+    // Check next section.\r
+    //\r
+    SectionHeaderOffset += sizeof (EFI_IMAGE_SECTION_HEADER);\r
+  }\r
+\r
   return RETURN_SUCCESS;\r
 }\r
 \r
@@ -185,6 +376,9 @@ PeCoffLoaderGetPeHeader (
   The ImageRead and Handle fields of ImageContext structure must be valid prior \r
   to invoking this service.\r
 \r
+  Also done many checks in PE image to make sure PE image DosHeader, PeOptionHeader, \r
+  SizeOfHeader, Section Data Region and Security Data Region be in PE image range. \r
+\r
   @param  ImageContext              The pointer to the image context structure that describes the PE/COFF\r
                                     image that needs to be examined by this function.\r
 \r

diff --git a/SecurityPkg/Library/DxeImageVerificationLib/DxeImageVerificationLib.c b/SecurityPkg/Library/DxeImageVerificationLib/DxeImageVerificationLib.c
index 19852dd483241b32f622327f3f5c57a2b4dc4644..6e3e9eea95200a7a5b73ea230ffc39de3bbae490 100644 (file)
--- a/SecurityPkg/Library/DxeImageVerificationLib/DxeImageVerificationLib.c
+++ b/SecurityPkg/Library/DxeImageVerificationLib/DxeImageVerificationLib.c
@@ -54,6 +54,50 @@ HASH_TABLE mHash[] = {
   { L"SHA512", 64, &mHashOidValue[40], 9, NULL,                NULL,       NULL,          NULL       }\r
 };\r
 \r
+/**\r
+  Reads contents of a PE/COFF image in memory buffer.\r
+\r
+  @param  FileHandle      Pointer to the file handle to read the PE/COFF image.\r
+  @param  FileOffset      Offset into the PE/COFF image to begin the read operation.\r
+  @param  ReadSize        On input, the size in bytes of the requested read operation.  \r
+                          On output, the number of bytes actually read.\r
+  @param  Buffer          Output buffer that contains the data read from the PE/COFF image.\r
+  \r
+  @retval EFI_SUCCESS     The specified portion of the PE/COFF image was read and the size \r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ImageRead (\r
+  IN     VOID    *FileHandle,\r
+  IN     UINTN   FileOffset,\r
+  IN OUT UINTN   *ReadSize,\r
+  OUT    VOID    *Buffer\r
+  )\r
+{\r
+  UINTN               EndPosition;\r
+\r
+  if (FileHandle == NULL || ReadSize == NULL || Buffer == NULL) {\r
+    return EFI_INVALID_PARAMETER;    \r
+  }\r
+\r
+  if (MAX_ADDRESS - FileOffset < *ReadSize) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  EndPosition = FileOffset + *ReadSize;\r
+  if (EndPosition > mImageSize) {\r
+    *ReadSize = (UINT32)(mImageSize - FileOffset);\r
+  }\r
+\r
+  if (FileOffset >= mImageSize) {\r
+    *ReadSize = 0;\r
+  }\r
+\r
+  CopyMem (Buffer, (UINT8 *)((UINTN) FileHandle + FileOffset), *ReadSize);\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
 \r
 /**\r
   Get the image type.\r
@@ -422,6 +466,10 @@ HashPeImage (
   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
@@ -430,6 +478,10 @@ HashPeImage (
                  mNtHeader.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size -\r
                  SumOfBytesHashed);\r
     } else {\r
+      if (mImageSize - SumOfBytesHashed < mNtHeader.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size) {\r
+        Status = FALSE;\r
+        goto Done;\r
+      }\r
       //\r
       // Use PE32+ offset.\r
       //\r
@@ -1130,6 +1182,7 @@ DxeImageVerificationHandler (
   WIN_CERTIFICATE             *WinCertificate;\r
   UINT32                      Policy;\r
   UINT8                       *SecureBootEnable;\r
+  PE_COFF_LOADER_IMAGE_CONTEXT      ImageContext;\r
 \r
   if (File == NULL) {\r
     return EFI_INVALID_PARAMETER;\r
@@ -1216,6 +1269,22 @@ DxeImageVerificationHandler (
   }\r
   mImageBase  = (UINT8 *) FileBuffer;\r
   mImageSize  = FileSize;\r
+\r
+  ZeroMem (&ImageContext, sizeof (ImageContext));\r
+  ImageContext.Handle    = (VOID *) FileBuffer;\r
+  ImageContext.ImageRead = (PE_COFF_LOADER_READ_FILE) ImageRead;\r
+\r
+  //\r
+  // Get information about the image being loaded\r
+  //\r
+  Status = PeCoffLoaderGetImageInfo (&ImageContext);\r
+  if (EFI_ERROR (Status)) {\r
+    //\r
+    // The information can't be got from the invalid PeImage\r
+    //\r
+    goto Done;\r
+  }\r
+\r
   DosHdr      = (EFI_IMAGE_DOS_HEADER *) mImageBase;\r
   if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {\r
     //\r

diff --git a/SecurityPkg/Library/DxeImageVerificationLib/DxeImageVerificationLib.h b/SecurityPkg/Library/DxeImageVerificationLib/DxeImageVerificationLib.h
index 881e4e5c35861b2100869dbc8da05d7a279fab5a..55371e90bfc1773521ddb27d8cf31c5f6f1e5280 100644 (file)
--- a/SecurityPkg/Library/DxeImageVerificationLib/DxeImageVerificationLib.h
+++ b/SecurityPkg/Library/DxeImageVerificationLib/DxeImageVerificationLib.h
@@ -28,6 +28,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Library/PcdLib.h>\r
 #include <Library/DevicePathLib.h>\r
 #include <Library/SecurityManagementLib.h>\r
+#include <Library/PeCoffLib.h>\r
 #include <Protocol/FirmwareVolume2.h>\r
 #include <Protocol/DevicePath.h>\r
 #include <Protocol/BlockIo.h>\r

diff --git a/SecurityPkg/Library/DxeImageVerificationLib/DxeImageVerificationLib.inf b/SecurityPkg/Library/DxeImageVerificationLib/DxeImageVerificationLib.inf
index 1dda6774fa031a6d6f4233f626dc6d7579d59149..860d64ba83fec2f771eee10c0dbd8d6799f1ea72 100644 (file)
--- a/SecurityPkg/Library/DxeImageVerificationLib/DxeImageVerificationLib.inf
+++ b/SecurityPkg/Library/DxeImageVerificationLib/DxeImageVerificationLib.inf
@@ -2,7 +2,7 @@
 #  The library instance provides security service of image verification.\r
 #  Image verification Library module supports UEFI2.3.1\r
 #\r
-# Copyright (c) 2009 - 2011, Intel Corporation. All rights reserved.<BR>\r
+# Copyright (c) 2009 - 2012, 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
@@ -48,6 +48,7 @@
   DevicePathLib\r
   BaseCryptLib\r
   SecurityManagementLib\r
+  PeCoffLib\r
 \r
 [Protocols]\r
   gEfiFirmwareVolume2ProtocolGuid\r

diff --git a/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.c b/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.c
index d3c7bfec627cc9a970e55b40e6dd7e8bbd9ae835..f0039c80483af7c1e56ca2958b1ba94ce75c6d6f 100644 (file)
--- a/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.c
+++ b/SecurityPkg/Library/DxeTpmMeasureBootLib/DxeTpmMeasureBootLib.c
@@ -36,6 +36,8 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 BOOLEAN                           mMeasureGptTableFlag = FALSE;\r
 EFI_GUID                          mZeroGuid = {0, 0, 0, {0, 0, 0, 0, 0, 0, 0, 0}};\r
 UINTN                             mMeasureGptCount = 0;\r
+VOID                              *mFileBuffer;\r
+UINTN                             mImageSize;\r
 \r
 /**\r
   Reads contents of a PE/COFF image in memory buffer.\r
@@ -57,7 +59,27 @@ ImageRead (
   OUT    VOID    *Buffer\r
   )\r
 {\r
+  UINTN               EndPosition;\r
+\r
+  if (FileHandle == NULL || ReadSize == NULL || Buffer == NULL) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  if (MAX_ADDRESS - FileOffset < *ReadSize) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  EndPosition = FileOffset + *ReadSize;\r
+  if (EndPosition > mImageSize) {\r
+    *ReadSize = (UINT32)(mImageSize - FileOffset);\r
+  }\r
+\r
+  if (FileOffset >= mImageSize) {\r
+    *ReadSize = 0;\r
+  }\r
+\r
   CopyMem (Buffer, (UINT8 *)((UINTN) FileHandle + FileOffset), *ReadSize);\r
+\r
   return EFI_SUCCESS;\r
 }\r
 \r
@@ -495,6 +517,10 @@ TcgMeasurePeImage (
   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
@@ -502,6 +528,10 @@ TcgMeasurePeImage (
                  Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_SECURITY].Size -\r
                  SumOfBytesHashed);\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
+      }\r
       //\r
       // Use PE32+ offset\r
       //\r
@@ -735,6 +765,9 @@ DxeTpmMeasureBootHandler (
     goto Finish;\r
   }\r
 \r
+  mImageSize  = FileSize;\r
+  mFileBuffer = FileBuffer;\r
+\r
   //\r
   // Measure PE Image\r
   //\r