From b25125e0e1973fae2be9c2348dce03e0af42b95d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marvin=20Ha=CC=88user?=
 <8659494+mhaeuser@users.noreply.github.com>
Date: Tue, 13 Jun 2023 02:17:36 +0200
Subject: [PATCH] ImageTool: Handle image address overflows

---
 BaseTools/ImageTool/Image.c         | 29 ++++++++++++++++++++++++-----
 BaseTools/ImageTool/ImageTool.h     |  3 ++-
 BaseTools/ImageTool/ImageToolEmit.c |  2 +-
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/BaseTools/ImageTool/Image.c b/BaseTools/ImageTool/Image.c
index b93438d224..e803d0f87a 100644
--- a/BaseTools/ImageTool/Image.c
+++ b/BaseTools/ImageTool/Image.c
@@ -105,6 +105,11 @@ CheckToolImageHeaderInfo (
     return false;
   }
 
+  if (HeaderInfo->BaseAddress + ImageSize < HeaderInfo->BaseAddress) {
+    DEBUG_RAISE ();
+    return false;
+  }
+
   return true;
 }
 
@@ -161,7 +166,6 @@ static
 bool
 CheckToolImageReloc (
   const image_tool_image_info_t *Image,
-  uint32_t                      ImageSize,
   const image_tool_reloc_t      *Reloc
   )
 {
@@ -224,8 +228,7 @@ CheckToolImageReloc (
 static
 bool
 CheckToolImageRelocInfo (
-  const image_tool_image_info_t *Image,
-  uint32_t                      ImageSize
+  const image_tool_image_info_t *Image
   )
 {
   const image_tool_reloc_info_t *RelocInfo;
@@ -310,7 +313,7 @@ CheckToolImage (
     return false;
   }
 
-  Result = CheckToolImageRelocInfo (Image, ImageSize);
+  Result = CheckToolImageRelocInfo (Image);
   if (!Result) {
     DEBUG_RAISE ();
     return false;
@@ -371,9 +374,12 @@ ToolImageDestruct (
 bool
 ToolImageRelocate (
   image_tool_image_info_t *Image,
-  uint64_t                BaseAddress
+  uint64_t                BaseAddress,
+  uint32_t                IgnorePrefix
   )
 {
+  const image_tool_segment_t *LastSegment;
+  uint32_t                   ImageSize;
   uint64_t                   Adjust;
   const image_tool_reloc_t   *Reloc;
   uint32_t                   RelocOffset;
@@ -394,6 +400,19 @@ ToolImageRelocate (
     return true;
   }
 
+  LastSegment = &Image->SegmentInfo.Segments[Image->SegmentInfo.NumSegments - 1];
+  ImageSize   = LastSegment->ImageAddress + LastSegment->ImageSize;
+
+  //
+  // When removing the image header prefix, BaseAddress + ImageSize may indeed
+  // overflow. The important part is that the address starting from the first
+  // image segment does not.
+  //
+  if (BaseAddress + ImageSize < BaseAddress + IgnorePrefix) {
+    DEBUG_RAISE ();
+    return false;
+  }
+
   for (Index = 0; Index < Image->RelocInfo.NumRelocs; ++Index) {
     Reloc = &Image->RelocInfo.Relocs[Index];
 
diff --git a/BaseTools/ImageTool/ImageTool.h b/BaseTools/ImageTool/ImageTool.h
index e49eb5e30b..a0382d79d8 100644
--- a/BaseTools/ImageTool/ImageTool.h
+++ b/BaseTools/ImageTool/ImageTool.h
@@ -110,7 +110,8 @@ ImageInitUnpaddedSize (
 bool
 ToolImageRelocate (
   image_tool_image_info_t *Image,
-  uint64_t                BaseAddress
+  uint64_t                BaseAddress,
+  uint32_t                IgnorePrefix
   );
 
 void
diff --git a/BaseTools/ImageTool/ImageToolEmit.c b/BaseTools/ImageTool/ImageToolEmit.c
index ba602f2fc8..138cd4715f 100644
--- a/BaseTools/ImageTool/ImageToolEmit.c
+++ b/BaseTools/ImageTool/ImageToolEmit.c
@@ -145,7 +145,7 @@ ToolImageEmit (
   }
 
   if (Relocate) {
-    Success = ToolImageRelocate (&ImageInfo, BaseAddress);
+    Success = ToolImageRelocate (&ImageInfo, BaseAddress, 0);
     if (!Success) {
       fprintf (stderr, "ImageTool: Failed to relocate input file %s\n", SymbolsPath);
       ToolImageDestruct (&ImageInfo);