[RFC PATCH 08/11] MdeModulePkg/DxeIpl: Relocate and remap XIP capable DXE drivers

["Ard Biesheuvel" <ardb@kernel.org>]

Before handing over to the DXE core, iterate over all known FFS files
and find the ones that can execute in place. These files are then
relocated in place and mapped with restricted permissions, allowing the
DXE core to dispatch them without the need to perform any manipulation
of the file contents or the page table permissions.

Signed-off-by: Ard Biesheuvel <ardb@kernel.org>
---
 MdeModulePkg/Core/DxeIplPeim/DxeIpl.inf |   1 +
 MdeModulePkg/Core/DxeIplPeim/DxeLoad.c  | 196 ++++++++++++++++++++
 2 files changed, 197 insertions(+)

diff --git a/MdeModulePkg/Core/DxeIplPeim/DxeIpl.inf b/MdeModulePkg/Core/DxeIplPeim/DxeIpl.inf
index f1990eac77607854..60112100df78b396 100644
--- a/MdeModulePkg/Core/DxeIplPeim/DxeIpl.inf
+++ b/MdeModulePkg/Core/DxeIplPeim/DxeIpl.inf
@@ -65,6 +65,7 @@ [LibraryClasses]
   PeimEntryPoint
   DebugLib
   DebugAgentLib
+  PeCoffLib
   PeiServicesTablePointerLib
   PerformanceLib
 
diff --git a/MdeModulePkg/Core/DxeIplPeim/DxeLoad.c b/MdeModulePkg/Core/DxeIplPeim/DxeLoad.c
index 2c19f1a507baf34a..1f20db1faffbd1d2 100644
--- a/MdeModulePkg/Core/DxeIplPeim/DxeLoad.c
+++ b/MdeModulePkg/Core/DxeIplPeim/DxeLoad.c
@@ -10,6 +10,9 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 
 #include "DxeIpl.h"
 
+#include <Library/PeCoffLib.h>
+#include <Ppi/MemoryAttribute.h>
+
 //
 // Module Globals used in the DXE to PEI hand off
 // These must be module globals, so the stack can be switched
@@ -228,6 +231,197 @@ ValidateMemoryTypeInfoVariable (
   return TRUE;
 }
 
