From 9f43ad06afb7bbf027ee315ca1c624f399729a11 Mon Sep 17 00:00:00 2001
From: belliash <belliash@codingworkshop.eu.org>
Date: Tue, 15 Nov 2022 23:31:35 +0100
Subject: [PATCH] Implement PeRelocateImage() routine

---
 xtldr/includes/blmod.h                 |  9 ++--
 xtldr/modules/pecoff/includes/pecoff.h |  6 ++-
 xtldr/modules/pecoff/pecoff.c          | 58 ++++++++++++++++++++++++--
 3 files changed, 66 insertions(+), 7 deletions(-)

diff --git a/xtldr/includes/blmod.h b/xtldr/includes/blmod.h
index 3d36606..66f415a 100644
--- a/xtldr/includes/blmod.h
+++ b/xtldr/includes/blmod.h
@@ -22,8 +22,10 @@ typedef struct _XT_PECOFFF_IMAGE_PROTOCOL XT_PECOFF_IMAGE_PROTOCOL, *PXT_PECOFF_
 /* Pointers to the routines provided by the modules */
 typedef EFI_STATUS (*PXT_BOOTPROTO_BOOT_SYSTEM)(IN PXT_BOOT_PROTOCOL_PARAMETERS Parameters);
 typedef EFI_STATUS (*PXT_PECOFF_GET_ENTRY_POINT)(IN PPECOFF_IMAGE_CONTEXT Image, OUT PVOID *EntryPoint);
-typedef EFI_STATUS (*PXT_PECOFF_PROTOCOL_LOAD)(IN PEFI_FILE_HANDLE FileHandle, IN LOADER_MEMORY_TYPE MemoryType,
-                                               IN PVOID VirtualAddress, OUT PPECOFF_IMAGE_CONTEXT *Image);
+typedef EFI_STATUS (*PXT_PECOFF_LOAD_IMAGE)(IN PEFI_FILE_HANDLE FileHandle, IN LOADER_MEMORY_TYPE MemoryType,
+                                            IN PVOID VirtualAddress, OUT PPECOFF_IMAGE_CONTEXT *Image);
+typedef EFI_STATUS (*PXT_PECOFF_RELOCATE_IMAGE)(IN PPECOFF_IMAGE_CONTEXT Image, IN EFI_VIRTUAL_ADDRESS Address);
+
 
 /* XT common boot protocols */
 typedef struct _XT_BOOT_PROTOCOL
@@ -47,7 +49,8 @@ typedef struct _XT_BOOT_PROTOCOL_PARAMETERS
 typedef struct _XT_PECOFFF_IMAGE_PROTOCOL
 {
     PXT_PECOFF_GET_ENTRY_POINT GetEntryPoint;
-    PXT_PECOFF_PROTOCOL_LOAD Load;
+    PXT_PECOFF_LOAD_IMAGE Load;
+    PXT_PECOFF_RELOCATE_IMAGE Relocate;
 } XT_PECOFF_IMAGE_PROTOCOL, *PXT_PECOFF_IMAGE_PROTOCOL;
 
 #endif /* __XTLDR_BLMOD_H */
diff --git a/xtldr/modules/pecoff/includes/pecoff.h b/xtldr/modules/pecoff/includes/pecoff.h
index 64c5123..885f885 100644
--- a/xtldr/modules/pecoff/includes/pecoff.h
+++ b/xtldr/modules/pecoff/includes/pecoff.h
@@ -24,7 +24,11 @@ PeLoadImage(IN PEFI_FILE_HANDLE FileHandle,
             OUT PPECOFF_IMAGE_CONTEXT *Image);
 
 EFI_STATUS
-PepRelocateImage(IN PPECOFF_IMAGE_CONTEXT Image);
+PeRelocateImage(IN PPECOFF_IMAGE_CONTEXT Image,
+                IN EFI_VIRTUAL_ADDRESS Address);
+
+EFI_STATUS
+PepRelocateLoadedImage(IN PPECOFF_IMAGE_CONTEXT Image);
 
 EFI_STATUS
 PepValidateImageHeaders(IN PPECOFF_IMAGE_DOS_HEADER DosHeader,
diff --git a/xtldr/modules/pecoff/pecoff.c b/xtldr/modules/pecoff/pecoff.c
index 862e16a..a815232 100644
--- a/xtldr/modules/pecoff/pecoff.c
+++ b/xtldr/modules/pecoff/pecoff.c
@@ -46,7 +46,7 @@ PeGetEntryPoint(IN PPECOFF_IMAGE_CONTEXT Image,
     }
 
     /* Set entry point and return success */
-    *EntryPoint = (UINT8*)Image->VirtualAddress + Image->PeHeader->OptionalHeader.AddressOfEntryPoint;
+    *EntryPoint = (PUINT8)Image->VirtualAddress + Image->PeHeader->OptionalHeader.AddressOfEntryPoint;
     return STATUS_EFI_SUCCESS;
 }
 
