[patch] MdePkg/BasePeCoffLib: Add more check for relocation data

[Dandan Bi <dandan.bi@intel.com>]

REF: https://bugzilla.tianocore.org/show_bug.cgi?id=1426

In function PeCoffLoaderRelocateImageForRuntime, it doesn't
do much check when applies relocation fixups. For API level
consideration, it's not safe enough.
This patch is to replace the same code logic with calling
function PeCoffLoaderImageAddress which will cover more
check and validation.

Cc: Michael D Kinney <michael.d.kinney@intel.com>
Cc: Liming Gao <liming.gao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Dandan Bi <dandan.bi@intel.com>
---
 MdePkg/Library/BasePeCoffLib/BasePeCoff.c | 29 +++++++++++++++++++----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/MdePkg/Library/BasePeCoffLib/BasePeCoff.c b/MdePkg/Library/BasePeCoffLib/BasePeCoff.c
index c57816a808..ae64705d7c 100644
--- a/MdePkg/Library/BasePeCoffLib/BasePeCoff.c
+++ b/MdePkg/Library/BasePeCoffLib/BasePeCoff.c
@@ -13,11 +13,11 @@
   This library will also do some additional check for PE header fields.
 
   PeCoffLoaderGetPeHeader() routine will do basic check for PE/COFF header.
   PeCoffLoaderGetImageInfo() routine will do basic check for whole PE/COFF image.
 
-  Copyright (c) 2006 - 2018, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2006 - 2019, 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
   which accompanies this distribution.  The full text of the license may be found at
   http://opensource.org/licenses/bsd-license.php.
@@ -1669,25 +1669,30 @@ PeCoffLoaderRelocateImageForRuntime (
   UINT32                              NumberOfRvaAndSizes;
   EFI_IMAGE_DATA_DIRECTORY            *DataDirectory;
   EFI_IMAGE_DATA_DIRECTORY            *RelocDir;
   EFI_IMAGE_BASE_RELOCATION           *RelocBase;
   EFI_IMAGE_BASE_RELOCATION           *RelocBaseEnd;
+  EFI_IMAGE_BASE_RELOCATION           *RelocBaseOrig;
   UINT16                              *Reloc;
   UINT16                              *RelocEnd;
   CHAR8                               *Fixup;
   CHAR8                               *FixupBase;
   UINT16                              *Fixup16;
   UINT32                              *Fixup32;
   UINT64                              *Fixup64;
   CHAR8                               *FixupData;
   UINTN                               Adjust;
   RETURN_STATUS                       Status;
+  PE_COFF_LOADER_IMAGE_CONTEXT        ImageContext;
 
   OldBase = (CHAR8 *)((UINTN)ImageBase);
   NewBase = (CHAR8 *)((UINTN)VirtImageBase);
   Adjust = (UINTN) NewBase - (UINTN) OldBase;
 
+  ImageContext.ImageAddress = ImageBase;
+  ImageContext.ImageSize = ImageSize;
+
   //
   // Find the image's relocate dir info
   //
   DosHdr = (EFI_IMAGE_DOS_HEADER *)OldBase;
   if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
@@ -1730,12 +1735,15 @@ PeCoffLoaderRelocateImageForRuntime (
   // is present in the image. You have to check the NumberOfRvaAndSizes in
   // the optional header to verify a desired directory entry is there.
   //
   if (NumberOfRvaAndSizes > EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC) {
     RelocDir      = DataDirectory + EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC;
-    RelocBase     = (EFI_IMAGE_BASE_RELOCATION *)(UINTN)(ImageBase + RelocDir->VirtualAddress);
-    RelocBaseEnd  = (EFI_IMAGE_BASE_RELOCATION *)(UINTN)(ImageBase + RelocDir->VirtualAddress + RelocDir->Size);
+    RelocBase     = (EFI_IMAGE_BASE_RELOCATION *) PeCoffLoaderImageAddress (&ImageContext, RelocDir->VirtualAddress, 0);
+    RelocBaseEnd  = (EFI_IMAGE_BASE_RELOCATION *) PeCoffLoaderImageAddress (&ImageContext,
+                                                                            RelocDir->VirtualAddress + RelocDir->Size,
+                                                                            0
+                                                                            );
   } else {
     //
     // Cannot find relocations, cannot continue to relocate the image, ASSERT for this invalid image.
     //
     ASSERT (FALSE);
@@ -1753,10 +1761,11 @@ PeCoffLoaderRelocateImageForRuntime (
   // since it was relocated. This is so data sections that have been updated
   // by code will not be fixed up, since that would set them back to
   // defaults.
   //
   FixupData = RelocationData;
+  RelocBaseOrig = RelocBase;
   while (RelocBase < RelocBaseEnd) {
     //
     // Add check for RelocBase->SizeOfBlock field.
     //
     if ((RelocBase->SizeOfBlock == 0) || (RelocBase->SizeOfBlock > RelocDir->Size)) {
@@ -1766,18 +1775,28 @@ PeCoffLoaderRelocateImageForRuntime (
       return;
     }
 
     Reloc     = (UINT16 *) ((UINT8 *) RelocBase + sizeof (EFI_IMAGE_BASE_RELOCATION));
     RelocEnd  = (UINT16 *) ((UINT8 *) RelocBase + RelocBase->SizeOfBlock);
-    FixupBase = (CHAR8 *) ((UINTN)ImageBase) + RelocBase->VirtualAddress;
+    if ((UINTN)RelocEnd > (UINTN)RelocBaseOrig + RelocDir->Size) {
+      return;
+    }
+
+    FixupBase = PeCoffLoaderImageAddress (&ImageContext, RelocBase->VirtualAddress, 0);
+    if (FixupBase == NULL) {
+      return;
+    }
 
     //
     // Run this relocation record
     //
     while (Reloc < RelocEnd) {
 
-      Fixup = FixupBase + (*Reloc & 0xFFF);
+      Fixup = PeCoffLoaderImageAddress (&ImageContext, RelocBase->VirtualAddress + (*Reloc & 0xFFF), 0);
+      if (Fixup == NULL) {
+        return;
+      }
       switch ((*Reloc) >> 12) {
 
       case EFI_IMAGE_REL_BASED_ABSOLUTE:
         break;
 
-- 
2.18.0.windows.1