From 6da370bc1198a374cc589ba6dfd6a33e313be889 Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Mon, 9 Dec 2024 17:28:43 +0300
Subject: [PATCH] Ring3: Added MakeUserPageTableTemplate() for AARCH64.

---
 ArmPkg/Include/Library/ArmMmuLib.h            |  8 ++++
 .../Library/ArmMmuLib/AArch64/ArmMmuLibCore.c | 47 +++++++++++++++++++
 MdeModulePkg/Core/Dxe/DxeMain.inf             |  6 +++
 .../Dxe/SysCall/AARCH64/InitializeAARCH64.c   | 29 ++++++++++++
 4 files changed, 90 insertions(+)

diff --git a/ArmPkg/Include/Library/ArmMmuLib.h b/ArmPkg/Include/Library/ArmMmuLib.h
index 2ce948e8db..49410a8829 100644
--- a/ArmPkg/Include/Library/ArmMmuLib.h
+++ b/ArmPkg/Include/Library/ArmMmuLib.h
@@ -21,6 +21,14 @@ ArmConfigureMmu (
   OUT UINTN                         *TranslationTableSize  OPTIONAL
   );
 
+EFI_STATUS
+EFIAPI
+ArmMakeUserPageTableTemplate (
+  IN  ARM_MEMORY_REGION_DESCRIPTOR  *MemoryTable,
+  OUT UINTN                         *TranslationTableBase,
+  OUT UINTN                         *TranslationTableSize
+  );
+
 VOID
 EFIAPI
 ArmReplaceLiveTranslationEntry (
diff --git a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c
index 2e8eea05b5..11a4571dee 100644
--- a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c
+++ b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c
@@ -772,6 +772,53 @@ FreeTranslationTable:
   return Status;
 }
 
+EFI_STATUS
+EFIAPI
+ArmMakeUserPageTableTemplate (
+  IN  ARM_MEMORY_REGION_DESCRIPTOR  *MemoryTable,
+  OUT UINTN                         *TranslationTableBase,
+  OUT UINTN                         *TranslationTableSize
+  )
+{
+  VOID        *TranslationTable;
+  UINTN       MaxAddressBits;
+  UINT64      MaxAddress;
+  UINTN       T0SZ;
+  UINTN       RootTableEntryCount;
+  EFI_STATUS  Status;
+
+  MaxAddressBits = MIN (ArmGetPhysicalAddressBits (), MAX_VA_BITS);
+  MaxAddress     = LShiftU64 (1ULL, MaxAddressBits) - 1;
+
+  T0SZ                = 64 - MaxAddressBits;
+  RootTableEntryCount = GetRootTableEntryCount (T0SZ);
+
+  TranslationTable = AllocatePages (1);
+  if (TranslationTable == NULL) {
+    return EFI_OUT_OF_RESOURCES;
+  }
+
+  ZeroMem (TranslationTable, RootTableEntryCount * sizeof (UINT64));
+
+  Status = UpdateRegionMapping (
+             MemoryTable->VirtualBase,
+             MemoryTable->Length,
+             ArmMemoryAttributeToPageAttribute (MemoryTable->Attributes),
+             0,
+             TranslationTable,
+             FALSE
+             );
+  if (EFI_ERROR (Status)) {
+    FreePages (TranslationTable, 1);
+    return Status;
+  }
+
+  *TranslationTableBase = (UINTN)TranslationTable;
+  *TranslationTableSize = RootTableEntryCount * sizeof (UINT64);
+
+  return EFI_SUCCESS;
+}
+
 RETURN_STATUS
 EFIAPI
 ArmMmuBaseLibConstructor (
diff --git a/MdeModulePkg/Core/Dxe/DxeMain.inf b/MdeModulePkg/Core/Dxe/DxeMain.inf
index c608cd4719..a6b4741e30 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain.inf
+++ b/MdeModulePkg/Core/Dxe/DxeMain.inf
@@ -151,6 +151,9 @@
   gEfiGlobalVariableGuid                        ## SOMETIMES_CONSUMES   ## SysCall
   gEarlyPL011BaseAddressGuid                    ## CONSUMES
 
+[Guids.ARM, Guids.AARCH64]
+  gArmVirtSystemMemorySizeGuid                  ## CONSUMES
+
 [Ppis]
   gEfiVectorHandoffInfoPpiGuid                  ## UNDEFINED # HOB
 
@@ -231,6 +234,9 @@
   gUefiOvmfPkgTokenSpaceGuid.PcdOvmfWorkAreaBase                            ## CONSUMES
   gUefiOvmfPkgTokenSpaceGuid.PcdOvmfWorkAreaSize                            ## CONSUMES
 
+[Pcd.ARM, Pcd.AARCH64]
+  gArmTokenSpaceGuid.PcdSystemMemoryBase                                    ## CONSUMES
+
 # [Hob]
 # RESOURCE_DESCRIPTOR   ## CONSUMES
 # MEMORY_ALLOCATION     ## CONSUMES
diff --git a/MdeModulePkg/Core/Dxe/SysCall/AARCH64/InitializeAARCH64.c b/MdeModulePkg/Core/Dxe/SysCall/AARCH64/InitializeAARCH64.c
index 85a0edf599..54dfd33a6f 100644
--- a/MdeModulePkg/Core/Dxe/SysCall/AARCH64/InitializeAARCH64.c
+++ b/MdeModulePkg/Core/Dxe/SysCall/AARCH64/InitializeAARCH64.c
@@ -7,6 +7,7 @@
 
 #include <Chipset/AArch64.h>
 #include <Library/ArmLib.h>
+#include <Library/ArmMmuLib.h>
 #include <Library/DefaultExceptionHandlerLib.h>
 
 #include "DxeMain.h"
@@ -89,6 +90,34 @@ SysCallBootService (
   return Status;
 }
 
+VOID
+EFIAPI
+MakeUserPageTableTemplate (
+  OUT UINTN  *UserPageTableTemplate,
+  OUT UINTN  *UserPageTableTemplateSize
+  )
+{
+  ARM_MEMORY_REGION_DESCRIPTOR  Descriptor;
+  VOID                          *MemorySizeHob;
+
+  MemorySizeHob = GetFirstGuidHob (&gArmVirtSystemMemorySizeGuid);
+  ASSERT (MemorySizeHob != NULL);
+  if (MemorySizeHob == NULL) {
+    return;
+  }
+
+  Descriptor.PhysicalBase = PcdGet64 (PcdSystemMemoryBase);
+  Descriptor.VirtualBase  = Descriptor.PhysicalBase;
+  Descriptor.Length       = *(UINT64 *)GET_GUID_HOB_DATA (MemorySizeHob);
+  Descriptor.Attributes   = ARM_MEMORY_REGION_ATTRIBUTE_WRITE_BACK;
+
+  ArmMakeUserPageTableTemplate (
+    &Descriptor,
+    UserPageTableTemplate,
+    UserPageTableTemplateSize
+    );
+}
+
 VOID
 EFIAPI
 InitializeMsr (