From b74a1336a558928bb55c8e66d19899c12c1c091e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marvin=20Ha=CC=88user?=
 <8659494+mhaeuser@users.noreply.github.com>
Date: Tue, 21 Mar 2023 16:19:54 +0100
Subject: [PATCH] MdePkg/BasePeCoffLib2: Move ExtraAction to UefiImageLib

---
 .../Library/BasePeCoffLib2/BasePeCoffLib2.inf |  1 -
 .../Library/BasePeCoffLib2/PeCoffRelocate.c   |  4 ----
 .../BaseUefiImageLibPeCoff.inf                |  1 +
 .../BaseUefiImageLib/BaseUefiImageLibUe.inf   |  1 +
 .../BaseUefiImageLib/UefiImageLibPeCoff.c     | 20 +++++++++++++------
 .../Library/BaseUefiImageLib/UefiImageLibUe.c | 13 ++++++++----
 6 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/MdePkg/Library/BasePeCoffLib2/BasePeCoffLib2.inf b/MdePkg/Library/BasePeCoffLib2/BasePeCoffLib2.inf
index 8428ccf694..1f1dcf29ca 100644
--- a/MdePkg/Library/BasePeCoffLib2/BasePeCoffLib2.inf
+++ b/MdePkg/Library/BasePeCoffLib2/BasePeCoffLib2.inf
@@ -34,7 +34,6 @@
   BaseOverflowLib
   DebugLib
   MemoryAllocationLib
-  UefiImageExtraActionLib
 
 [FixedPcd]
   gEfiMdePkgTokenSpaceGuid.PcdImageLoaderRtRelocAllowTargetMismatch
diff --git a/MdePkg/Library/BasePeCoffLib2/PeCoffRelocate.c b/MdePkg/Library/BasePeCoffLib2/PeCoffRelocate.c
index 40a79ff993..34a598cb73 100644
--- a/MdePkg/Library/BasePeCoffLib2/PeCoffRelocate.c
+++ b/MdePkg/Library/BasePeCoffLib2/PeCoffRelocate.c
@@ -19,7 +19,6 @@
 #include <Library/BaseMemoryLib.h>
 #include <Library/BaseOverflowLib.h>
 #include <Library/DebugLib.h>
-#include <Library/UefiImageExtraActionLib.h>
 #include <Library/PcdLib.h>
 #include <Library/PeCoffLib2.h>
 
@@ -451,7 +450,6 @@ PeCoffRelocateImage (
   // Verify the Relocation Directory is not empty.
   //
   if (Context->RelocDirSize == 0) {
-    UefiImageLoaderRelocateImageExtraAction (Context);
     return RETURN_SUCCESS;
   }
   //
@@ -467,7 +465,6 @@ PeCoffRelocateImage (
   // prefered location.
   //
   if (RuntimeContext == NULL && Adjust == 0) {
-    UefiImageLoaderRelocateImageExtraAction (Context);
     return RETURN_SUCCESS;
   }
   //
@@ -623,7 +620,6 @@ PeCoffRelocateImage (
       );
   }
 
-  UefiImageLoaderRelocateImageExtraAction (Context);
   return RETURN_SUCCESS;
 }
 
diff --git a/MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibPeCoff.inf b/MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibPeCoff.inf
index cf9bb407dd..85cd2831e8 100644
--- a/MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibPeCoff.inf
+++ b/MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibPeCoff.inf
@@ -31,6 +31,7 @@
   DebugLib
   MemoryAllocationLib
   PeCoffLib2
+  UefiImageExtraActionLib
 
 [FixedPcd]
   gEfiMdePkgTokenSpaceGuid.PcdImageLoaderAlignmentPolicy
diff --git a/MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibUe.inf b/MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibUe.inf
index c10cca03d5..abb5354557 100644
--- a/MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibUe.inf
+++ b/MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibUe.inf
@@ -31,6 +31,7 @@
   DebugLib
   MemoryAllocationLib
   UeImageLib
+  UefiImageExtraActionLib
 
 [FixedPcd]
   gEfiMdePkgTokenSpaceGuid.PcdDebugRaisePropertyMask
diff --git a/MdePkg/Library/BaseUefiImageLib/UefiImageLibPeCoff.c b/MdePkg/Library/BaseUefiImageLib/UefiImageLibPeCoff.c
index 1dc3e9ada4..6dde99baa1 100644
--- a/MdePkg/Library/BaseUefiImageLib/UefiImageLibPeCoff.c
+++ b/MdePkg/Library/BaseUefiImageLib/UefiImageLibPeCoff.c
@@ -22,6 +22,7 @@
 #include <Library/PcdLib.h>
 #include <Library/PeCoffLib2.h>
 #include <Library/UefiImageLib.h>
+#include <Library/UefiImageExtraActionLib.h>
 
 #include "PeCoffSupport.h"
 
@@ -106,12 +107,19 @@ UefiImageRelocateImage (
   IN     UINT32                             RuntimeContextSize
   )
 {
-  return PeCoffRelocateImage (
-           Context,
-           BaseAddress,
-           RuntimeContext,
-           RuntimeContextSize
-           );
+  RETURN_STATUS Status;
+
+  Status = PeCoffRelocateImage (
+             Context,
+             BaseAddress,
+             RuntimeContext,
+             RuntimeContextSize
+             );
+  if (!RETURN_ERROR (Status)) {
+    UefiImageLoaderRelocateImageExtraAction (Context);
+  }
+
+  return Status;
 }
 
 RETURN_STATUS
diff --git a/MdePkg/Library/BaseUefiImageLib/UefiImageLibUe.c b/MdePkg/Library/BaseUefiImageLib/UefiImageLibUe.c
index 984e8c9309..c63dbea75c 100644
--- a/MdePkg/Library/BaseUefiImageLib/UefiImageLibUe.c
+++ b/MdePkg/Library/BaseUefiImageLib/UefiImageLibUe.c
@@ -22,6 +22,7 @@
 #include <Library/PcdLib.h>
 #include <Library/UeImageLib.h>
 #include <Library/UefiImageLib.h>
+#include <Library/UefiImageExtraActionLib.h>
 
 #include "UeSupport.h"
 
@@ -100,10 +101,14 @@ UefiImageRelocateImage (
   IN     UINT32                             RuntimeContextSize
   )
 {
-  return UeRelocateImage (
-           Context,
-           BaseAddress
-           );
+  RETURN_STATUS Status;
+
+  Status = UeRelocateImage (Context, BaseAddress);
+  if (!RETURN_ERROR (Status)) {
+    UefiImageLoaderRelocateImageExtraAction (Context);
+  }
+
+  return Status;
 }
 
 RETURN_STATUS