[RFC 03/13] MdePkg/BasePeCoffLib: Add API to keep track of relocation range

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

Add a library call to obtain the start and end of the region covered by
relocation fixups. This will be used in a future patch to limit the
range of memory that needs to be remapped with read-write-execute
permissions at ExitBootServices() time.

Signed-off-by: Ard Biesheuvel <ardb@kernel.org>
---
 MdePkg/Include/Library/PeCoffLib.h        | 25 ++++++
 MdePkg/Library/BasePeCoffLib/BasePeCoff.c | 83 +++++++++++++++++++-
 2 files changed, 107 insertions(+), 1 deletion(-)

diff --git a/MdePkg/Include/Library/PeCoffLib.h b/MdePkg/Include/Library/PeCoffLib.h
index b45879453785..3706c8a4e858 100644
--- a/MdePkg/Include/Library/PeCoffLib.h
+++ b/MdePkg/Include/Library/PeCoffLib.h
@@ -382,4 +382,29 @@ PeCoffLoaderUnloadImage (
   IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
   );
 
+/**
+  Retrieve the range subject to relocation fixups from the recorded fixup data
+  of a runtime image
+
+  @param       ImageBase           The base address of a PE/COFF image that has been loaded
+                                   and relocated into system memory.
+  @param       VirtImageBase       The request virtual address that the PE/COFF image is to
+                                   be fixed up for.
+  @param       ImageSize           The size, in bytes, of the PE/COFF image.
+  @param       RelocationData      A pointer to the relocation data that was collected when the
+                                   PE/COFF image was relocated using PeCoffLoaderRelocateImage().
+  @param[out]  RelocationRangeMin  The start of the relocated range.
+  @param[out]  RelocationRangeMax  The end of the relocated range.
+
+**/
+VOID
+EFIAPI
+PeCoffLoaderGetRelocationRange (
+  IN  PHYSICAL_ADDRESS  ImageBase,
+  IN  UINTN             ImageSize,
+  IN  VOID              *RelocationData,
+  OUT PHYSICAL_ADDRESS  *RelocationRangeMin,
+  OUT PHYSICAL_ADDRESS  *RelocationRangeMax
+  );
+
 #endif
