From ed6dde5a8e0483cadfb82f1ab094038629d3fe16 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marvin=20Ha=CC=88user?=
 <8659494+mhaeuser@users.noreply.github.com>
Date: Tue, 11 Apr 2023 09:47:52 +0200
Subject: [PATCH] ArmPkg/ArmMmuLib: ASSERT ArmReplaceLiveTranslationEntry
 alignment

---
 ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c   | 10 ++++++++++
 .../ArmMmuLib/AArch64/ArmMmuPeiLibConstructor.c    | 14 ++++++++++++--
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c
index 45b7c81341..cb45102bcc 100644
--- a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c
+++ b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c
@@ -755,12 +755,22 @@ ArmMmuBaseLibConstructor (
   )
 {
   extern UINT32  ArmReplaceLiveTranslationEntrySize;
+  UINTN          ArmReplaceLiveTranslationEntryEnd;
   VOID           *Hob;
 
   Hob = GetFirstGuidHob (&gArmMmuReplaceLiveTranslationEntryFuncGuid);
   if (Hob != NULL) {
     mReplaceLiveEntryFunc = *(ARM_REPLACE_LIVE_TRANSLATION_ENTRY *)GET_GUID_HOB_DATA (Hob);
   } else {
+    ArmReplaceLiveTranslationEntryEnd = (UINTN)ArmReplaceLiveTranslationEntry + ArmReplaceLiveTranslationEntrySize;
+    //
+    // Align this routine to a log2 upper bound of its size, so that it is
+    // guaranteed not to cross a page or block boundary
+    // (see ArmMmuLibReplaceEntry.S).
+    //
+    ASSERT (IS_ALIGNED ((UINTN)ArmReplaceLiveTranslationEntry, 0x200));
+    ASSERT (((UINTN)ArmReplaceLiveTranslationEntry >> EFI_PAGE_SHIFT) == ((ArmReplaceLiveTranslationEntryEnd - 1) >> EFI_PAGE_SHIFT));
+
     //
     // The ArmReplaceLiveTranslationEntry () helper function may be invoked
     // with the MMU off so we have to ensure that it gets cleaned to the PoC
diff --git a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuPeiLibConstructor.c b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuPeiLibConstructor.c
index a0c9facdea..eb7011b3c5 100644
--- a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuPeiLibConstructor.c
+++ b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuPeiLibConstructor.c
@@ -24,6 +24,7 @@ ArmMmuPeiLibConstructor (
 {
   extern UINT32                       ArmReplaceLiveTranslationEntrySize;
   ARM_REPLACE_LIVE_TRANSLATION_ENTRY  ArmReplaceLiveTranslationEntryFunc;
+  UINTN                               ArmReplaceLiveTranslationEntryEnd;
   VOID                                *Hob;
 
   EFI_FV_FILE_INFO  FileInfo;
@@ -31,6 +32,16 @@ ArmMmuPeiLibConstructor (
 
   ASSERT (FileHandle != NULL);
 
+  ArmReplaceLiveTranslationEntryEnd = (UINTN)ArmReplaceLiveTranslationEntry + ArmReplaceLiveTranslationEntrySize;
+
+  //
+  // Align this routine to a log2 upper bound of its size, so that it is
+  // guaranteed not to cross a page or block boundary
+  // (see ArmMmuLibReplaceEntry.S).
+  //
+  ASSERT (IS_ALIGNED ((UINTN)ArmReplaceLiveTranslationEntry, 0x200));
+  ASSERT (((UINTN)ArmReplaceLiveTranslationEntry >> EFI_PAGE_SHIFT) == ((ArmReplaceLiveTranslationEntryEnd - 1) >> EFI_PAGE_SHIFT));
+
   Status = (*PeiServices)->FfsGetFileInfo (FileHandle, &FileInfo);
   ASSERT_EFI_ERROR (Status);
 
@@ -42,8 +53,7 @@ ArmMmuPeiLibConstructor (
   // when not executing in place.
   //
   if (((UINTN)FileInfo.Buffer <= (UINTN)ArmReplaceLiveTranslationEntry) &&
-      ((UINTN)FileInfo.Buffer + FileInfo.BufferSize >=
-       (UINTN)ArmReplaceLiveTranslationEntry + ArmReplaceLiveTranslationEntrySize))
+      ((UINTN)FileInfo.Buffer + FileInfo.BufferSize >= ArmReplaceLiveTranslationEntryEnd))
   {
     DEBUG ((DEBUG_INFO, "ArmMmuLib: skipping cache maintenance on XIP PEIM\n"));