[PATCH v2 1/5] MdePkg/PeCoffGetEntryPointLib: Add PeCoffSerachImageBase()

[Jeff Fan <jeff.fan@intel.com>]

This new API only works on DEBUG build. It will search the PE/COFF image base
forward the input address in this PE/COFF image and returns it.

Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Michael Kinney <michael.d.kinney@intel.com>
Cc: Liming Gao <liming.gao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jeff Fan <jeff.fan@intel.com>
---
 MdePkg/Include/Library/PeCoffGetEntryPointLib.h    | 20 +++++-
 .../PeCoffGetEntryPoint.c                          | 72 +++++++++++++++++++++-
 2 files changed, 90 insertions(+), 2 deletions(-)

diff --git a/MdePkg/Include/Library/PeCoffGetEntryPointLib.h b/MdePkg/Include/Library/PeCoffGetEntryPointLib.h
index e517ca2..647503b 100644
--- a/MdePkg/Include/Library/PeCoffGetEntryPointLib.h
+++ b/MdePkg/Include/Library/PeCoffGetEntryPointLib.h
@@ -1,7 +1,7 @@
 /** @file
   Provides a service to retrieve the PE/COFF entry point from a PE/COFF image.
 
-Copyright (c) 2006 - 2010, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2006 - 2017, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials are licensed and made available under 
 the terms and conditions of the BSD License that accompanies this distribution.  
 The full text of the license may be found at
@@ -101,4 +101,22 @@ PeCoffGetSizeOfHeaders (
   IN VOID     *Pe32Data
   );
 
+/**
+  Returns PE/COFF image base specified by the address in this PE/COFF image.
+
+  On DEBUG build, searches the PE/COFF image base forward the address in this
+  PE/COFF image and returns it.
+
+  @param  Address    Address located in one PE/COFF image.
+
+  @retval 0          RELEASE build or cannot find the PE/COFF image base.
+  @retval others     PE/COFF image base found.
+
+**/
+UINTN
+EFIAPI
+PeCoffSerachImageBase (
+  IN UINTN    Address
+  );
+
 #endif
diff --git a/MdePkg/Library/BasePeCoffGetEntryPointLib/PeCoffGetEntryPoint.c b/MdePkg/Library/BasePeCoffGetEntryPointLib/PeCoffGetEntryPoint.c
index 0fb7e84..00f6d7d 100644
--- a/MdePkg/Library/BasePeCoffGetEntryPointLib/PeCoffGetEntryPoint.c
+++ b/MdePkg/Library/BasePeCoffGetEntryPointLib/PeCoffGetEntryPoint.c
@@ -2,7 +2,7 @@
   Provides the services to get the entry point to a PE/COFF image that has either been 
   loaded into memory or is executing at it's linked address.
 
-  Copyright (c) 2006 - 2010, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2006 - 2017, Intel Corporation. All rights reserved.<BR>
   Portions copyright (c) 2008 - 2009, Apple Inc. All rights reserved.<BR>
   This program and the accompanying materials
   are licensed and made available under the terms and conditions of the BSD License
@@ -22,6 +22,8 @@
 
 #include <IndustryStandard/PeImage.h>
 
+#define PE_COFF_IMAGE_ALIGN_SIZE        4
+
 /**
   Retrieves and returns a pointer to the entry point to a PE/COFF image that has been loaded
   into system memory with the PE/COFF Loader Library functions.
@@ -316,3 +318,71 @@ PeCoffGetSizeOfHeaders (
   return (UINT32) SizeOfHeaders;
 }
 
+/**
+  Returns PE/COFF image base is loaded in system memory where the input address is in.
+
+  On DEBUG build, searches the PE/COFF image base forward the input address and
+  returns it.
+
+  @param  Address    Address located in one PE/COFF image.
+
+  @retval 0          RELEASE build or cannot find the PE/COFF image base.
+  @retval others     PE/COFF image base found.
+
+**/
+UINTN
+EFIAPI
+PeCoffSerachImageBase (
+  IN UINTN    Address
+  )
+{
+  UINTN                                Pe32Data;
+
+  Pe32Data = 0;
+
+  DEBUG_CODE (
+    EFI_IMAGE_DOS_HEADER                 *DosHdr;
+    EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
+
+    //
+    // Find Image Base
+    //
+    Pe32Data = Address & ~(PE_COFF_IMAGE_ALIGN_SIZE - 1);
+    while (Pe32Data != 0) {
+      DosHdr = (EFI_IMAGE_DOS_HEADER *) Pe32Data;
+      if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
+        //
+        // DOS image header is present, so read the PE header after the DOS image header.
+        //
+        Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)(Pe32Data + (UINTN) ((DosHdr->e_lfanew) & 0x0ffff));
+        //
+        // Make sure PE header address does not overflow and is less than the initial address.
+        //
+        if (((UINTN)Hdr.Pe32 > Pe32Data) && ((UINTN)Hdr.Pe32 < Address)) {
+          if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {
+            break;
+          }
+        }
+      } else {
+        //
+        // DOS image header is not present, TE header is at the image base.
+        //
+        Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)Pe32Data;
+        if ((Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) &&
+            ((Hdr.Te->Machine == IMAGE_FILE_MACHINE_I386)  || (Hdr.Te->Machine == IMAGE_FILE_MACHINE_IA64) ||
+             (Hdr.Te->Machine == IMAGE_FILE_MACHINE_EBC)   || (Hdr.Te->Machine == IMAGE_FILE_MACHINE_X64)  ||
+             (Hdr.Te->Machine == IMAGE_FILE_MACHINE_ARM64) || (Hdr.Te->Machine == IMAGE_FILE_MACHINE_ARMTHUMB_MIXED))
+             ) {
+          break;
+        }
+      }
+
+      //
+      // Not found the image base, check the previous aligned address
+      //
+      Pe32Data -= PE_COFF_IMAGE_ALIGN_SIZE;
+    }
+  );
+
+  return Pe32Data;
+}
-- 
2.9.3.windows.2