+/**
+  Support routine for the PE/COFF Loader that reads a buffer from a PE/COFF file.
+  The function is used for XIP code to have optimized memory copy.
+
+  @param FileHandle      The handle to the PE/COFF file
+  @param FileOffset      The offset, in bytes, into the file to read
+  @param ReadSize        The number of bytes to read from the file starting at FileOffset
+  @param Buffer          A pointer to the buffer to read the data into.
+
+  @return EFI_SUCCESS    ReadSize bytes of data were read into Buffer from the
+                         PE/COFF file starting at FileOffset
+
+**/
+STATIC
+EFI_STATUS
+EFIAPI
+PeiImageRead (
+  IN     VOID   *FileHandle,
+  IN     UINTN  FileOffset,
+  IN     UINTN  *ReadSize,
+  OUT    VOID   *Buffer
+  )
+{
+  CHAR8  *Destination8;
+  CHAR8  *Source8;
+
+  Destination8 = Buffer;
+  Source8      = (CHAR8 *)((UINTN)FileHandle + FileOffset);
+  if (Destination8 != Source8) {
+    CopyMem (Destination8, Source8, *ReadSize);
+  }
+
+  return EFI_SUCCESS;
+}
+
+STATIC
+VOID
+RemapImage (
+  IN  EDKII_MEMORY_ATTRIBUTE_PPI          *MemoryPpi,
+  IN  CONST PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  )
+{
+  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
+  EFI_IMAGE_SECTION_HEADER             *Section;
+  PHYSICAL_ADDRESS                     SectionAddress;
+  EFI_STATUS                           Status;
+  UINT64                               Permissions;
+  UINTN                                Index;
+
+  if (MemoryPpi == NULL) {
+    return;
+  }
+
+  Hdr.Union = (EFI_IMAGE_OPTIONAL_HEADER_UNION *)((UINT8 *)ImageContext->Handle +
+                                                  ImageContext->PeCoffHeaderOffset);
+  ASSERT (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE);
+
+  Section = (EFI_IMAGE_SECTION_HEADER *)((UINT8 *)Hdr.Union + sizeof (UINT32) +
+                                         sizeof (EFI_IMAGE_FILE_HEADER) +
+                                         Hdr.Pe32->FileHeader.SizeOfOptionalHeader
+                                         );
+
+  for (Index = 0; Index < Hdr.Pe32->FileHeader.NumberOfSections; Index++) {
+    SectionAddress = ImageContext->ImageAddress + Section[Index].VirtualAddress;
+    Permissions    = 0;
+
+    if ((Section[Index].Characteristics & EFI_IMAGE_SCN_MEM_WRITE) == 0) {
+      Permissions |= EFI_MEMORY_RO;
+    }
+
+    if ((Section[Index].Characteristics & EFI_IMAGE_SCN_MEM_EXECUTE) == 0) {
+      Permissions |= EFI_MEMORY_XP;
+    }
+
+    Status = MemoryPpi->SetPermissions (
+                          MemoryPpi,
+                          SectionAddress,
+                          Section[Index].Misc.VirtualSize,
+                          Permissions,
+                          Permissions ^ EFI_MEMORY_RO ^ EFI_MEMORY_XP
+                          );
+    ASSERT_EFI_ERROR (Status);
+  }
+}
+
+STATIC
+VOID
+RelocateAndRemapDriversInPlace (
+  VOID
+  )
+{
+  EFI_STATUS                    Status;
+  UINTN                         Instance;
+  EFI_PEI_FV_HANDLE             VolumeHandle;
+  EFI_PEI_FILE_HANDLE           FileHandle;
+  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;
+  UINT32                        AuthenticationState;
+  EDKII_MEMORY_ATTRIBUTE_PPI    *MemoryPpi;
+
+  MemoryPpi = NULL;
+  PeiServicesLocatePpi (&gEdkiiMemoryAttributePpiGuid, 0, NULL, (VOID **)&MemoryPpi);
+
+  Instance = 0;
+  do {
+    //
+    // Traverse all firmware volume instances
+    //
+    Status = PeiServicesFfsFindNextVolume (Instance, &VolumeHandle);
+    if (Status == EFI_NOT_FOUND) {
+      return;
+    }
+
+    ASSERT_EFI_ERROR (Status);
+
+    FileHandle = NULL;
+    do {
+      Status = PeiServicesFfsFindNextFile (
+                 EFI_FV_FILETYPE_DRIVER,
+                 VolumeHandle,
+                 &FileHandle);
+      if (Status == EFI_NOT_FOUND) {
+        break;
+      }
+
+      ASSERT_EFI_ERROR (Status);
+
+      ZeroMem (&ImageContext, sizeof (ImageContext));
+
+      Status = PeiServicesFfsFindSectionData3 (
+                 EFI_SECTION_PE32,
+                 0,
+                 FileHandle,
+                 &ImageContext.Handle,
+                 &AuthenticationState
+                 );
+      if (Status == EFI_NOT_FOUND) {
+        continue;
+      }
+
+      ASSERT_EFI_ERROR (Status);
+
+      ImageContext.ImageRead = PeiImageRead;
+      Status                 = PeCoffLoaderGetImageInfo (&ImageContext);
+      if (EFI_ERROR (Status)) {
+        ASSERT_EFI_ERROR (Status);
+        continue;
+      }
+
+      ImageContext.ImageAddress = (PHYSICAL_ADDRESS)(UINTN)ImageContext.Handle;
+      if ((ImageContext.ImageAddress & (ImageContext.SectionAlignment - 1)) != 0) {
+        DEBUG ((DEBUG_VERBOSE, "%a: skip PE image at %p\n", __func__, ImageContext.Handle));
+        continue;
+      }
+
+      DEBUG ((
+        DEBUG_INFO,
+        "%a: relocate PE image at %p for execution in place\n",
+        __func__,
+        ImageContext.Handle
+        ));
+
+      //
+      // 'Load' the image in-place - this just performs a sanity check on
+      // the PE metadata but does not actually move any data
+      //
+      Status = PeCoffLoaderLoadImage (&ImageContext);
+      if (EFI_ERROR (Status)) {
+        ASSERT_EFI_ERROR (Status);
+        continue;
+      }
+
+      //
+      // Relocate this driver in place
+      //
+      Status = PeCoffLoaderRelocateImage (&ImageContext);
+      if (EFI_ERROR (Status)) {
+        ASSERT_EFI_ERROR (Status);
+        continue;
+      }
+
+      //
+      // Apply section permissions to the page tables
+      //
+      RemapImage (MemoryPpi, &ImageContext);
+
+    } while (TRUE);
+
+    Instance++;
+  } while (TRUE);
+}
+
 /**
    Main entry point to last PEIM.
 
@@ -436,6 +630,8 @@ DxeLoadCore (
     DxeCoreEntryPoint
     );
 
+  RelocateAndRemapDriversInPlace ();
+
   //
   // Report Status Code EFI_SW_PEI_PC_HANDOFF_TO_NEXT
   //
-- 
2.39.2