@@ -259,7 +259,7 @@ PeLoadImage(IN PEFI_FILE_HANDLE FileHandle,
     XtLdrProtocol->FreePages((EFI_PHYSICAL_ADDRESS)(UINT_PTR)Data, Pages);
 
     /* Perform relocation fixups */
-    Status = PepRelocateImage(ImageData);
+    Status = PepRelocateLoadedImage(ImageData);
     if(Status != STATUS_EFI_SUCCESS)
     {
         XtLdrProtocol->DbgPrint(L"ERROR: PE/COFF image relocation failed\n");
@@ -273,6 +273,57 @@ PeLoadImage(IN PEFI_FILE_HANDLE FileHandle,
     return STATUS_EFI_SUCCESS;
 }
 
+/**
+ * Relocates PE/COFF image to the specified address.
+ *
+ * @param Image
+ *        A pointer to the PE/COFF context structure representing the loaded image.
+ *
+ * @param Address
+ *        Destination address of memory region, where image should be relocated.
+ *
+ * @return This routine returns status code.
+ *
+ * @since XT 1.0
+ */
+EFI_STATUS
+PeRelocateImage(IN PPECOFF_IMAGE_CONTEXT Image,
+                IN EFI_VIRTUAL_ADDRESS Address)
+{
+    UINT64 ImageBase, OldVirtualAddress;
+    EFI_STATUS Status;
+
+    /* Store original virtual address */
+    OldVirtualAddress = (UINT_PTR)Image->VirtualAddress;
+
+    /* Check PE/COFF image type */
+    if(Image->PeHeader->OptionalHeader.Magic == PECOFF_IMAGE_PE_OPTIONAL_HDR64_MAGIC)
+    {
+        /* This is 64-bit PE32+, store its image base address */
+        ImageBase = Image->PeHeader->OptionalHeader.ImageBase64;
+    }
+    else
+    {
+        /* This is 32-bit PE32, store its image base address */
+        ImageBase = Image->PeHeader->OptionalHeader.ImageBase32;
+    }
+
+    /* Overwrite virtual address and relocate image once again */
+    Image->VirtualAddress = (PVOID)(Address - OldVirtualAddress + ImageBase);
+    Status = PepRelocateLoadedImage(Image);
+    if(Status != STATUS_EFI_SUCCESS)
+    {
+        /* Relocation failed */
+        return Status;
+    }
+
+    /* Store new image virtual address */
+    Image->VirtualAddress = (PVOID)Address;
+
+    /* Return success */
+    return STATUS_EFI_SUCCESS;
+}
+
 /**
  * Relocates a loaded PE/COFF image.
  *
@@ -284,7 +335,7 @@ PeLoadImage(IN PEFI_FILE_HANDLE FileHandle,
  * @since XT 1.0
  */
 EFI_STATUS
-PepRelocateImage(IN PPECOFF_IMAGE_CONTEXT Image)
+PepRelocateLoadedImage(IN PPECOFF_IMAGE_CONTEXT Image)
 {
     PPECOFF_IMAGE_BASE_RELOCATION RelocationDir, RelocationEnd;
     PPECOFF_IMAGE_DATA_DIRECTORY DataDirectory;
@@ -476,6 +527,7 @@ BlXtLdrModuleMain(IN EFI_HANDLE ImageHandle,
     /* Set routines available via PE/COFF image protocol */
     XtPeCoffProtocol.GetEntryPoint = PeGetEntryPoint;
     XtPeCoffProtocol.Load = PeLoadImage;
+    XtPeCoffProtocol.Relocate = PeRelocateImage;
 
     /* Register PE/COFF protocol */
     return EfiSystemTable->BootServices->InstallProtocolInterface(&Handle, &Guid, EFI_NATIVE_INTERFACE, &XtPeCoffProtocol);