BaseTools/PeCoffLib: Check 'RelocDir' before finding relocation block

To match the code logics in MdePkg/Library/BasePeCoffLib, add checks for
'RelocDir' before finding the relocation block.

Cc: Liming Gao <liming.gao@intel.com>
Cc: Yonghong Zhu <yonghong.zhu@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Hao Wu <hao.a.wu@intel.com>
Reviewed-by: Liming Gao <liming.gao@intel.com>

diff --git a/BaseTools/Source/C/Common/BasePeCoff.c b/BaseTools/Source/C/Common/BasePeCoff.c
index 9adbdfaf09da202c25a939ec914ac45cccb6de18..692663d0e497715d4e0d6a8544ecd5bffe0b8afd 100644 (file)
--- a/BaseTools/Source/C/Common/BasePeCoff.c
+++ b/BaseTools/Source/C/Common/BasePeCoff.c
@@ -645,14 +645,21 @@ Returns:
       //\r
       if (OptionHeader.Optional32->NumberOfRvaAndSizes > EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC) {\r
         RelocDir  = &OptionHeader.Optional32->DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC];\r
-        RelocBase = PeCoffLoaderImageAddress (ImageContext, RelocDir->VirtualAddress);\r
-        RelocBaseEnd = PeCoffLoaderImageAddress (\r
-                        ImageContext,\r
-                        RelocDir->VirtualAddress + RelocDir->Size - 1\r
-                        );\r
-        if (RelocBase == NULL || RelocBaseEnd == NULL || RelocBaseEnd < RelocBase) {\r
-          ImageContext->ImageError = IMAGE_ERROR_FAILED_RELOCATION;\r
-          return RETURN_LOAD_ERROR;\r
+        if ((RelocDir != NULL) && (RelocDir->Size > 0)) {\r
+          RelocBase = PeCoffLoaderImageAddress (ImageContext, RelocDir->VirtualAddress);\r
+          RelocBaseEnd = PeCoffLoaderImageAddress (\r
+                           ImageContext,\r
+                           RelocDir->VirtualAddress + RelocDir->Size - 1\r
+                           );\r
+          if (RelocBase == NULL || RelocBaseEnd == NULL || RelocBaseEnd < RelocBase) {\r
+            ImageContext->ImageError = IMAGE_ERROR_FAILED_RELOCATION;\r
+            return RETURN_LOAD_ERROR;\r
+          }\r
+        } else {\r
+          //\r
+          // Set base and end to bypass processing below.\r
+          //\r
+          RelocBase = RelocBaseEnd = 0;\r
         }\r
       } else {\r
         //\r
@@ -673,14 +680,21 @@ Returns:
       //\r
       if (OptionHeader.Optional64->NumberOfRvaAndSizes > EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC) {\r
         RelocDir  = &OptionHeader.Optional64->DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC];\r
-        RelocBase = PeCoffLoaderImageAddress (ImageContext, RelocDir->VirtualAddress);\r
-        RelocBaseEnd = PeCoffLoaderImageAddress (\r
-                        ImageContext,\r
-                        RelocDir->VirtualAddress + RelocDir->Size - 1\r
-                        );\r
-        if (RelocBase == NULL || RelocBaseEnd == NULL || RelocBaseEnd < RelocBase) {\r
-          ImageContext->ImageError = IMAGE_ERROR_FAILED_RELOCATION;\r
-          return RETURN_LOAD_ERROR;\r
+        if ((RelocDir != NULL) && (RelocDir->Size > 0)) {\r
+          RelocBase = PeCoffLoaderImageAddress (ImageContext, RelocDir->VirtualAddress);\r
+          RelocBaseEnd = PeCoffLoaderImageAddress (\r
+                           ImageContext,\r
+                           RelocDir->VirtualAddress + RelocDir->Size - 1\r
+                          );\r
+          if (RelocBase == NULL || RelocBaseEnd == NULL || RelocBaseEnd < RelocBase) {\r
+            ImageContext->ImageError = IMAGE_ERROR_FAILED_RELOCATION;\r
+            return RETURN_LOAD_ERROR;\r
+          }\r
+        } else {\r
+          //\r
+          // Set base and end to bypass processing below.\r
+          //\r
+          RelocBase = RelocBaseEnd = 0;\r
         }\r
       } else {\r
         //\r