Merged back the Yizhong's fix which was overwritten by check-in of r2157,2158.

[mirror_edk2.git] / Tools / CCode / Source / PeiRebase / PeiRebaseExe.c

diff --git a/Tools/CCode/Source/PeiRebase/PeiRebaseExe.c b/Tools/CCode/Source/PeiRebase/PeiRebaseExe.c
index 8537d69f68fd9b0d13a6003ca40f5a426f291658..5f7dfc25c71cd47a3e78d2c8baedd4e5fa87bc84 100644 (file)
--- a/Tools/CCode/Source/PeiRebase/PeiRebaseExe.c
+++ b/Tools/CCode/Source/PeiRebase/PeiRebaseExe.c
@@ -1,8 +1,8 @@
 /*++\r
 \r
 Copyright (c)  1999-2006 Intel Corporation. All rights reserved\r
 /*++\r
 \r
 Copyright (c)  1999-2006 Intel Corporation. All rights reserved\r

-This program and the accompanying materials are licensed and made available \r
-under the terms and conditions of the BSD License which accompanies this \r
+This program and the accompanying materials are licensed and made available\r
+under the terms and conditions of the BSD License which accompanies this\r

 distribution.  The full text of the license may be found at\r
 http://opensource.org/licenses/bsd-license.php\r
 \r
 distribution.  The full text of the license may be found at\r
 http://opensource.org/licenses/bsd-license.php\r
 \r


@@ -59,23 +59,21 @@ Arguments:
 \r
   argc          - Number of command line arguments\r
   argv[]:\r
 \r
   argc          - Number of command line arguments\r
   argv[]:\r

-  BaseAddress     The base address to use for rebasing the FV.  The correct \r
+  BaseAddress     The base address to use for rebasing the FV.  The correct\r

                   format is a hex number preceded by 0x.\r
   InputFileName   The name of the input FV file.\r
   OutputFileName  The name of the output FV file.\r
                   format is a hex number preceded by 0x.\r
   InputFileName   The name of the input FV file.\r
   OutputFileName  The name of the output FV file.\r

-  MapFileName     The name of the map file of relocation info.\r

 \r
   Arguments come in pair in any order.\r
 \r
   Arguments come in pair in any order.\r

-    -I InputFileName \r
+    -I InputFileName\r

     -O OutputFileName\r
     -O OutputFileName\r

-    -B BaseAddress \r
-    -M MapFileName \r
+    -B BaseAddress\r

 \r
 Returns:\r
 \r
   0   No error conditions detected.\r
   1   One or more of the input parameters is invalid.\r
 \r
 Returns:\r
 \r
   0   No error conditions detected.\r
   1   One or more of the input parameters is invalid.\r

-  2   A resource required by the utility was unavailable.  \r
+  2   A resource required by the utility was unavailable.\r

       Most commonly this will be memory allocation or file creation.\r
   3   PeiRebase.dll could not be loaded.\r
   4   Error executing the PEI rebase.\r
       Most commonly this will be memory allocation or file creation.\r
   3   PeiRebase.dll could not be loaded.\r
   4   Error executing the PEI rebase.\r


@@ -84,14 +82,13 @@ Returns:
 {\r
   UINT8                       Index;\r
   CHAR8                       InputFileName[_MAX_PATH];\r
 {\r
   UINT8                       Index;\r
   CHAR8                       InputFileName[_MAX_PATH];\r

-  CHAR8                       OutputFileName[_MAX_PATH];\r
-  CHAR8                       MapFileName[_MAX_PATH];\r
-  EFI_PHYSICAL_ADDRESS        BaseAddress;\r
-  BOOLEAN                     BaseAddressSet;\r
+  CHAR8                       *OutputFileName;\r
+  EFI_PHYSICAL_ADDRESS        XipBase, BsBase, RtBase;\r
+  UINT32                      BaseTypes;\r

   EFI_STATUS                  Status;\r
   FILE                        *InputFile;\r
   FILE                        *OutputFile;\r
   EFI_STATUS                  Status;\r
   FILE                        *InputFile;\r
   FILE                        *OutputFile;\r

-  FILE                        *MapFile;\r
+  FILE                        *LogFile;\r

   UINT64                      FvOffset;\r
   UINT32                      FileCount;\r
   int                         BytesRead;\r
   UINT64                      FvOffset;\r
   UINT32                      FileCount;\r
   int                         BytesRead;\r


@@ -99,11 +96,8 @@ Returns:
   UINT32                      FvSize;\r
   EFI_FFS_FILE_HEADER         *CurrentFile;\r
   BOOLEAN                     ErasePolarity;\r
   UINT32                      FvSize;\r
   EFI_FFS_FILE_HEADER         *CurrentFile;\r
   BOOLEAN                     ErasePolarity;\r

-  EFI_PHYSICAL_ADDRESS        CurrentFileBaseAddress;\r
-  CHAR8                       InfFileName[_MAX_PATH];\r
-  CHAR8                       *InfFileImage;\r
-  UINTN                       InfFileSize;\r

   MEMORY_FILE                 InfMemoryFile;\r
   MEMORY_FILE                 InfMemoryFile;\r

+  CHAR8                       StringBuffer[0x100];\r

 \r
   ErasePolarity = FALSE;\r
   //\r
 \r
   ErasePolarity = FALSE;\r
   //\r


@@ -121,21 +115,17 @@ Returns:
   //\r
   // Initialize variables\r
   //\r
   //\r
   // Initialize variables\r
   //\r

-  InputFileName[0]  = 0;\r
-  OutputFileName[0] = 0;\r
-  MapFileName[0]    = 0;\r
-  BaseAddress       = 0;\r
-  BaseAddressSet    = FALSE;\r
+  InputFileName[0]  = '\0';\r
+  OutputFileName    = NULL;\r
+  XipBase = BsBase = RtBase = 0;\r
+  BaseTypes         = 0;\r

   FvOffset          = 0;\r
   FileCount         = 0;\r
   ErasePolarity     = FALSE;\r
   InputFile         = NULL;\r
   OutputFile        = NULL;\r
   FvOffset          = 0;\r
   FileCount         = 0;\r
   ErasePolarity     = FALSE;\r
   InputFile         = NULL;\r
   OutputFile        = NULL;\r

-  MapFile           = NULL;\r
+  LogFile           = NULL;\r

   FvImage           = NULL;\r
   FvImage           = NULL;\r

-  InfFileImage      = NULL;\r
-  InfFileSize       = 0;\r
-  strcpy (InfFileName, "");\r

 \r
   //\r
   // Parse the command line arguments\r
 \r
   //\r
   // Parse the command line arguments\r


@@ -156,7 +146,7 @@ Returns:
       PrintUsage ();\r
       Error (NULL, 0, 0, argv[Index], "unrecognized option");\r
       return STATUS_ERROR;\r
       PrintUsage ();\r
       Error (NULL, 0, 0, argv[Index], "unrecognized option");\r
       return STATUS_ERROR;\r

-    }    \r
+    }\r

     //\r
     // Determine argument to read\r
     //\r
     //\r
     // Determine argument to read\r
     //\r


@@ -174,8 +164,8 @@ Returns:
 \r
     case 'O':\r
     case 'o':\r
 \r
     case 'O':\r
     case 'o':\r

-      if (strlen (OutputFileName) == 0) {\r
-        strcpy (OutputFileName, argv[Index + 1]);\r
+      if (OutputFileName == NULL) {\r
+        OutputFileName = argv[Index + 1];\r

       } else {\r
         PrintUsage ();\r
         Error (NULL, 0, 0, argv[Index + 1], "only one -o OutputFileName may be specified");\r
       } else {\r
         PrintUsage ();\r
         Error (NULL, 0, 0, argv[Index + 1], "only one -o OutputFileName may be specified");\r


@@ -183,79 +173,86 @@ Returns:
       }\r
       break;\r
 \r
       }\r
       break;\r
 \r

+    case 'F':\r
+    case 'f':\r
+      //\r
+      // Load INF file into memory & initialize MEMORY_FILE structure\r
+      //\r
+      Status = GetFileImage (argv[Index + 1], &InfMemoryFile.FileImage, (UINT32*)&InfMemoryFile.Eof);\r
+      InfMemoryFile.Eof = InfMemoryFile.FileImage + (UINT32)(UINTN)InfMemoryFile.Eof;\r
+      InfMemoryFile.CurrentFilePointer = InfMemoryFile.FileImage;\r
+      if (EFI_ERROR (Status)) {\r
+        Error (NULL, 0, 0, argv[Index + 1], "Error opening FvInfFile");\r
+        return STATUS_ERROR;\r
+      }\r
+\r
+      //\r
+      // Read BaseAddress from fv.inf file\r
+      //\r
+      FindToken (&InfMemoryFile, "[options]", "EFI_BASE_ADDRESS", 0, StringBuffer);\r
+\r
+      //\r
+      // Free INF file image\r
+      //\r
+      free (InfMemoryFile.FileImage);\r
+\r
+      //\r
+      // Point argv[Index + 1] to StringBuffer so that it could be processed as "-b"\r
+      //\r
+      argv[Index + 1] = StringBuffer;\r
+\r

     case 'B':\r
     case 'b':\r
     case 'B':\r
     case 'b':\r

-      if (!BaseAddressSet) {\r
-        Status = AsciiStringToUint64 (argv[Index + 1], FALSE, &BaseAddress);\r
-        if (EFI_ERROR (Status)) {\r
-          PrintUsage ();\r
-          Error (NULL, 0, 0, argv[Index + 1], "invalid hex digit given for the base address");\r
-          return STATUS_ERROR;\r
-        }\r
+      if (BaseTypes & 1) {\r
+        PrintUsage ();\r
+        Error (NULL, 0, 0, argv[Index + 1], "XipBaseAddress may be specified only once by either -b or -f");\r
+        return STATUS_ERROR;\r
+      }\r

 \r
 \r

-        BaseAddressSet = TRUE;\r
-      } else {\r
+      Status = AsciiStringToUint64 (argv[Index + 1], FALSE, &XipBase);\r
+      if (EFI_ERROR (Status)) {\r

         PrintUsage ();\r
         PrintUsage ();\r

-        Error (NULL, 0, 0, argv[Index + 1], "-b BaseAddress may only be specified once");\r
+        Error (NULL, 0, 0, argv[Index + 1], "invalid hex digit given for XIP base address");\r

         return STATUS_ERROR;\r
       }\r
         return STATUS_ERROR;\r
       }\r

+\r
+      BaseTypes |= 1;\r

       break;\r
 \r
       break;\r
 \r

-    case 'F':\r
-    case 'f':\r
-      if (!BaseAddressSet) {\r
-        strcpy (InfFileName, argv[Index + 1]);\r
-        //\r
-        // Read the INF file image\r
-        //\r
-        Status = GetFileImage (InfFileName, &InfFileImage, &InfFileSize);\r
-        if (EFI_ERROR (Status)) {\r
-          PrintUsage ();\r
-          Error (NULL, 0, 0, argv[Index + 1], "-f FvInfFile can't be opened.");\r
-          return STATUS_ERROR;\r
-        }\r
-        //\r
-        // Initialize file structures\r
-        //\r
-        InfMemoryFile.FileImage           = InfFileImage;\r
-        InfMemoryFile.CurrentFilePointer  = InfFileImage;\r
-        InfMemoryFile.Eof                 = InfFileImage + InfFileSize;\r
-        //\r
-        // Read BaseAddress from fv.inf file.\r
-        //\r
-        FindToken (&InfMemoryFile, "[options]", "EFI_BASE_ADDRESS", 0, InfFileName);\r
-        //\r
-        // free Inf File Image\r
-        //\r
-        free (InfFileImage);\r
-        \r
-        //\r
-        // Convert string to UINT64 base address.\r
-        //\r
-        Status = AsciiStringToUint64 (InfFileName, FALSE, &BaseAddress);\r
-        if (EFI_ERROR (Status)) {\r
-          PrintUsage ();\r
-          Error (NULL, 0, 0, argv[Index + 1], "can't find the base address in the specified fv.inf file.");\r
-          return STATUS_ERROR;\r
-        }\r
+    case 'D':\r
+    case 'd':\r
+      if (BaseTypes & 2) {\r
+        PrintUsage ();\r
+        Error (NULL, 0, 0, argv[Index + 1], "-d BsBaseAddress may be specified only once");\r
+        return STATUS_ERROR;\r
+      }\r

 \r
 \r

-        BaseAddressSet = TRUE;\r
-      } else {\r
+      Status = AsciiStringToUint64 (argv[Index + 1], FALSE, &BsBase);\r
+      if (EFI_ERROR (Status)) {\r

         PrintUsage ();\r
         PrintUsage ();\r

-        Error (NULL, 0, 0, argv[Index + 1], "BaseAddress has been got once from fv.inf or the specified base address.");\r
+        Error (NULL, 0, 0, argv[Index + 1], "invalid hex digit given for BS_DRIVER base address");\r

         return STATUS_ERROR;\r
       }\r
         return STATUS_ERROR;\r
       }\r

+\r
+      BaseTypes |= 2;\r

       break;\r
 \r
       break;\r
 \r

-    case 'M':\r
-    case 'm':\r
-      if (strlen (MapFileName) == 0) {\r
-        strcpy (MapFileName, argv[Index + 1]);\r
-      } else {\r
+    case 'R':\r
+    case 'r':\r
+      if (BaseTypes & 4) {\r

         PrintUsage ();\r
         PrintUsage ();\r

-        Error (NULL, 0, 0, argv[Index + 1], "only one -m MapFileName may be specified");\r
+        Error (NULL, 0, 0, argv[Index + 1], "-r RtBaseAddress may be specified only once");\r

         return STATUS_ERROR;\r
       }\r
         return STATUS_ERROR;\r
       }\r

+\r
+      Status = AsciiStringToUint64 (argv[Index + 1], FALSE, &RtBase);\r
+      if (EFI_ERROR (Status)) {\r
+        PrintUsage ();\r
+        Error (NULL, 0, 0, argv[Index + 1], "invalid hex digit given for RT_DRIVER base address");\r
+        return STATUS_ERROR;\r
+      }\r
+\r
+      BaseTypes |= 4;\r

       break;\r
 \r
     default:\r
       break;\r
 \r
     default:\r


@@ -265,18 +262,6 @@ Returns:
       break;\r
     }\r
   }\r
       break;\r
     }\r
   }\r

-  \r
-  //\r
-  // Create the Map file if we need it\r
-  //\r
-  if (strlen (MapFileName) != 0) {\r
-    MapFile = fopen (MapFileName, "w");\r
-    if (MapFile == NULL) {\r
-      Error (NULL, 0, 0, MapFileName, "failed to open map file");\r
-      goto Finish;\r
-    }\r
-  } \r
-\r

   //\r
   // Open the file containing the FV\r
   //\r
   //\r
   // Open the file containing the FV\r
   //\r


@@ -285,6 +270,16 @@ Returns:
     Error (NULL, 0, 0, InputFileName, "could not open input file for reading");\r
     return STATUS_ERROR;\r
   }\r
     Error (NULL, 0, 0, InputFileName, "could not open input file for reading");\r
     return STATUS_ERROR;\r
   }\r

+\r
+  //\r
+  // Open the log file\r
+  //\r
+  strcat (InputFileName, ".log");\r
+  LogFile = fopen (InputFileName, "a");\r
+  if (LogFile == NULL) {\r
+    Error (NULL, 0, 0, InputFileName, "could not append to log file");\r
+  }\r
+\r

   //\r
   // Determine size of FV\r
   //\r
   //\r
   // Determine size of FV\r
   //\r


@@ -330,8 +325,14 @@ Returns:
     //\r
     // Rebase this file\r
     //\r
     //\r
     // Rebase this file\r
     //\r

-    CurrentFileBaseAddress  = BaseAddress + ((UINTN) CurrentFile - (UINTN) FvImage);\r
-    Status                  = FfsRebase (CurrentFile, CurrentFileBaseAddress, MapFile);\r
+    FfsRebase (\r
+      CurrentFile,\r
+      BaseTypes,\r
+      XipBase + (UINTN)CurrentFile - (UINTN)FvImage,\r
+      &BsBase,\r
+      &RtBase,\r
+      LogFile\r
+      );\r

 \r
     if (EFI_ERROR (Status)) {\r
       switch (Status) {\r
 \r
     if (EFI_ERROR (Status)) {\r
       switch (Status) {\r


@@ -359,7 +360,6 @@ Returns:
 \r
       goto Finish;\r
     }\r
 \r
       goto Finish;\r
     }\r

-\r

     //\r
     // Get the next file\r
     //\r
     //\r
     // Get the next file\r
     //\r


@@ -399,8 +399,8 @@ Finish:
     fclose (OutputFile);\r
   }\r
 \r
     fclose (OutputFile);\r
   }\r
 \r

-  if (MapFile != NULL) {\r
-    fclose (MapFile);\r
+  if (LogFile != NULL) {\r
+    fclose (LogFile);\r

   }\r
 \r
   if (FvImage != NULL) {\r
   }\r
 \r
   if (FvImage != NULL) {\r


@@ -420,7 +420,7 @@ ReadHeader (
 \r
 Routine Description:\r
 \r
 \r
 Routine Description:\r
 \r

-  This function determines the size of the FV and the erase polarity.  The \r
+  This function determines the size of the FV and the erase polarity.  The\r

   erase polarity is the FALSE value for file state.\r
 \r
 Arguments:\r
   erase polarity is the FALSE value for file state.\r
 \r
 Arguments:\r


@@ -428,9 +428,9 @@ Arguments:
   InputFile       The file that contains the FV image.\r
   FvSize          The size of the FV.\r
   ErasePolarity   The FV erase polarity.\r
   InputFile       The file that contains the FV image.\r
   FvSize          The size of the FV.\r
   ErasePolarity   The FV erase polarity.\r

-    \r
+\r

 Returns:\r
 Returns:\r

- \r
+\r

   EFI_SUCCESS             Function completed successfully.\r
   EFI_INVALID_PARAMETER   A required parameter was NULL or is out of range.\r
   EFI_ABORTED             The function encountered an error.\r
   EFI_SUCCESS             Function completed successfully.\r
   EFI_INVALID_PARAMETER   A required parameter was NULL or is out of range.\r
   EFI_ABORTED             The function encountered an error.\r


@@ -539,38 +539,36 @@ Returns:
 --*/\r
 {\r
   printf (\r
 --*/\r
 {\r
   printf (\r

-    "Usage: %s -I InputFileName -O OutputFileName [-B BaseAddress] -F FvInfFileName -M MapFile\n",\r
+    "Usage: %s -I InputFileName -O OutputFileName -B BaseAddress\n",\r

     UTILITY_NAME\r
     );\r
   printf ("  Where:\n");\r
     UTILITY_NAME\r
     );\r
   printf ("  Where:\n");\r

-  printf ("    InputFileName  is the name of the EFI FV file to rebase.\n");\r
+  printf ("    InputFileName is the name of the EFI FV file to rebase.\n");\r

   printf ("    OutputFileName is the desired output file name.\n");\r
   printf ("    OutputFileName is the desired output file name.\n");\r

-  printf ("    BaseAddress    is the FV base address to rebase agains.\n");\r
-  printf ("    FvInfFileName  is the fv.inf to be used to generate this fv image.\n");\r
-  printf ("                   BaseAddress can also be got from the fv.inf file.\n");\r
-  printf ("                   Choose only one method to input BaseAddress.\n");\r
-  printf ("    MapFileName    is an optional map file of the relocations\n");\r
-  printf ("    Argument pair may be in any order.\n\n");\r
+  printf ("    BaseAddress is the FV base address to rebase against.\n");\r
+  printf ("  Argument pair may be in any order.\n\n");\r

 }\r
 \r
 EFI_STATUS\r
 FfsRebase (\r
 }\r
 \r
 EFI_STATUS\r
 FfsRebase (\r

-  IN OUT EFI_FFS_FILE_HEADER    *FfsFile,\r
-  IN EFI_PHYSICAL_ADDRESS       BaseAddress,\r
-  IN FILE                       *MapFile      OPTIONAL\r
+  IN OUT  EFI_FFS_FILE_HEADER       *FfsFile,\r
+  IN      UINT32                    Flags,\r
+  IN OUT  EFI_PHYSICAL_ADDRESS      XipBase,\r
+  IN OUT  EFI_PHYSICAL_ADDRESS      *BsBase,\r
+  IN OUT  EFI_PHYSICAL_ADDRESS      *RtBase,\r
+  OUT     FILE                      *LogFile\r

   )\r
 /*++\r
 \r
 Routine Description:\r
 \r
   )\r
 /*++\r
 \r
 Routine Description:\r
 \r

-  This function determines if a file is XIP and should be rebased.  It will \r
+  This function determines if a file is XIP and should be rebased.  It will\r

   rebase any PE32 sections found in the file using the base address.\r
   rebase any PE32 sections found in the file using the base address.\r

-  \r
+\r

 Arguments:\r
 \r
   FfsFile           A pointer to Ffs file image.\r
   BaseAddress       The base address to use for rebasing the file image.\r
 Arguments:\r
 \r
   FfsFile           A pointer to Ffs file image.\r
   BaseAddress       The base address to use for rebasing the file image.\r

-  MapFile           Optional file to dump relocation information into\r

 \r
 Returns:\r
 \r
 \r
 Returns:\r
 \r


@@ -590,20 +588,21 @@ Returns:
   UINT64                                ImageSize;\r
   EFI_PHYSICAL_ADDRESS                  EntryPoint;\r
   UINT32                                Pe32ImageSize;\r
   UINT64                                ImageSize;\r
   EFI_PHYSICAL_ADDRESS                  EntryPoint;\r
   UINT32                                Pe32ImageSize;\r

-  UINT32                                NewPe32BaseAddress;\r
+  EFI_PHYSICAL_ADDRESS                  NewPe32BaseAddress;\r

   UINTN                                 Index;\r
   EFI_FILE_SECTION_POINTER              CurrentPe32Section;\r
   EFI_FFS_FILE_STATE                    SavedState;\r
   UINTN                                 Index;\r
   EFI_FILE_SECTION_POINTER              CurrentPe32Section;\r
   EFI_FFS_FILE_STATE                    SavedState;\r

-  EFI_IMAGE_NT_HEADERS                  *PeHdr;\r
+  EFI_IMAGE_NT_HEADERS32                *PeHdr;\r
+  EFI_IMAGE_NT_HEADERS64                *PePlusHdr;\r

   UINT32                                *PeHdrSizeOfImage;\r
   UINT32                                *PeHdrChecksum;\r
   UINT32                                *PeHdrSizeOfImage;\r
   UINT32                                *PeHdrChecksum;\r

-  UINT32                                FoundCount;\r

   EFI_TE_IMAGE_HEADER                   *TEImageHeader;\r
   UINT8                                 *TEBuffer;\r
   EFI_IMAGE_DOS_HEADER                  *DosHeader;\r
   UINT8                                 FileGuidString[80];\r
   UINT32                                TailSize;\r
   EFI_FFS_FILE_TAIL                     TailValue;\r
   EFI_TE_IMAGE_HEADER                   *TEImageHeader;\r
   UINT8                                 *TEBuffer;\r
   EFI_IMAGE_DOS_HEADER                  *DosHeader;\r
   UINT8                                 FileGuidString[80];\r
   UINT32                                TailSize;\r
   EFI_FFS_FILE_TAIL                     TailValue;\r

+  EFI_PHYSICAL_ADDRESS                  *BaseToUpdate;\r

 \r
   //\r
   // Verify input parameters\r
 \r
   //\r
   // Verify input parameters\r


@@ -611,7 +610,6 @@ Returns:
   if (FfsFile == NULL) {\r
     return EFI_INVALID_PARAMETER;\r
   }\r
   if (FfsFile == NULL) {\r
     return EFI_INVALID_PARAMETER;\r
   }\r

-  \r

   //\r
   // Convert the GUID to a string so we can at least report which file\r
   // if we find an error.\r
   //\r
   // Convert the GUID to a string so we can at least report which file\r
   // if we find an error.\r


@@ -622,7 +620,6 @@ Returns:
   } else {\r
     TailSize = 0;\r
   }\r
   } else {\r
     TailSize = 0;\r
   }\r

-  \r

   //\r
   // Do some cursory checks on the FFS file contents\r
   //\r
   //\r
   // Do some cursory checks on the FFS file contents\r
   //\r


@@ -632,36 +629,31 @@ Returns:
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r

-  memset (&ImageContext, 0, sizeof (ImageContext));\r
-\r

   //\r
   //\r

-  // Check if XIP file type. If not XIP, don't rebase.\r
+  // We only process files potentially containing PE32 sections.\r

   //\r
   //\r

-  if (FfsFile->Type != EFI_FV_FILETYPE_PEI_CORE &&\r
-      FfsFile->Type != EFI_FV_FILETYPE_PEIM &&\r
-      FfsFile->Type != EFI_FV_FILETYPE_SECURITY_CORE &&\r
-      FfsFile->Type != EFI_FV_FILETYPE_COMBINED_PEIM_DRIVER\r
-      ) {\r
-    return EFI_SUCCESS;\r
+  switch (FfsFile->Type) {\r
+    case EFI_FV_FILETYPE_SECURITY_CORE:\r
+    case EFI_FV_FILETYPE_PEI_CORE:\r
+    case EFI_FV_FILETYPE_PEIM:\r
+    case EFI_FV_FILETYPE_COMBINED_PEIM_DRIVER:\r
+    case EFI_FV_FILETYPE_DRIVER:\r
+    case EFI_FV_FILETYPE_DXE_CORE:\r
+      break;\r
+    default:\r
+      return EFI_SUCCESS;\r

   }\r
 \r
   //\r
   // Rebase each PE32 section\r
   //\r
   Status      = EFI_SUCCESS;\r
   }\r
 \r
   //\r
   // Rebase each PE32 section\r
   //\r
   Status      = EFI_SUCCESS;\r

-  FoundCount  = 0;\r

   for (Index = 1;; Index++) {\r
     Status = GetSectionByType (FfsFile, EFI_SECTION_PE32, Index, &CurrentPe32Section);\r
     if (EFI_ERROR (Status)) {\r
       break;\r
     }\r
 \r
   for (Index = 1;; Index++) {\r
     Status = GetSectionByType (FfsFile, EFI_SECTION_PE32, Index, &CurrentPe32Section);\r
     if (EFI_ERROR (Status)) {\r
       break;\r
     }\r
 \r

-    FoundCount++;\r
-\r
-    //\r
-    // Calculate the PE32 base address, the FFS file base plus the offset of the PE32 section\r
-    //\r
-    NewPe32BaseAddress = ((UINT32) BaseAddress) + ((UINTN) CurrentPe32Section.Pe32Section + sizeof (EFI_COMMON_SECTION_HEADER) - (UINTN) FfsFile);\r

 \r
     //\r
     // Initialize context\r
 \r
     //\r
     // Initialize context\r


@@ -669,13 +661,90 @@ Returns:
     memset (&ImageContext, 0, sizeof (ImageContext));\r
     ImageContext.Handle     = (VOID *) ((UINTN) CurrentPe32Section.Pe32Section + sizeof (EFI_PE32_SECTION));\r
     ImageContext.ImageRead  = (PE_COFF_LOADER_READ_FILE) FfsRebaseImageRead;\r
     memset (&ImageContext, 0, sizeof (ImageContext));\r
     ImageContext.Handle     = (VOID *) ((UINTN) CurrentPe32Section.Pe32Section + sizeof (EFI_PE32_SECTION));\r
     ImageContext.ImageRead  = (PE_COFF_LOADER_READ_FILE) FfsRebaseImageRead;\r

-\r

     Status                  = PeCoffLoaderGetImageInfo (&ImageContext);\r
     Status                  = PeCoffLoaderGetImageInfo (&ImageContext);\r

-\r

     if (EFI_ERROR (Status)) {\r
       Error (NULL, 0, 0, "GetImageInfo() call failed on rebase", FileGuidString);\r
       return Status;\r
     }\r
     if (EFI_ERROR (Status)) {\r
       Error (NULL, 0, 0, "GetImageInfo() call failed on rebase", FileGuidString);\r
       return Status;\r
     }\r

+\r
+    //\r
+    // Calculate the PE32 base address, based on file type\r
+    //\r
+    switch (FfsFile->Type) {\r
+      case EFI_FV_FILETYPE_SECURITY_CORE:\r
+      case EFI_FV_FILETYPE_PEI_CORE:\r
+      case EFI_FV_FILETYPE_PEIM:\r
+      case EFI_FV_FILETYPE_COMBINED_PEIM_DRIVER:\r
+        if ((Flags & 1) == 0) {\r
+          //\r
+          // We aren't relocating XIP code, so skip it.\r
+          //\r
+          return EFI_SUCCESS;\r
+        }\r
+\r
+        NewPe32BaseAddress =\r
+          XipBase +\r
+          (UINTN)CurrentPe32Section.Pe32Section +\r
+          sizeof (EFI_COMMON_SECTION_HEADER) -\r
+          (UINTN)FfsFile;\r
+        BaseToUpdate = &XipBase;\r
+        break;\r
+\r
+      case EFI_FV_FILETYPE_DRIVER:\r
+        PeHdr = (EFI_IMAGE_NT_HEADERS32*)(\r
+          (UINTN)CurrentPe32Section.Pe32Section +\r
+          sizeof (EFI_COMMON_SECTION_HEADER) +\r
+          ImageContext.PeCoffHeaderOffset\r
+          );\r
+        switch (PeHdr->OptionalHeader.Subsystem) {\r
+          case EFI_IMAGE_SUBSYSTEM_EFI_RUNTIME_DRIVER:\r
+            if ((Flags & 4) == 0) {\r
+              //\r
+              // RT drivers aren't supposed to be relocated\r
+              //\r
+              continue;\r
+            }\r
+\r
+            NewPe32BaseAddress = *RtBase;\r
+            BaseToUpdate = RtBase;\r
+            break;\r
+\r
+          default:\r
+            //\r
+            // We treat all other subsystems the same as BS_DRIVER\r
+            //\r
+            if ((Flags & 2) == 0) {\r
+              //\r
+              // Skip all BS_DRIVER's\r
+              //\r
+              continue;\r
+            }\r
+\r
+            NewPe32BaseAddress = *BsBase;\r
+            BaseToUpdate = BsBase;\r
+            break;\r
+        }\r
+        break;\r
+\r
+      case EFI_FV_FILETYPE_DXE_CORE:\r
+        if ((Flags & 2) == 0) {\r
+          //\r
+          // Skip DXE core\r
+          //\r
+          return EFI_SUCCESS;\r
+        }\r
+\r
+        NewPe32BaseAddress = *BsBase;\r
+        BaseToUpdate = BsBase;\r
+        break;\r
+\r
+      default:\r
+        //\r
+        // Not supported file type\r
+        //\r
+        return EFI_SUCCESS;\r
+    }\r
+\r

     //\r
     // Allocate a buffer for the image to be loaded into.\r
     //\r
     //\r
     // Allocate a buffer for the image to be loaded into.\r
     //\r


@@ -706,7 +775,6 @@ Returns:
                                               ImageContext.PeCoffHeaderOffset);\r
       if (PeHdr->OptionalHeader.SectionAlignment != PeHdr->OptionalHeader.FileAlignment) {\r
         Error (NULL, 0, 0, "Section-Alignment and File-Alignment does not match", FileGuidString);\r
                                               ImageContext.PeCoffHeaderOffset);\r
       if (PeHdr->OptionalHeader.SectionAlignment != PeHdr->OptionalHeader.FileAlignment) {\r
         Error (NULL, 0, 0, "Section-Alignment and File-Alignment does not match", FileGuidString);\r

-               _asm int 3;\r

         free ((VOID *) MemoryImagePointer);\r
         return EFI_ABORTED;\r
       }\r
         free ((VOID *) MemoryImagePointer);\r
         return EFI_ABORTED;\r
       }\r


@@ -743,21 +811,35 @@ Returns:
       free ((VOID *) MemoryImagePointer);\r
       return EFI_ABORTED;\r
     }\r
       free ((VOID *) MemoryImagePointer);\r
       return EFI_ABORTED;\r
     }\r

+\r
+    //\r
+    // Update BASE address\r
+    //\r
+    fprintf (\r
+      LogFile,\r
+      "%s %016I64X %s\n",\r
+      FileGuidString,\r
+      ImageContext.DestinationAddress,\r
+      ImageContext.PdbPointer == NULL ? "*" : ImageContext.PdbPointer\r
+      );\r
+    *BaseToUpdate += EFI_SIZE_TO_PAGES (ImageContext.ImageSize) * EFI_PAGE_SIZE;\r
+\r

     //\r
     // Since we may have updated the Codeview RVA, we need to insure the PE\r
     // header indicates the image is large enough to contain the Codeview data\r
     // so it will be loaded properly later if the PEIM is reloaded into memory...\r
     //\r
     PeHdr = (VOID *) ((UINTN) ImageAddress + ImageContext.PeCoffHeaderOffset);\r
     //\r
     // Since we may have updated the Codeview RVA, we need to insure the PE\r
     // header indicates the image is large enough to contain the Codeview data\r
     // so it will be loaded properly later if the PEIM is reloaded into memory...\r
     //\r
     PeHdr = (VOID *) ((UINTN) ImageAddress + ImageContext.PeCoffHeaderOffset);\r

+    PePlusHdr = (EFI_IMAGE_NT_HEADERS64*)PeHdr;\r

     if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_IA32) {\r
     if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_IA32) {\r

-      PeHdrSizeOfImage     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER32 *) &PeHdr->OptionalHeader).SizeOfImage);\r
-      PeHdrChecksum        = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER32 *) &PeHdr->OptionalHeader).CheckSum);\r
+      PeHdrSizeOfImage  = (UINT32 *) (&(PeHdr->OptionalHeader).SizeOfImage);\r
+      PeHdrChecksum     = (UINT32 *) (&(PeHdr->OptionalHeader).CheckSum);\r

     } else if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_IA64) {\r
     } else if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_IA64) {\r

-      PeHdrSizeOfImage     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).SizeOfImage);\r
-      PeHdrChecksum        = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).CheckSum);\r
+      PeHdrSizeOfImage  = (UINT32 *) (&(PePlusHdr->OptionalHeader).SizeOfImage);\r
+      PeHdrChecksum     = (UINT32 *) (&(PePlusHdr->OptionalHeader).CheckSum);\r

     } else if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_X64) {\r
     } else if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_X64) {\r

-      PeHdrSizeOfImage     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).SizeOfImage);\r
-      PeHdrChecksum        = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).CheckSum);\r
+      PeHdrSizeOfImage  = (UINT32 *) (&(PePlusHdr->OptionalHeader).SizeOfImage);\r
+      PeHdrChecksum     = (UINT32 *) (&(PePlusHdr->OptionalHeader).CheckSum);\r

     } else {\r
       Error (\r
         NULL,\r
     } else {\r
       Error (\r
         NULL,\r


@@ -780,24 +862,6 @@ Returns:
     }\r
 \r
     memcpy (CurrentPe32Section.Pe32Section + 1, (VOID *) MemoryImagePointerAligned, (UINT32) ImageSize);\r
     }\r
 \r
     memcpy (CurrentPe32Section.Pe32Section + 1, (VOID *) MemoryImagePointerAligned, (UINT32) ImageSize);\r

-    \r
-    //\r
-    // Get EntryPoint in Flash Region.\r
-    //\r
-    EntryPoint = NewPe32BaseAddress + EntryPoint - ImageAddress;\r
-\r
-    //\r
-    // If a map file was selected output mapping information for any file that\r
-    // was rebased.\r
-    //\r
-    if (MapFile != NULL) {\r
-      fprintf (MapFile, "PE32 File: %s Base:%08lx", FileGuidString, BaseAddress);\r
-      fprintf (MapFile, " EntryPoint:%08lx", EntryPoint);\r
-      if (ImageContext.PdbPointer != NULL) {\r
-        fprintf (MapFile, " FileName: %s", ImageContext.PdbPointer);\r
-      }\r
-      fprintf (MapFile, "\n");\r
-    }\r

 \r
     free ((VOID *) MemoryImagePointer);\r
 \r
 \r
     free ((VOID *) MemoryImagePointer);\r
 \r


@@ -833,6 +897,20 @@ Returns:
       *(EFI_FFS_FILE_TAIL *) (((UINTN) FfsFile + GetLength (FfsFile->Size) - sizeof (EFI_FFS_FILE_TAIL))) = TailValue;\r
     }\r
   }\r
       *(EFI_FFS_FILE_TAIL *) (((UINTN) FfsFile + GetLength (FfsFile->Size) - sizeof (EFI_FFS_FILE_TAIL))) = TailValue;\r
     }\r
   }\r

+\r
+  if ((Flags & 1) == 0 || (\r
+      FfsFile->Type != EFI_FV_FILETYPE_SECURITY_CORE &&\r
+      FfsFile->Type != EFI_FV_FILETYPE_PEI_CORE &&\r
+\r
+      FfsFile->Type != EFI_FV_FILETYPE_PEIM &&\r
+      FfsFile->Type != EFI_FV_FILETYPE_COMBINED_PEIM_DRIVER\r
+      )) {\r
+    //\r
+    // Only XIP code may have a TE section\r
+    //\r
+    return EFI_SUCCESS;\r
+  }\r
+\r

   //\r
   // Now process TE sections\r
   //\r
   //\r
   // Now process TE sections\r
   //\r


@@ -842,15 +920,13 @@ Returns:
       break;\r
     }\r
 \r
       break;\r
     }\r
 \r

-    FoundCount++;\r
-\r

     //\r
     // Calculate the TE base address, the FFS file base plus the offset of the TE section less the size stripped off\r
     // by GenTEImage\r
     //\r
     TEImageHeader = (EFI_TE_IMAGE_HEADER *) ((UINT8 *) CurrentPe32Section.Pe32Section + sizeof (EFI_COMMON_SECTION_HEADER));\r
 \r
     //\r
     // Calculate the TE base address, the FFS file base plus the offset of the TE section less the size stripped off\r
     // by GenTEImage\r
     //\r
     TEImageHeader = (EFI_TE_IMAGE_HEADER *) ((UINT8 *) CurrentPe32Section.Pe32Section + sizeof (EFI_COMMON_SECTION_HEADER));\r
 \r

-    NewPe32BaseAddress = ((UINT32) BaseAddress) +\r
+    NewPe32BaseAddress = ((UINT32) XipBase) +\r

       (\r
         (UINTN) CurrentPe32Section.Pe32Section +\r
         sizeof (EFI_COMMON_SECTION_HEADER) +\r
       (\r
         (UINTN) CurrentPe32Section.Pe32Section +\r
         sizeof (EFI_COMMON_SECTION_HEADER) +\r


@@ -881,6 +957,7 @@ Returns:
     DosHeader->e_magic = EFI_IMAGE_DOS_SIGNATURE;\r
     *(UINT32 *) (TEBuffer + 0x3C) = 0x40;\r
     PeHdr = (EFI_IMAGE_NT_HEADERS *) (TEBuffer + 0x40);\r
     DosHeader->e_magic = EFI_IMAGE_DOS_SIGNATURE;\r
     *(UINT32 *) (TEBuffer + 0x3C) = 0x40;\r
     PeHdr = (EFI_IMAGE_NT_HEADERS *) (TEBuffer + 0x40);\r

+    PePlusHdr = (EFI_IMAGE_NT_HEADERS64*)PeHdr;\r

     PeHdr->Signature = EFI_IMAGE_NT_SIGNATURE;\r
     PeHdr->FileHeader.Machine = TEImageHeader->Machine;\r
     PeHdr->FileHeader.NumberOfSections = TEImageHeader->NumberOfSections;\r
     PeHdr->Signature = EFI_IMAGE_NT_SIGNATURE;\r
     PeHdr->FileHeader.Machine = TEImageHeader->Machine;\r
     PeHdr->FileHeader.NumberOfSections = TEImageHeader->NumberOfSections;\r


@@ -890,39 +967,89 @@ Returns:
     // the 0x40 bytes for our DOS header.\r
     //\r
     PeHdr->FileHeader.SizeOfOptionalHeader = (UINT16) (TEImageHeader->StrippedSize - 0x40 - sizeof (UINT32) - sizeof (EFI_IMAGE_FILE_HEADER));\r
     // the 0x40 bytes for our DOS header.\r
     //\r
     PeHdr->FileHeader.SizeOfOptionalHeader = (UINT16) (TEImageHeader->StrippedSize - 0x40 - sizeof (UINT32) - sizeof (EFI_IMAGE_FILE_HEADER));\r

-    PeHdr->OptionalHeader.ImageBase = (UINTN) (TEImageHeader->ImageBase - TEImageHeader->StrippedSize + sizeof (EFI_TE_IMAGE_HEADER));\r
-    PeHdr->OptionalHeader.AddressOfEntryPoint = TEImageHeader->AddressOfEntryPoint;\r
-    PeHdr->OptionalHeader.BaseOfCode  = TEImageHeader->BaseOfCode;\r
-    PeHdr->OptionalHeader.SizeOfImage = Pe32ImageSize;\r
-    PeHdr->OptionalHeader.Subsystem   = TEImageHeader->Subsystem;\r
-    PeHdr->OptionalHeader.SizeOfImage = Pe32ImageSize;\r
-    PeHdr->OptionalHeader.SizeOfHeaders = TEImageHeader->StrippedSize + TEImageHeader->NumberOfSections *\r
-    sizeof (EFI_IMAGE_SECTION_HEADER) - 12;\r
-\r
-    //\r
-    // Set NumberOfRvaAndSizes in the optional header to what we had available in the original image\r
-    //\r
-    if ((TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress != 0) ||\r
-        (TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_BASERELOC].Size != 0)\r
-        ) {\r
-      PeHdr->OptionalHeader.NumberOfRvaAndSizes = EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC + 1;\r
-      PeHdr->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress = TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress;\r
-      PeHdr->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC].Size = TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_BASERELOC].Size;\r
+    if (TEImageHeader->Machine == EFI_IMAGE_MACHINE_IA32) {\r
+      PeHdr->OptionalHeader.Magic = EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC;\r
+    } else if (TEImageHeader->Machine == EFI_IMAGE_MACHINE_IA64) {\r
+      PePlusHdr->OptionalHeader.Magic = EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC;\r
+    } else if (TEImageHeader->Machine == EFI_IMAGE_MACHINE_X64) {\r
+      PePlusHdr->OptionalHeader.Magic = EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC;\r
+    } else {\r
+      Error (\r
+        NULL,\r
+        0,\r
+        0,\r
+        "unknown machine type in TE image",\r
+        "machine type=0x%X, file=%s",\r
+        (UINT32) TEImageHeader->Machine,\r
+        FileGuidString\r
+        );\r
+      free (TEBuffer);\r
+      return EFI_ABORTED;\r

     }\r
 \r
     }\r
 \r

-    if ((TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].VirtualAddress != 0) ||\r
-        (TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].Size != 0)\r
-        ) {\r
-      PeHdr->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_DEBUG].VirtualAddress = TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].VirtualAddress;\r
-      PeHdr->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_DEBUG].Size = TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].Size;\r
-      if (PeHdr->OptionalHeader.NumberOfRvaAndSizes < EFI_IMAGE_DIRECTORY_ENTRY_DEBUG + 1) {\r
-        PeHdr->OptionalHeader.NumberOfRvaAndSizes = EFI_IMAGE_DIRECTORY_ENTRY_DEBUG + 1;\r
+    if (PeHdr->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+      PeHdr->OptionalHeader.ImageBase     = (UINTN) (TEImageHeader->ImageBase - TEImageHeader->StrippedSize + sizeof (EFI_TE_IMAGE_HEADER));\r
+      PeHdr->OptionalHeader.SizeOfImage   = Pe32ImageSize;\r
+      PeHdr->OptionalHeader.Subsystem     = TEImageHeader->Subsystem;\r
+      PeHdr->OptionalHeader.SizeOfHeaders = TEImageHeader->StrippedSize + TEImageHeader->NumberOfSections *\r
+                                            sizeof (EFI_IMAGE_SECTION_HEADER) - 12;\r
+\r
+      //\r
+      // Set NumberOfRvaAndSizes in the optional header to what we had available in the original image\r
+      //\r
+      if ((TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress != 0) ||\r
+          (TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_BASERELOC].Size != 0)\r
+          ) {\r
+        PeHdr->OptionalHeader.NumberOfRvaAndSizes = EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC + 1;\r
+        PeHdr->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress = TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress;\r
+        PeHdr->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC].Size = TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_BASERELOC].Size;\r

       }\r
       }\r

+\r
+      if ((TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].VirtualAddress != 0) ||\r
+          (TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].Size != 0)\r
+          ) {\r
+        PeHdr->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_DEBUG].VirtualAddress = TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].VirtualAddress;\r
+        PeHdr->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_DEBUG].Size = TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].Size;\r
+        if (PeHdr->OptionalHeader.NumberOfRvaAndSizes < EFI_IMAGE_DIRECTORY_ENTRY_DEBUG + 1) {\r
+          PeHdr->OptionalHeader.NumberOfRvaAndSizes = EFI_IMAGE_DIRECTORY_ENTRY_DEBUG + 1;\r
+        }\r
+      }\r
+      //\r
+      // NOTE: These values are defaults, and should be verified to be correct in the GenTE utility\r
+      //\r
+      PeHdr->OptionalHeader.SectionAlignment = 0x10;\r
+    } else {\r
+      PePlusHdr->OptionalHeader.ImageBase     = (UINTN) (TEImageHeader->ImageBase - TEImageHeader->StrippedSize + sizeof (EFI_TE_IMAGE_HEADER));\r
+      PePlusHdr->OptionalHeader.SizeOfImage   = Pe32ImageSize;\r
+      PePlusHdr->OptionalHeader.Subsystem     = TEImageHeader->Subsystem;\r
+      PePlusHdr->OptionalHeader.SizeOfHeaders = TEImageHeader->StrippedSize + TEImageHeader->NumberOfSections *\r
+                                                sizeof (EFI_IMAGE_SECTION_HEADER) - 12;\r
+\r
+      //\r
+      // Set NumberOfRvaAndSizes in the optional header to what we had available in the original image\r
+      //\r
+      if ((TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress != 0) ||\r
+          (TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_BASERELOC].Size != 0)\r
+          ) {\r
+        PePlusHdr->OptionalHeader.NumberOfRvaAndSizes = EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC + 1;\r
+        PePlusHdr->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress = TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress;\r
+        PePlusHdr->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC].Size = TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_BASERELOC].Size;\r
+      }\r
+\r
+      if ((TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].VirtualAddress != 0) ||\r
+          (TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].Size != 0)\r
+          ) {\r
+        PePlusHdr->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_DEBUG].VirtualAddress = TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].VirtualAddress;\r
+        PePlusHdr->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_DEBUG].Size = TEImageHeader->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].Size;\r
+        if (PePlusHdr->OptionalHeader.NumberOfRvaAndSizes < EFI_IMAGE_DIRECTORY_ENTRY_DEBUG + 1) {\r
+          PePlusHdr->OptionalHeader.NumberOfRvaAndSizes = EFI_IMAGE_DIRECTORY_ENTRY_DEBUG + 1;\r
+        }\r
+      }\r
+      //\r
+      // NOTE: These values are defaults, and should be verified to be correct in the GenTE utility\r
+      //\r
+      PePlusHdr->OptionalHeader.SectionAlignment = 0x10;\r

     }\r
     }\r

-    //\r
-    // NOTE: These values are defaults, and should be verified to be correct in the GenTE utility\r
-    //\r
-    PeHdr->OptionalHeader.SectionAlignment = 0x10;\r

 \r
     //\r
     // Copy the rest of the image to its original offset\r
 \r
     //\r
     // Copy the rest of the image to its original offset\r


@@ -969,11 +1096,6 @@ Returns:
       free ((VOID *) MemoryImagePointer);\r
       return Status;\r
     }\r
       free ((VOID *) MemoryImagePointer);\r
       return Status;\r
     }\r

-    \r
-    //\r
-    // Check if section-alignment and file-alignment match or not\r
-       // BUGBUG: TE Image Header lack section-alignment and file-alignment info\r
-    //\r

 \r
     ImageContext.DestinationAddress = NewPe32BaseAddress;\r
     Status                          = PeCoffLoaderRelocateImage (&ImageContext);\r
 \r
     ImageContext.DestinationAddress = NewPe32BaseAddress;\r
     Status                          = PeCoffLoaderRelocateImage (&ImageContext);\r


@@ -994,12 +1116,16 @@ Returns:
     // so it will be loaded properly later if the PEIM is reloaded into memory...\r
     //\r
     PeHdr = (VOID *) ((UINTN) ImageAddress + ImageContext.PeCoffHeaderOffset);\r
     // so it will be loaded properly later if the PEIM is reloaded into memory...\r
     //\r
     PeHdr = (VOID *) ((UINTN) ImageAddress + ImageContext.PeCoffHeaderOffset);\r

+    PePlusHdr = (EFI_IMAGE_NT_HEADERS64*)PeHdr;\r

     if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_IA32) {\r
     if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_IA32) {\r

-      PeHdrSizeOfImage     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER32 *) &PeHdr->OptionalHeader).SizeOfImage);\r
-      PeHdrChecksum        = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER32 *) &PeHdr->OptionalHeader).CheckSum);\r
+      PeHdrSizeOfImage  = (UINT32 *) (&(PeHdr->OptionalHeader).SizeOfImage);\r
+      PeHdrChecksum     = (UINT32 *) (&(PeHdr->OptionalHeader).CheckSum);\r

     } else if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_IA64) {\r
     } else if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_IA64) {\r

-      PeHdrSizeOfImage     = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).SizeOfImage);\r
-      PeHdrChecksum        = (UINT32 *) (&(*(EFI_IMAGE_OPTIONAL_HEADER64 *) &PeHdr->OptionalHeader).CheckSum);\r
+      PeHdrSizeOfImage  = (UINT32 *) (&(PePlusHdr->OptionalHeader).SizeOfImage);\r
+      PeHdrChecksum     = (UINT32 *) (&(PePlusHdr->OptionalHeader).CheckSum);\r
+    } else if (PeHdr->FileHeader.Machine == EFI_IMAGE_MACHINE_X64) {\r
+      PeHdrSizeOfImage  = (UINT32 *) (&(PePlusHdr->OptionalHeader).SizeOfImage);\r
+      PeHdrChecksum     = (UINT32 *) (&(PePlusHdr->OptionalHeader).CheckSum);\r

     } else {\r
       Error (\r
         NULL,\r
     } else {\r
       Error (\r
         NULL,\r


@@ -1029,27 +1155,6 @@ Returns:
       GetLength (CurrentPe32Section.Pe32Section->CommonHeader.Size) - sizeof (EFI_PE32_SECTION) -\r
       sizeof (EFI_TE_IMAGE_HEADER)\r
       );\r
       GetLength (CurrentPe32Section.Pe32Section->CommonHeader.Size) - sizeof (EFI_PE32_SECTION) -\r
       sizeof (EFI_TE_IMAGE_HEADER)\r
       );\r

-    \r
-    //\r
-    // Get EntryPoint in Flash Region.\r
-    //\r
-    EntryPoint = NewPe32BaseAddress + EntryPoint - ImageAddress;\r
-\r
-    //\r
-    // If a map file was selected output mapping information for any file that\r
-    // was rebased.\r
-    //\r
-    if (MapFile != NULL) {\r
-      fprintf (MapFile, "TE   File: %s Base:%08lx", FileGuidString, BaseAddress);\r
-      fprintf (MapFile, " EntryPoint:%08lx", EntryPoint);\r
-      if (ImageContext.PdbPointer != NULL) {\r
-        fprintf (MapFile, " FileName: %s", ImageContext.PdbPointer);\r
-      }\r
-      fprintf (MapFile, "\n");\r
-    }\r
-\r
-    free ((VOID *) MemoryImagePointer);\r
-    free (TEBuffer);\r

     if (FfsFile->Attributes & FFS_ATTRIB_TAIL_PRESENT) {\r
       TailSize = sizeof (EFI_FFS_FILE_TAIL);\r
     } else {\r
     if (FfsFile->Attributes & FFS_ATTRIB_TAIL_PRESENT) {\r
       TailSize = sizeof (EFI_FFS_FILE_TAIL);\r
     } else {\r


@@ -1080,18 +1185,22 @@ Returns:
       TailValue = (EFI_FFS_FILE_TAIL) (~(FfsFile->IntegrityCheck.TailReference));\r
       *(EFI_FFS_FILE_TAIL *) (((UINTN) FfsFile + GetLength (FfsFile->Size) - sizeof (EFI_FFS_FILE_TAIL))) = TailValue;\r
     }\r
       TailValue = (EFI_FFS_FILE_TAIL) (~(FfsFile->IntegrityCheck.TailReference));\r
       *(EFI_FFS_FILE_TAIL *) (((UINTN) FfsFile + GetLength (FfsFile->Size) - sizeof (EFI_FFS_FILE_TAIL))) = TailValue;\r
     }\r

+\r
+    fprintf (\r
+      LogFile,\r
+      "%s %016I64X %s\n",\r
+      FileGuidString,\r
+      ImageContext.DestinationAddress,\r
+      ImageContext.PdbPointer == NULL ? "*" : ImageContext.PdbPointer\r
+      );\r
+\r
+    //\r
+    // Free buffers\r
+    //\r
+    free ((VOID *) MemoryImagePointer);\r
+    free (TEBuffer);\r

   }\r
   }\r

-  //\r
-  // If we found no files, then emit an error if no compressed sections either\r
-  //\r
-  if (FoundCount == 0) {\r
-    Status = GetSectionByType (FfsFile, EFI_SECTION_COMPRESSION, Index, &CurrentPe32Section);\r
-    if (EFI_ERROR (Status)) {\r
-      Error (NULL, 0, 0, "no PE32, TE, nor compressed section found in FV file", FileGuidString);\r
-      return EFI_NOT_FOUND;\r
-    }\r
-  }\r
-  \r
+\r

   return EFI_SUCCESS;\r
 }\r
 \r
   return EFI_SUCCESS;\r
 }\r
 \r