[PATCH 5/5] StandaloneMmPkg/StandaloneMmCoreEntryPoint: relocate StMM core on the fly

["Ard Biesheuvel" <ard.biesheuvel@arm.com>]

Apply PE/COFF fixups when starting up the standalone MM core, so that
it can execute at any address regardless of the link time address.

Note that this requires the PE/COFF image to be emitted with its
relocation section preserved. Special care is taken to ensure that
TE images are dealt with correctly as well.

Signed-off-by: Ard Biesheuvel <ard.biesheuvel@arm.com>
---
 StandaloneMmPkg/Include/Library/AArch64/StandaloneMmCoreEntryPoint.h                    |  2 ++
 StandaloneMmPkg/Library/StandaloneMmCoreEntryPoint/AArch64/SetPermissions.c             | 11 +++++++---
 StandaloneMmPkg/Library/StandaloneMmCoreEntryPoint/AArch64/StandaloneMmCoreEntryPoint.c | 22 ++++++++++++++++++++
 3 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/StandaloneMmPkg/Include/Library/AArch64/StandaloneMmCoreEntryPoint.h b/StandaloneMmPkg/Include/Library/AArch64/StandaloneMmCoreEntryPoint.h
index 494bcf3dc28f..a3420699e6f1 100644
--- a/StandaloneMmPkg/Include/Library/AArch64/StandaloneMmCoreEntryPoint.h
+++ b/StandaloneMmPkg/Include/Library/AArch64/StandaloneMmCoreEntryPoint.h
@@ -82,6 +82,7 @@ EFI_STATUS
 EFIAPI
 UpdateMmFoundationPeCoffPermissions (
   IN  CONST PE_COFF_LOADER_IMAGE_CONTEXT      *ImageContext,
+  IN  EFI_PHYSICAL_ADDRESS                    ImageBase,
   IN  UINT32                                  SectionHeaderOffset,
   IN  CONST  UINT16                           NumberOfSections,
   IN  REGION_PERMISSION_UPDATE_FUNC           TextUpdater,
@@ -107,6 +108,7 @@ EFIAPI
 GetStandaloneMmCorePeCoffSections (
   IN        VOID                            *TeData,
   IN  OUT   PE_COFF_LOADER_IMAGE_CONTEXT    *ImageContext,
+      OUT   EFI_PHYSICAL_ADDRESS            *ImageBase,
   IN  OUT   UINT32                          *SectionHeaderOffset,
   IN  OUT   UINT16                          *NumberOfSections
   );
diff --git a/StandaloneMmPkg/Library/StandaloneMmCoreEntryPoint/AArch64/SetPermissions.c b/StandaloneMmPkg/Library/StandaloneMmCoreEntryPoint/AArch64/SetPermissions.c
index 00f49c9d0558..bf9650d54629 100644
--- a/StandaloneMmPkg/Library/StandaloneMmCoreEntryPoint/AArch64/SetPermissions.c
+++ b/StandaloneMmPkg/Library/StandaloneMmCoreEntryPoint/AArch64/SetPermissions.c
@@ -29,6 +29,7 @@ EFI_STATUS
 EFIAPI
 UpdateMmFoundationPeCoffPermissions (
   IN  CONST PE_COFF_LOADER_IMAGE_CONTEXT      *ImageContext,
+  IN  EFI_PHYSICAL_ADDRESS                    ImageBase,
   IN  UINT32                                  SectionHeaderOffset,
   IN  CONST  UINT16                           NumberOfSections,
   IN  REGION_PERMISSION_UPDATE_FUNC           TextUpdater,
@@ -87,7 +88,7 @@ UpdateMmFoundationPeCoffPermissions (
     // if it is a writeable section then mark it appropriately as well.
     //
     if ((SectionHeader.Characteristics & EFI_IMAGE_SCN_MEM_EXECUTE) == 0) {
-      Base = ImageContext->ImageAddress + SectionHeader.VirtualAddress;
+      Base = ImageBase + SectionHeader.VirtualAddress;
 
       TextUpdater (Base, SectionHeader.Misc.VirtualSize);
 
@@ -153,6 +154,7 @@ STATIC
 EFI_STATUS
 GetPeCoffSectionInformation (
   IN  OUT   PE_COFF_LOADER_IMAGE_CONTEXT      *ImageContext,
+      OUT   EFI_PHYSICAL_ADDRESS              *ImageBase,
       OUT   UINT32                            *SectionHeaderOffset,
       OUT   UINT16                            *NumberOfSections
   )
@@ -212,6 +214,7 @@ GetPeCoffSectionInformation (
     return Status;
   }
 
+  *ImageBase = ImageContext->ImageAddress;
   if (!ImageContext->IsTeImage) {
     ASSERT (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE);
 
@@ -232,7 +235,7 @@ GetPeCoffSectionInformation (
   } else {
     *SectionHeaderOffset = (UINTN)(sizeof (EFI_TE_IMAGE_HEADER));
     *NumberOfSections = Hdr.Te->NumberOfSections;
-    ImageContext->ImageAddress -= (UINT32)Hdr.Te->StrippedSize - sizeof (EFI_TE_IMAGE_HEADER);
+    *ImageBase -= (UINT32)Hdr.Te->StrippedSize - sizeof (EFI_TE_IMAGE_HEADER);
   }
   return RETURN_SUCCESS;
 }
@@ -242,6 +245,7 @@ EFIAPI
 GetStandaloneMmCorePeCoffSections (
   IN        VOID                            *TeData,
   IN  OUT   PE_COFF_LOADER_IMAGE_CONTEXT    *ImageContext,
+      OUT   EFI_PHYSICAL_ADDRESS            *ImageBase,
   IN  OUT   UINT32                          *SectionHeaderOffset,
   IN  OUT   UINT16                          *NumberOfSections
   )
@@ -255,7 +259,8 @@ GetStandaloneMmCorePeCoffSections (
 
   DEBUG ((DEBUG_INFO, "Found Standalone MM PE data - 0x%x\n", TeData));
 
-  Status = GetPeCoffSectionInformation (ImageContext, SectionHeaderOffset, NumberOfSections);
+  Status = GetPeCoffSectionInformation (ImageContext, ImageBase,
+             SectionHeaderOffset, NumberOfSections);
   if (EFI_ERROR (Status)) {
     DEBUG ((DEBUG_ERROR, "Unable to locate Standalone MM Core PE-COFF Section information - %r\n", Status));
     return Status;
diff --git a/StandaloneMmPkg/Library/StandaloneMmCoreEntryPoint/AArch64/StandaloneMmCoreEntryPoint.c b/StandaloneMmPkg/Library/StandaloneMmCoreEntryPoint/AArch64/StandaloneMmCoreEntryPoint.c
index 20723385113f..9cecfa667b90 100644
--- a/StandaloneMmPkg/Library/StandaloneMmCoreEntryPoint/AArch64/StandaloneMmCoreEntryPoint.c
+++ b/StandaloneMmPkg/Library/StandaloneMmCoreEntryPoint/AArch64/StandaloneMmCoreEntryPoint.c
@@ -225,6 +225,7 @@ _ModuleEntryPoint (
   VOID                                    *HobStart;
   VOID                                    *TeData;
   UINTN                                   TeDataSize;
+  EFI_PHYSICAL_ADDRESS                    ImageBase;
 
   // Get Secure Partition Manager Version Information
   Status = GetSpmVersion ();
@@ -253,6 +254,7 @@ _ModuleEntryPoint (
   Status = GetStandaloneMmCorePeCoffSections (
              TeData,
              &ImageContext,
+             &ImageBase,
              &SectionHeaderOffset,
              &NumberOfSections
              );
@@ -261,10 +263,21 @@ _ModuleEntryPoint (
     goto finish;
   }
 
+  //
+  // ImageBase may deviate from ImageContext.ImageAddress if we are dealing
+  // with a TE image, in which case the latter points to the actual offset
+  // of the image, whereas ImageBase refers to the address where the image
+  // would start if the stripped PE headers were still in place. In either
+  // case, we need to fix up ImageBase so it refers to the actual current
+  // load address.
+  //
+  ImageBase += (UINTN)TeData - ImageContext.ImageAddress;
+
   // Update the memory access permissions of individual sections in the
   // Standalone MM core module
   Status = UpdateMmFoundationPeCoffPermissions (
              &ImageContext,
+             ImageBase,
              SectionHeaderOffset,
              NumberOfSections,
              ArmSetMemoryRegionNoExec,
@@ -276,6 +289,15 @@ _ModuleEntryPoint (
     goto finish;
   }
 
+  if (ImageContext.ImageAddress != (UINTN)TeData) {
+    ImageContext.ImageAddress = (UINTN)TeData;
+    ArmSetMemoryRegionNoExec (ImageBase, SIZE_4KB);
+    ArmClearMemoryRegionReadOnly (ImageBase, SIZE_4KB);
+
+    Status = PeCoffLoaderRelocateImage (&ImageContext);
+    ASSERT_EFI_ERROR (Status);
+  }
+
   //
   // Create Hoblist based upon boot information passed by privileged software
   //
-- 
2.26.2