From 8957da50bb2ed1ce4c7d390ec8e435cfbbe2280d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marvin=20Ha=CC=88user?=
 <8659494+mhaeuser@users.noreply.github.com>
Date: Sat, 3 Jun 2023 15:54:08 +0200
Subject: [PATCH] ImageTool: Update relocation logic

---
 BaseTools/ImageTool/Image.c         | 115 +++++++++++++++-------------
 BaseTools/ImageTool/ImageTool.h     |  14 +++-
 BaseTools/ImageTool/ImageToolEmit.c |   1 +
 3 files changed, 75 insertions(+), 55 deletions(-)

diff --git a/BaseTools/ImageTool/Image.c b/BaseTools/ImageTool/Image.c
index 0c32659c17..22e9488e04 100644
--- a/BaseTools/ImageTool/Image.c
+++ b/BaseTools/ImageTool/Image.c
@@ -92,7 +92,6 @@ CheckToolImageSegmentInfo (
   return true;
 }
 
-static
 const image_tool_segment_t *
 ImageGetSegmentByAddress (
   uint32_t                        *Address,
@@ -117,6 +116,34 @@ ImageGetSegmentByAddress (
   return NULL;
 }
 
+uint8_t
+ToolImageGetRelocSize (
+  uint8_t  Type
+  )
+{
+  switch (Type) {
+    case EFI_IMAGE_REL_BASED_HIGHLOW:
+    {
+      return sizeof (UINT32);
+    }
+
+    case EFI_IMAGE_REL_BASED_DIR64:
+    {
+      return sizeof (UINT64);
+    }
+
+    case EFI_IMAGE_REL_BASED_ARM_MOV32T:
+    {
+      return sizeof (UINT32);
+    }
+
+    default:
+    {
+      return 0;
+    }
+  }
+}
+
 static
 bool
 CheckToolImageReloc (
@@ -126,6 +153,7 @@ CheckToolImageReloc (
   )
 {
   uint32_t                   RelocOffset;
+  uint8_t                    RelocSize;
   uint32_t                   RemainingSize;
   const image_tool_segment_t *Segment;
   uint16_t                   MovHigh;
@@ -145,60 +173,40 @@ CheckToolImageReloc (
     return false;
   }
 
-  switch (Reloc->Type) {
-    case EFI_IMAGE_REL_BASED_HIGHLOW:
-    {
-      if (RemainingSize < sizeof (UINT32)) {
-        raise ();
-        return false;
-      }
+  RelocSize = ToolImageGetRelocSize (Reloc->Type);
+  if (RelocSize == 0) {
+    raise ();
+    return false;
+  }
 
-      break;
+  if (RelocSize > RemainingSize) {
+    raise ();
+    return false;
+  }
+
+  if (Reloc->Type == EFI_IMAGE_REL_BASED_ARM_MOV32T) {
+    if (!IS_ALIGNED (Reloc->Target, ALIGNOF (UINT16))) {
+      raise ();
+      return false;
     }
 
-    case EFI_IMAGE_REL_BASED_DIR64:
-    {
-      if (RemainingSize < sizeof (UINT64)) {
-        raise ();
-        return false;
-      }
-
-      break;
-    }
-
-    case EFI_IMAGE_REL_BASED_ARM_MOV32T:
-    {
-      if (RemainingSize < sizeof (UINT32)) {
-        raise ();
-        return false;
-      }
-
-      if (!IS_ALIGNED (Reloc->Target, ALIGNOF (UINT16))) {
-        raise ();
-        return false;
-      }
-
-      MovHigh = *(const uint16_t *)&Segment->Data[RelocOffset];
-      MovLow  = *(const uint16_t *)&Segment->Data[RelocOffset + 2];
-      if (((MovHigh & 0xFBF0U) != 0xF200U && (MovHigh & 0xFBF0U) != 0xF2C0U) ||
-        (MovLow & 0x8000U) != 0) {
-        raise ();
-        return false;
-      }
-
-      break;
-    }
-
-    default:
-    {
+    MovHigh = *(const uint16_t *)&Segment->Data[RelocOffset];
+    MovLow  = *(const uint16_t *)&Segment->Data[RelocOffset + 2];
+    if (((MovHigh & 0xFBF0U) != 0xF200U && (MovHigh & 0xFBF0U) != 0xF2C0U) ||
+      (MovLow & 0x8000U) != 0) {
       raise ();
       return false;
     }
   }
 
-  /*if (Segment->Write) {
+  // FIXME: Update drivers?
+  if ((Image->HeaderInfo.Subsystem == EFI_IMAGE_SUBSYSTEM_EFI_RUNTIME_DRIVER ||
+       Image->HeaderInfo.Subsystem == EFI_IMAGE_SUBSYSTEM_SAL_RUNTIME_DRIVER) &&
+      Segment->Write) {
     printf("!!! writable reloc at %x !!!\n", Reloc->Target);
-  }*/
+    //raise ();
+    //return false;
+  }
 
   return true;
 }
@@ -211,6 +219,7 @@ CheckToolImageRelocInfo (
   )
 {
   const image_tool_reloc_info_t *RelocInfo;
+  uint32_t                      PrevTarget;
   uint32_t                      Index;
   bool                          Result;
 
@@ -222,7 +231,7 @@ CheckToolImageRelocInfo (
     return true;
   }
 
-  if (RelocInfo->RelocsStripped && (RelocInfo->NumRelocs > 0)) {
+  if (RelocInfo->RelocsStripped) {
     raise ();
     return false;
   }
@@ -232,14 +241,10 @@ CheckToolImageRelocInfo (
     return false;
   }
 
-  Result = CheckToolImageReloc (Image, ImageSize, &RelocInfo->Relocs[0]);
-  if (!Result) {
-    raise ();
-    return false;
-  }
+  PrevTarget = 0;
 
-  for (Index = 1; Index < RelocInfo->NumRelocs; ++Index) {
-    if (RelocInfo->Relocs[Index].Target < RelocInfo->Relocs[Index - 1].Target) {
+  for (Index = 0; Index < RelocInfo->NumRelocs; ++Index) {
+    if (RelocInfo->Relocs[Index].Target < PrevTarget) {
       assert (false);
       return false;
     }
@@ -249,6 +254,8 @@ CheckToolImageRelocInfo (
       raise ();
       return false;
     }
+
+    PrevTarget = RelocInfo->Relocs[Index].Target;
   }
 
   return true;
diff --git a/BaseTools/ImageTool/ImageTool.h b/BaseTools/ImageTool/ImageTool.h
index cae2838ced..99da2b29e8 100644
--- a/BaseTools/ImageTool/ImageTool.h
+++ b/BaseTools/ImageTool/ImageTool.h
@@ -34,7 +34,7 @@ typedef struct {
   uint16_t Machine;
   uint16_t Subsystem;
   uint8_t  FixedAddress;
-  uint8_t  Reserved[6];
+  uint8_t  Reserved[7];
 } image_tool_header_info_t;
 
 typedef struct {
@@ -130,6 +130,11 @@ ToolImageStripRelocs (
   image_tool_image_info_t  *Image
   );
 
+uint8_t
+ToolImageGetRelocSize (
+  uint8_t  Type
+  );
+
 RETURN_STATUS
 ToolContextConstructUefiImage (
   OUT image_tool_image_info_t *Image,
@@ -176,4 +181,11 @@ ConstructHii (
   OUT UINT32     *HiiSize
   );
 
+const image_tool_segment_t *
+ImageGetSegmentByAddress (
+  uint32_t                        *Address,
+  uint32_t                        *RemainingSize,
+  const image_tool_segment_info_t *SegmentInfo
+  );
+
 #endif // IMAGE_TOOL_H
diff --git a/BaseTools/ImageTool/ImageToolEmit.c b/BaseTools/ImageTool/ImageToolEmit.c
index 657ad8ce96..459b217ac2 100644
--- a/BaseTools/ImageTool/ImageToolEmit.c
+++ b/BaseTools/ImageTool/ImageToolEmit.c
@@ -138,6 +138,7 @@ ToolImageEmit (
 
   Success = CheckToolImage (&ImageInfo);
   if (!Success) {
+    raise ();
     ToolImageDestruct (&ImageInfo);
     return NULL;
   }