diff --git a/MdePkg/Library/BasePeCoffLib/BasePeCoff.c b/MdePkg/Library/BasePeCoffLib/BasePeCoff.c
index 97a8aaf8c73d..10f3d04d2490 100644
--- a/MdePkg/Library/BasePeCoffLib/BasePeCoff.c
+++ b/MdePkg/Library/BasePeCoffLib/BasePeCoff.c
@@ -936,6 +936,8 @@ PeCoffLoaderRelocateImage (
   PHYSICAL_ADDRESS                     BaseAddress;
   UINT32                               NumberOfRvaAndSizes;
   UINT32                               TeStrippedOffset;
+  PHYSICAL_ADDRESS                     *RelocRangeStart;
+  PHYSICAL_ADDRESS                     *RelocRangeEnd;
 
   ASSERT (ImageContext != NULL);
 
@@ -1043,6 +1045,21 @@ PeCoffLoaderRelocateImage (
     // Run the relocation information and apply the fixups
     //
     FixupData = ImageContext->FixupData;
+    if (FixupData != NULL) {
+      FixupData = ALIGN_POINTER (FixupData, sizeof (PHYSICAL_ADDRESS));
+
+      //
+      // Use the first two UINT64s in the fixup data to keep track of the start
+      // and end of the region that is subject to relocation fixups.
+      //
+      RelocRangeStart = (PHYSICAL_ADDRESS *)FixupData;
+      RelocRangeEnd   = RelocRangeStart + 1;
+      FixupData      += 2 * sizeof (PHYSICAL_ADDRESS);
+
+      *RelocRangeStart = MAX_UINT64;
+      *RelocRangeEnd   = 0;
+    }
+
     while ((UINTN)RelocBase < (UINTN)RelocBaseEnd) {
       Reloc = (UINT16 *)((CHAR8 *)RelocBase + sizeof (EFI_IMAGE_BASE_RELOCATION));
       //
@@ -1070,6 +1087,14 @@ PeCoffLoaderRelocateImage (
         return RETURN_LOAD_ERROR;
       }
 
+      //
+      // Capture this page in the recorded relocation range
+      //
+      if (FixupData != NULL) {
+        *RelocRangeStart = MIN (*RelocRangeStart, (UINTN)FixupBase);
+        *RelocRangeEnd   = MAX (*RelocRangeEnd, (UINTN)FixupBase + SIZE_4KB);
+      }
+
       //
       // Run this relocation record
       //
@@ -1470,6 +1495,9 @@ PeCoffLoaderLoadImage (
   //
   ImageContext->FixupData = NULL;
 
+  // Add two UINT64s at the start to carry the min/max of the relocated range
+  ImageContext->FixupDataSize += 2 * sizeof (PHYSICAL_ADDRESS);
+
   //
   // Load the Codeview information if present
   //
@@ -1824,7 +1852,8 @@ PeCoffLoaderRelocateImageForRuntime (
     // by code will not be fixed up, since that would set them back to
     // defaults.
     //
-    FixupData     = RelocationData;
+    FixupData     = ALIGN_POINTER (RelocationData, sizeof (PHYSICAL_ADDRESS));
+    FixupData    += 2 * sizeof (PHYSICAL_ADDRESS);
     RelocBaseOrig = RelocBase;
     while ((UINTN)RelocBase < (UINTN)RelocBaseEnd) {
       //
@@ -1994,3 +2023,55 @@ PeCoffLoaderUnloadImage (
   PeCoffLoaderUnloadImageExtraAction (ImageContext);
   return RETURN_SUCCESS;
 }
+
+/**
+  Retrieve the range subject to relocation fixups from the recorded fixup data
+  of a runtime image
+
+  @param       ImageBase           The base address of a PE/COFF image that has been loaded
+                                   and relocated into system memory.
+  @param       VirtImageBase       The request virtual address that the PE/COFF image is to
+                                   be fixed up for.
+  @param       ImageSize           The size, in bytes, of the PE/COFF image.
+  @param       RelocationData      A pointer to the relocation data that was collected when the
+                                   PE/COFF image was relocated using PeCoffLoaderRelocateImage().
+  @param[out]  RelocationRangeMin  The start of the relocated range.
+  @param[out]  RelocationRangeMax  The end of the relocated range.
+
+**/
+VOID
+EFIAPI
+PeCoffLoaderGetRelocationRange (
+  IN  PHYSICAL_ADDRESS  ImageBase,
+  IN  UINTN             ImageSize,
+  IN  VOID              *RelocationData,
+  OUT PHYSICAL_ADDRESS  *RelocationRangeMin,
+  OUT PHYSICAL_ADDRESS  *RelocationRangeMax
+  )
+{
+  PHYSICAL_ADDRESS  *FixupData;
+
+  if ((RelocationData == NULL) || (ImageBase == 0x0)) {
+    return;
+  }
+
+  FixupData = ALIGN_POINTER (RelocationData, sizeof (PHYSICAL_ADDRESS));
+
+  if ((FixupData[0] == MAX_UINT64) && (FixupData[1] == 0)) {
+    // No fixups recorded
+    *RelocationRangeMin = ImageBase;
+    *RelocationRangeMax = ImageBase;
+    return;
+  }
+
+  if ((FixupData[0] < ImageBase) ||
+      (FixupData[1] > (ImageBase + ImageSize))) {
+    ASSERT (FALSE);
+    *RelocationRangeMin = ImageBase;
+    *RelocationRangeMax = ImageBase + ImageSize;
+    return;
+  }
+
+  *RelocationRangeMin = FixupData[0];
+  *RelocationRangeMax = FixupData[1];
+}
-- 
2.39.1