From f374c858dfd05c9f3593af15381cc13df9c12d16 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 01:25:08 +0200
Subject: [PATCH] ImageTool: Check HeaderInfo

---
 BaseTools/ImageTool/Image.c | 39 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 38 insertions(+), 1 deletion(-)

diff --git a/BaseTools/ImageTool/Image.c b/BaseTools/ImageTool/Image.c
index 45a1d02a97..5f684849d1 100644
--- a/BaseTools/ImageTool/Image.c
+++ b/BaseTools/ImageTool/Image.c
@@ -92,6 +92,28 @@ CheckToolImageSegmentInfo (
   return true;
 }
 
+static
+bool
+CheckToolImageHeaderInfo (
+  const image_tool_header_info_t   *HeaderInfo,
+  const image_tool_segment_info_t  *SegmentInfo,
+  uint32_t                         ImageSize
+  )
+{
+  if (SegmentInfo->Segments[0].ImageAddress > HeaderInfo->EntryPointAddress ||
+      HeaderInfo->EntryPointAddress > ImageSize) {
+    DEBUG_RAISE ();
+    return false;
+  }
+
+  if (!IS_ALIGNED (HeaderInfo->BaseAddress, SegmentInfo->SegmentAlignment)) {
+    DEBUG_RAISE ();
+    return false;
+  }
+
+  return true;
+}
+
 const image_tool_segment_t *
 ImageGetSegmentByAddress (
   uint32_t                        *Address,
@@ -296,6 +318,16 @@ CheckToolImage (
     return false;
   }
 
+  Result = CheckToolImageHeaderInfo (
+             &Image->HeaderInfo,
+             &Image->SegmentInfo,
+             ImageSize
+             );
+  if (!Result) {
+    DEBUG_RAISE ();
+    return false;
+  }
+
   Result = CheckToolImageRelocInfo (Image, ImageSize);
   if (!Result) {
     DEBUG_RAISE ();
@@ -371,10 +403,15 @@ ToolImageRelocate (
   uint32_t                   RelocTarget32;
   uint64_t                   RelocTarget64;
 
+  if (!IS_ALIGNED (BaseAddress, Image->SegmentInfo.SegmentAlignment)) {
+    DEBUG_RAISE ();
+    return false;
+  }
+
   Adjust = BaseAddress - Image->HeaderInfo.BaseAddress;
 
   if (Adjust == 0) {
-    return TRUE;
+    return true;
   }
 
   for (Index = 0; Index < Image->RelocInfo.NumRelocs; ++Index) {