]> git.proxmox.com Git - mirror_edk2.git/blobdiff - Tools/CCode/Source/PeiRebase/PeiRebaseExe.c
Merged back the Yizhong's fix which was overwritten by check-in of r2157,2158.
[mirror_edk2.git] / Tools / CCode / Source / PeiRebase / PeiRebaseExe.c
index 8537d69f68fd9b0d13a6003ca40f5a426f291658..5f7dfc25c71cd47a3e78d2c8baedd4e5fa87bc84 100644 (file)
@@ -1,8 +1,8 @@
 /*++\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
@@ -59,23 +59,21 @@ Arguments:
 \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
-  MapFileName     The name of the map file of relocation info.\r
 \r
   Arguments come in pair in any order.\r
-    -I InputFileName \r
+    -I InputFileName\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
-  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
@@ -84,14 +82,13 @@ Returns:
 {\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
-  FILE                        *MapFile;\r
+  FILE                        *LogFile;\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
-  EFI_PHYSICAL_ADDRESS        CurrentFileBaseAddress;\r
-  CHAR8                       InfFileName[_MAX_PATH];\r
-  CHAR8                       *InfFileImage;\r
-  UINTN                       InfFileSize;\r
   MEMORY_FILE                 InfMemoryFile;\r
+  CHAR8                       StringBuffer[0x100];\r
 \r
   ErasePolarity = FALSE;\r
   //\r
@@ -121,21 +115,17 @@ Returns:
   //\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
-  MapFile           = NULL;\r
+  LogFile           = NULL;\r
   FvImage           = NULL;\r
-  InfFileImage      = NULL;\r
-  InfFileSize       = 0;\r
-  strcpy (InfFileName, "");\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
-    }    \r
+    }\r
     //\r
     // Determine argument to read\r
     //\r
@@ -174,8 +164,8 @@ Returns:
 \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
@@ -183,79 +173,86 @@ Returns:
       }\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
-      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
-        BaseAddressSet = TRUE;\r
-      } else {\r
+      Status = AsciiStringToUint64 (argv[Index + 1], FALSE, &XipBase);\r
+      if (EFI_ERROR (Status)) {\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
+\r
+      BaseTypes |= 1;\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
-        BaseAddressSet = TRUE;\r
-      } else {\r
+      Status = AsciiStringToUint64 (argv[Index + 1], FALSE, &BsBase);\r
+      if (EFI_ERROR (Status)) {\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
+\r
+      BaseTypes |= 2;\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
-        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
+\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
@@ -265,18 +262,6 @@ Returns:
       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
@@ -285,6 +270,16 @@ Returns:
     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
@@ -330,8 +325,14 @@ Returns:
     //\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
@@ -359,7 +360,6 @@ Returns:
 \r
       goto Finish;\r
     }\r
-\r
     //\r
     // Get the next file\r
     //\r
@@ -399,8 +399,8 @@ Finish:
     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
@@ -420,7 +420,7 @@ ReadHeader (
 \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
@@ -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
-    \r
+\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
@@ -539,38 +539,36 @@ Returns:
 --*/\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
-  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 ("    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
-  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
-  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
-  \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
-  MapFile           Optional file to dump relocation information into\r
 \r
 Returns:\r
 \r
@@ -590,20 +588,21 @@ Returns:
   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
-  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                                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_PHYSICAL_ADDRESS                  *BaseToUpdate;\r
 \r
   //\r
   // Verify input parameters\r
@@ -611,7 +610,6 @@ Returns:
   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
@@ -622,7 +620,6 @@ Returns:
   } else {\r
     TailSize = 0;\r
   }\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
-  memset (&ImageContext, 0, sizeof (ImageContext));\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
-  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
-  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
-    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
@@ -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
-\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
+\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
@@ -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
-               _asm int 3;\r
         free ((VOID *) MemoryImagePointer);\r
         return EFI_ABORTED;\r
       }\r
@@ -743,21 +811,35 @@ Returns:
       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
+    PePlusHdr = (EFI_IMAGE_NT_HEADERS64*)PeHdr;\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
-      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 *) (&(*(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
@@ -780,24 +862,6 @@ Returns:
     }\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
@@ -833,6 +897,20 @@ Returns:
       *(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
@@ -842,15 +920,13 @@ Returns:
       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
-    NewPe32BaseAddress = ((UINT32) BaseAddress) +\r
+    NewPe32BaseAddress = ((UINT32) XipBase) +\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
+    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
@@ -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
-    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
-    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
+      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
-    // 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
@@ -969,11 +1096,6 @@ Returns:
       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
@@ -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
+    PePlusHdr = (EFI_IMAGE_NT_HEADERS64*)PeHdr;\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
-      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
@@ -1029,27 +1155,6 @@ Returns:
       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
@@ -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
+\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
-  // 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