From 9fe980971fee64b0247d0085d898da504a3c9b0b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marvin=20Ha=CC=88user?=
 <8659494+mhaeuser@users.noreply.github.com>
Date: Tue, 25 Apr 2023 10:08:34 +0200
Subject: [PATCH] ImageTool: Support FixedAddress

---
 BaseTools/ImageTool/ImageTool.c     | 10 ++++++----
 BaseTools/ImageTool/ImageTool.h     |  3 ++-
 BaseTools/ImageTool/ImageToolEmit.c |  7 ++++++-
 BaseTools/ImageTool/ImageToolEmit.h |  3 ++-
 BaseTools/ImageTool/PeEmit.c        | 17 +++++++++++++++++
 BaseTools/ImageTool/UefiImageScan.c | 20 +++++++++++++++++++-
 6 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/BaseTools/ImageTool/ImageTool.c b/BaseTools/ImageTool/ImageTool.c
index a20ed81a94..07f2f3bb51 100644
--- a/BaseTools/ImageTool/ImageTool.c
+++ b/BaseTools/ImageTool/ImageTool.c
@@ -292,7 +292,8 @@ GenExecutable (
   IN const char  *TypeName,
   IN const char  *HiiFileName,
   IN const char  *BaseAddress,
-  IN bool        Strip
+  IN bool        Strip,
+  IN bool        FixedAddress
   )
 {
   UINT32         InputFileSize;
@@ -360,7 +361,8 @@ GenExecutable (
                  BaseAddress != NULL,
                  NewBaseAddress,
                  InputFileName,
-                 Strip
+                 Strip,
+                 FixedAddress
                  );
 
   if (OutputFile == NULL) {
@@ -394,7 +396,7 @@ int main (int argc, const char *argv[])
       return -1;
     }
 
-    Status = GenExecutable (argv[3], argv[2], "PE", argv[4], argc >= 6 ? argv[5] : NULL, NULL, FALSE);
+    Status = GenExecutable (argv[3], argv[2], "PE", argv[4], argc >= 6 ? argv[5] : NULL, NULL, FALSE, FALSE);
     if (RETURN_ERROR (Status)) {
       raise ();
       return -1;
@@ -422,7 +424,7 @@ int main (int argc, const char *argv[])
       return -1;
     }
 
-    Status = GenExecutable (argv[4], argv[3], "PE", NULL, NULL, argv[2], FALSE);
+    Status = GenExecutable (argv[4], argv[3], "PE", NULL, NULL, argv[2], FALSE, FALSE);
     if (RETURN_ERROR (Status)) {
       raise ();
       return -1;
diff --git a/BaseTools/ImageTool/ImageTool.h b/BaseTools/ImageTool/ImageTool.h
index 8426d8f5b1..320edcc086 100644
--- a/BaseTools/ImageTool/ImageTool.h
+++ b/BaseTools/ImageTool/ImageTool.h
@@ -40,7 +40,8 @@ typedef struct {
   uint16_t Machine;
   uint16_t Subsystem;
   uint8_t  IsXip;
-  uint8_t  Reserved[7];
+  uint8_t  FixedAddress;
+  uint8_t  Reserved[6];
 } image_tool_header_info_t;
 
 typedef struct {
diff --git a/BaseTools/ImageTool/ImageToolEmit.c b/BaseTools/ImageTool/ImageToolEmit.c
index 8ddda2354c..f1ae55b507 100644
--- a/BaseTools/ImageTool/ImageToolEmit.c
+++ b/BaseTools/ImageTool/ImageToolEmit.c
@@ -96,7 +96,8 @@ ToolImageEmit (
   IN  bool        Relocate,
   IN  uint64_t    BaseAddress,
   IN  const char  *SymbolsPath OPTIONAL,
-  IN  bool        Strip
+  IN  bool        Strip,
+  IN  bool        FixedAddress
   )
 {
   RETURN_STATUS            Status;
@@ -159,6 +160,10 @@ ToolImageEmit (
     }
   }
 
+  if (FixedAddress) {
+    ImageInfo.HeaderInfo.FixedAddress = true;
+  }
+
   OutputFile = NULL;
   if (Format == UefiImageFormatPe) {
     OutputFile = ToolImageEmitPe (&ImageInfo, OutputFileSize, Strip);
diff --git a/BaseTools/ImageTool/ImageToolEmit.h b/BaseTools/ImageTool/ImageToolEmit.h
index 11b8b3fa97..32828feaa7 100644
--- a/BaseTools/ImageTool/ImageToolEmit.h
+++ b/BaseTools/ImageTool/ImageToolEmit.h
@@ -22,7 +22,8 @@ ToolImageEmit (
   IN  bool        Relocate,
   IN  uint64_t    BaseAddress,
   IN  const char  *SymbolsPath OPTIONAL,
-  IN  bool        Strip
+  IN  bool        Strip,
+  IN  bool        FixedAddress
   );
 
 #endif // IMAGE_TOOL_EMIT_H
diff --git a/BaseTools/ImageTool/PeEmit.c b/BaseTools/ImageTool/PeEmit.c
index 42197aa143..ab2a18c5f1 100644
--- a/BaseTools/ImageTool/PeEmit.c
+++ b/BaseTools/ImageTool/PeEmit.c
@@ -337,6 +337,23 @@ ToolImageEmitPeSectionHeaders (
     SectionOffset += Sections[Index].SizeOfRawData;
   }
 
+  if (Image->HeaderInfo.FixedAddress) {
+    for (Index = 0; Index < Image->SegmentInfo.NumSegments; ++Index) {
+      if ((Sections[Index].Characteristics & EFI_IMAGE_SCN_MEM_EXECUTE) == 0) {
+        WriteUnaligned64 (
+          (VOID *)&Sections[Index].PointerToRelocations,
+          Image->HeaderInfo.BaseAddress
+          );
+        break;
+      }
+    }
+
+    if (Index == Image->SegmentInfo.NumSegments) {
+      raise ();
+      return false;
+    }
+  }
+
   *BufferSize -= SectionHeadersSize;
   *Buffer     += SectionHeadersSize;
 
diff --git a/BaseTools/ImageTool/UefiImageScan.c b/BaseTools/ImageTool/UefiImageScan.c
index c4c97f5405..60eda3f915 100644
--- a/BaseTools/ImageTool/UefiImageScan.c
+++ b/BaseTools/ImageTool/UefiImageScan.c
@@ -18,16 +18,34 @@ ScanUefiImageGetHeaderInfo (
   IN  UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *Context
   )
 {
+  RETURN_STATUS  Status;
+  UINT64         Address;
+
   assert (HeaderInfo != NULL);
   assert (Context    != NULL);
 
-  HeaderInfo->BaseAddress       = (uint64_t)UefiImageGetPreferredAddress (Context);
+  HeaderInfo->BaseAddress       = UefiImageGetPreferredAddress (Context);
   HeaderInfo->EntryPointAddress = UefiImageGetEntryPointAddress (Context);
   HeaderInfo->Machine           = UefiImageGetMachine (Context);
   HeaderInfo->Subsystem         = UefiImageGetSubsystem (Context);
   // FIXME:
   HeaderInfo->IsXip             = true;
 
+  Status = UefiImageGetFixedAddress (Context, &Address);
+  if (!RETURN_ERROR (Status)) {
+    if (Address != HeaderInfo->BaseAddress) {
+      fprintf (
+        stderr,
+        "ImageTool: Images with a fixed address different from their base address are not supported.\n"
+        );
+      return false;
+    }
+
+    HeaderInfo->FixedAddress = true;
+  } else if (Status != RETURN_NOT_FOUND) {
+    return false;
+  }
+
   return true;
 }