[PATCH 4/9] UefiCpuPkg/CpuExceptionHandlerLib: Add DumpModuleImageInfo()

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

Add internal DumpModuleImageInfo() to replace FindModuleImageBase(). It will
consume PeCoffGetEntrypointLib's PeCoffSerachImageBase() to get PE/COFF image
base.

Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Michael Kinney <michael.d.kinney@intel.com>
Cc: Feng Tian <feng.tian@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jeff Fan <jeff.fan@intel.com>
---
 .../CpuExceptionHandlerLib/CpuExceptionCommon.c    | 75 ++++++----------------
 .../CpuExceptionHandlerLib/CpuExceptionCommon.h    | 11 +---
 .../Library/CpuExceptionHandlerLib/DxeException.c  |  7 +-
 .../Ia32/ArchExceptionHandler.c                    | 15 +----
 .../CpuExceptionHandlerLib/PeiCpuException.c       |  6 +-
 .../CpuExceptionHandlerLib/SecPeiCpuException.c    |  4 --
 .../Library/CpuExceptionHandlerLib/SmmException.c  |  7 +-
 .../X64/ArchExceptionHandler.c                     | 12 +---
 8 files changed, 30 insertions(+), 107 deletions(-)

diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/CpuExceptionCommon.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/CpuExceptionCommon.c
index 3d85b0c..6080d1e 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/CpuExceptionCommon.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/CpuExceptionCommon.c
@@ -113,75 +113,40 @@ InternalPrintMessage (
   @return !0     Image base address.
   @return 0      Image header cannot be found.
 **/
-UINTN
-FindModuleImageBase (
-  IN  UINTN              CurrentEip,
-  OUT UINTN              *EntryPoint
+VOID
+DumpModuleImageInfo (
+  IN  UINTN              CurrentEip
   )
 {
+  EFI_STATUS                           Status;
   UINTN                                Pe32Data;
-  EFI_IMAGE_DOS_HEADER                 *DosHdr;
-  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
   VOID                                 *PdbPointer;
+  VOID                                 *EntryPoint;
 
-  //
-  // Find Image Base
-  //
-  Pe32Data = CurrentEip & ~(mImageAlignSize - 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 < CurrentEip)) {
-        if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {
-          //
-          // It's PE image.
-          //
-          InternalPrintMessage ("!!!! Find PE image ");
-          *EntryPoint = (UINTN)Pe32Data + (UINTN)(Hdr.Pe32->OptionalHeader.AddressOfEntryPoint & 0x0ffffffff);
-          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_X64)) {
-        //
-        // It's TE image, it TE header and Machine type match
-        //
-        InternalPrintMessage ("!!!! Find TE image ");
-        *EntryPoint = (UINTN)Pe32Data + (UINTN)(Hdr.Te->AddressOfEntryPoint & 0x0ffffffff) + sizeof(EFI_TE_IMAGE_HEADER) - Hdr.Te->StrippedSize;
-        break;
-      }
-    }
-
+  Pe32Data = PeCoffSerachImageBase (CurrentEip);
+  if (Pe32Data == 0) {
+    InternalPrintMessage ("!!!! Can't find image information. !!!!\n");
+  } else {
     //
-    // Not found the image base, check the previous aligned address
+    // Find Image Base entry point
     //
-    Pe32Data -= mImageAlignSize;
-  }
-
-  if (Pe32Data != 0) {
+    Status = PeCoffLoaderGetEntryPoint ((VOID *) Pe32Data, &EntryPoint);
+    if (EFI_ERROR (Status)) {
+      EntryPoint = NULL;
+    }
+    InternalPrintMessage ("!!!! Find image ");
     PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *) Pe32Data);
     if (PdbPointer != NULL) {
       InternalPrintMessage ("%a", PdbPointer);
     } else {
       InternalPrintMessage ("(No PDB) " );
     }
-  } else {
-    InternalPrintMessage ("!!!! Can't find image information. !!!!\n");
+    InternalPrintMessage (
+      " (ImageBase=%016lp, EntryPoint=%016p) !!!!\n",
+      (VOID *) Pe32Data,
+      EntryPoint
+      );
   }
-
-  return Pe32Data;
 }
 
 /**
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/CpuExceptionCommon.h b/UefiCpuPkg/Library/CpuExceptionHandlerLib/CpuExceptionCommon.h
index 9adb6a1..0047ad6 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/CpuExceptionCommon.h
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/CpuExceptionCommon.h
@@ -53,7 +53,6 @@ typedef struct {
 } EXCEPTION_HANDLER_DATA;
 
 extern CONST UINT32                mErrorCodeFlag;
-extern CONST UINTN                 mImageAlignSize;
 extern CONST UINTN                 mDoFarReturnFlag;
 
 /**
@@ -112,15 +111,11 @@ InternalPrintMessage (
   Find and display image base address and return image base and its entry point.
   
   @param CurrentEip      Current instruction pointer.
-  @param EntryPoint      Return module entry point if module header is found.
   
-  @return !0     Image base address.
-  @return 0      Image header cannot be found.
 **/
-UINTN 
-FindModuleImageBase (
-  IN  UINTN              CurrentEip,
-  OUT UINTN              *EntryPoint
+VOID 
+DumpModuleImageInfo (
+  IN  UINTN              CurrentEip
   );
 
 /**
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
index a61a52b..ab13e5e 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
@@ -1,7 +1,7 @@
 /** @file
   CPU exception handler library implemenation for DXE modules.
 
-  Copyright (c) 2013 - 2016, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2013 - 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
   which accompanies this distribution.  The full text of the license may be found at
@@ -19,11 +19,6 @@
 
 CONST UINTN    mDoFarReturnFlag  = 0;
 
-//
-// Image align size for DXE/SMM
-//
-CONST UINTN      mImageAlignSize = SIZE_4KB;
-
 RESERVED_VECTORS_DATA       mReservedVectorsData[CPU_EXCEPTION_NUM];
 EFI_CPU_INTERRUPT_HANDLER   mExternalInterruptHandlerTable[CPU_EXCEPTION_NUM];
 UINTN                       mEnabledInterruptNum = 0;
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c
index aaf90f6..d3b26d3 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c
@@ -119,9 +119,6 @@ DumpImageAndCpuContent (
   IN EFI_SYSTEM_CONTEXT   SystemContext
   )
 {
-  UINTN                   ImageBase;
-  UINTN                   EntryPoint;
-
   InternalPrintMessage (
     "!!!! IA32 Exception Type - %02x(%a)  CPU Apic ID - %08x !!!!\n",
     ExceptionType,
@@ -200,14 +197,8 @@ DumpImageAndCpuContent (
     );
 
   //
-  // Find module image base and module entry point by RIP
+  // Find module image base and module entry point by EIP
   //
-  ImageBase = FindModuleImageBase (SystemContext.SystemContextIa32->Eip, &EntryPoint);
-  if (ImageBase != 0) {
-    InternalPrintMessage (
-      " (ImageBase=%08x, EntryPoint=%08x) !!!!\n",
-      ImageBase,
-      EntryPoint
-      );
-  }
+  DumpModuleImageInfo (SystemContext.SystemContextIa32->Eip);
+
 }
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
index c3fd8ae..53fa3c6 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/PeiCpuException.c
@@ -1,7 +1,7 @@
 /** @file
   CPU exception handler library implementation for PEIM module.
 
-Copyright (c) 2016, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2016 - 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
@@ -18,10 +18,6 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Library/HobLib.h>
 #include <Library/MemoryAllocationLib.h>
 
-//
-// Image Alignment size for PEI phase
-//
-CONST UINTN    mImageAlignSize   = 4;
 CONST UINTN    mDoFarReturnFlag  = 0;
 
 EFI_GUID mCpuExceptrionHandlerLibHobGuid = CPU_EXCEPTION_HANDLER_LIB_HOB_GUID;
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/SecPeiCpuException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/SecPeiCpuException.c
index 7ac3fc2..5d6807b 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/SecPeiCpuException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/SecPeiCpuException.c
@@ -15,10 +15,6 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <PiPei.h>
 #include "CpuExceptionCommon.h"
 
-//
-// Image Aglinment size for SEC/PEI phase
-//
-CONST UINTN    mImageAlignSize   = 4;
 CONST UINTN    mDoFarReturnFlag  = 0;
 
 /**
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/SmmException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/SmmException.c
index 7ad228c..5a3d416 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/SmmException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/SmmException.c
@@ -1,7 +1,7 @@
 /** @file
   CPU exception handler library implemenation for SMM modules.
 
-  Copyright (c) 2013 - 2016, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2013 - 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
   which accompanies this distribution.  The full text of the license may be found at
@@ -22,11 +22,6 @@ CONST UINTN   mDoFarReturnFlag   = 1;
 //
 SPIN_LOCK        mDisplayMessageSpinLock;
 
-//
-// Image align size for DXE/SMM
-//
-CONST UINTN      mImageAlignSize = SIZE_4KB;
-
 RESERVED_VECTORS_DATA       mReservedVectorsData[CPU_EXCEPTION_NUM];
 EFI_CPU_INTERRUPT_HANDLER   mExternalInterruptHandlerTable[CPU_EXCEPTION_NUM];
 EXCEPTION_HANDLER_DATA      mExceptionHandlerData;
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c
index 3cda7d5..9cd2cc2 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c
@@ -124,9 +124,6 @@ DumpImageAndCpuContent (
   IN EFI_SYSTEM_CONTEXT   SystemContext
   )
 {
-  UINTN                   ImageBase;
-  UINTN                   EntryPoint;
-
   InternalPrintMessage (
     "!!!! X64 Exception Type - %02x(%a)  CPU Apic ID - %08x !!!!\n",
     ExceptionType,
@@ -234,12 +231,5 @@ DumpImageAndCpuContent (
   //
   // Find module image base and module entry point by RIP
   //
-  ImageBase = FindModuleImageBase (SystemContext.SystemContextX64->Rip, &EntryPoint);
-  if (ImageBase != 0) {
-    InternalPrintMessage (
-      " (ImageBase=%016lx, EntryPoint=%016lx) !!!!\n",
-      ImageBase,
-      EntryPoint
-      );
-  }
+  DumpModuleImageInfo (SystemContext.SystemContextX64->Rip);
 }
-- 
2.9.3.windows.2