From b6a7e2b3be7e101094086a194e00b8f56f12d3b0 Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Thu, 30 Jan 2025 15:45:49 +0300
Subject: [PATCH] Ring3: Properly refactored MemoryPoolLib to support User
 spaces.

---
 MdeModulePkg/Core/Dxe/DxeRing3/DxeRing3.c     |  2 +-
 .../Core/Dxe/DxeRing3/Ring3UefiBootServices.c | 14 ++++----
 MdeModulePkg/Core/Dxe/Gcd/Gcd.c               |  2 +-
 MdeModulePkg/Include/Library/MemoryPoolLib.h  |  8 ++++-
 MdeModulePkg/Library/MemoryPoolLib/Pool.c     | 34 ++++++++++++++-----
 5 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/DxeRing3/DxeRing3.c b/MdeModulePkg/Core/Dxe/DxeRing3/DxeRing3.c
index 579d86f9df..5bbc706bf5 100644
--- a/MdeModulePkg/Core/Dxe/DxeRing3/DxeRing3.c
+++ b/MdeModulePkg/Core/Dxe/DxeRing3/DxeRing3.c
@@ -259,7 +259,7 @@ Ring3Initialization (
   gBS = &mBootServices;
   gRT = &mRuntimeServices;
 
-  CoreInitializePool ();
+  CoreInitializePool (FALSE);
 
   return EFI_SUCCESS;
 }
diff --git a/MdeModulePkg/Core/Dxe/DxeRing3/Ring3UefiBootServices.c b/MdeModulePkg/Core/Dxe/DxeRing3/Ring3UefiBootServices.c
index 8047fa314e..5cc2c3dc53 100644
--- a/MdeModulePkg/Core/Dxe/DxeRing3/Ring3UefiBootServices.c
+++ b/MdeModulePkg/Core/Dxe/DxeRing3/Ring3UefiBootServices.c
@@ -18,7 +18,7 @@
 
 BOOLEAN  mOnGuarding = FALSE;
 
-STATIC UINTN  mMemoryTypes[EfiMaxMemoryType];
+STATIC UINTN  mMemoryTypes[MAX_MEMORY_TYPE];
 STATIC UINTN  mNumberOfUsers = 0;
 
 STATIC
@@ -32,18 +32,16 @@ GetMemoryType (
 
   for (Index = 0; Index < mNumberOfUsers; ++Index) {
     if (mMemoryTypes[Index] == UserPageTable) {
-      break;
+      return Index;
     }
   }
 
-  if (Index == mNumberOfUsers) {
-    ++mNumberOfUsers;
-    mMemoryTypes[Index] = UserPageTable;
-  }
+  ASSERT (mNumberOfUsers < MAX_MEMORY_TYPE);
 
-  ASSERT (mNumberOfUsers <= EfiMaxMemoryType);
+  mMemoryTypes[mNumberOfUsers] = UserPageTable;
+  ++mNumberOfUsers;
 
-  return Index;
+  return (mNumberOfUsers - 1);
 }
 
 STATIC
diff --git a/MdeModulePkg/Core/Dxe/Gcd/Gcd.c b/MdeModulePkg/Core/Dxe/Gcd/Gcd.c
index d43970665f..4c4414cdde 100644
--- a/MdeModulePkg/Core/Dxe/Gcd/Gcd.c
+++ b/MdeModulePkg/Core/Dxe/Gcd/Gcd.c
@@ -2261,7 +2261,7 @@ CoreInitializeMemoryServices (
   // Initialize the spin locks and maps in the memory services.
   // Also fill in the memory services into the EFI Boot Services Table
   //
-  CoreInitializePool ();
+  CoreInitializePool (TRUE);
 
   //
   // Initialize Local Variables
diff --git a/MdeModulePkg/Include/Library/MemoryPoolLib.h b/MdeModulePkg/Include/Library/MemoryPoolLib.h
index 9b3e5b257f..c3ee7c8e62 100644
--- a/MdeModulePkg/Include/Library/MemoryPoolLib.h
+++ b/MdeModulePkg/Include/Library/MemoryPoolLib.h
@@ -36,13 +36,19 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #define GUARD_HEAP_TYPE_ALL         \
         (GUARD_HEAP_TYPE_PAGE|GUARD_HEAP_TYPE_POOL|GUARD_HEAP_TYPE_FREED)
 
+#define MAX_MEMORY_TYPE        100
+
 /**
   Called to initialize the pool.
 
+  @param  IsCore                 Selects between the purposes of mPoolHead array.
+                                 In DxeCore each element describes EFI_MEMORY_TYPE.
+                                 In DxeRing3 each element describes User space.
+
 **/
 VOID
 CoreInitializePool (
-  VOID
+  IN BOOLEAN  IsCore
   );
 
 /**
diff --git a/MdeModulePkg/Library/MemoryPoolLib/Pool.c b/MdeModulePkg/Library/MemoryPoolLib/Pool.c
index 47d744c052..190ea17f55 100644
--- a/MdeModulePkg/Library/MemoryPoolLib/Pool.c
+++ b/MdeModulePkg/Library/MemoryPoolLib/Pool.c
@@ -80,7 +80,9 @@ typedef struct {
 //
 // Pool header for each memory type.
 //
-POOL  mPoolHead[EfiMaxMemoryType];
+POOL  mPoolHead[MAX_MEMORY_TYPE];
+
+STATIC BOOLEAN  mIsCore;
 
 //
 // List of pool header to search for the appropriate memory type.
@@ -185,12 +187,14 @@ GetPoolIndexFromSize (
 **/
 VOID
 CoreInitializePool (
-  VOID
+  IN BOOLEAN  IsCore
   )
 {
   UINTN  Type;
   UINTN  Index;
 
+  mIsCore = IsCore;
+
   for (Type = 0; Type < EfiMaxMemoryType; Type++) {
     mPoolHead[Type].Signature  = 0;
     mPoolHead[Type].Used       = 0;
@@ -219,6 +223,14 @@ LookupPoolHead (
   POOL        *Pool;
   UINTN       Index;
 
+  if (!mIsCore) {
+    if ((UINT32)MemoryType < MAX_MEMORY_TYPE) {
+      return &mPoolHead[MemoryType];
+    }
+
+    return NULL;
+  }
+
   if ((UINT32)MemoryType < EfiMaxMemoryType) {
     return &mPoolHead[MemoryType];
   }
@@ -285,8 +297,8 @@ CoreInternalAllocatePool (
   //
   // If it's not a valid type, fail it
   //
-  if (((PoolType >= EfiMaxMemoryType) && (PoolType < MEMORY_TYPE_OEM_RESERVED_MIN)) ||
-      (PoolType == EfiConventionalMemory) || (PoolType == EfiPersistentMemory) || (PoolType == EfiUnacceptedMemoryType))
+  if (mIsCore && (((PoolType >= EfiMaxMemoryType) && (PoolType < MEMORY_TYPE_OEM_RESERVED_MIN)) ||
+      (PoolType == EfiConventionalMemory) || (PoolType == EfiPersistentMemory) || (PoolType == EfiUnacceptedMemoryType)))
   {
     return EFI_INVALID_PARAMETER;
   }
@@ -395,10 +407,11 @@ CoreAllocatePoolI (
 
   ASSERT_LOCKED (&mPoolMemoryLock);
 
-  if ((PoolType == EfiReservedMemoryType) ||
+  if (mIsCore &&
+     ((PoolType == EfiReservedMemoryType) ||
       (PoolType == EfiACPIMemoryNVS) ||
       (PoolType == EfiRuntimeServicesCode) ||
-      (PoolType == EfiRuntimeServicesData))
+      (PoolType == EfiRuntimeServicesData)))
   {
     Granularity = RUNTIME_PAGE_ALLOCATION_GRANULARITY;
   } else {
@@ -724,10 +737,11 @@ CoreFreePoolI (
   Pool->Used -= Size;
   DEBUG ((DEBUG_POOL, "FreePool: %p (len %lx) %,ld\n", Head->Data, (UINT64)(Head->Size - POOL_OVERHEAD), (UINT64)Pool->Used));
 
-  if ((Head->Type == EfiReservedMemoryType) ||
+  if (mIsCore &&
+     ((Head->Type == EfiReservedMemoryType) ||
       (Head->Type == EfiACPIMemoryNVS) ||
       (Head->Type == EfiRuntimeServicesCode) ||
-      (Head->Type == EfiRuntimeServicesData))
+      (Head->Type == EfiRuntimeServicesData)))
   {
     Granularity = RUNTIME_PAGE_ALLOCATION_GRANULARITY;
   } else {
@@ -898,6 +912,10 @@ IsMemoryTypeToGuard (
   UINT64  TestBit;
   UINT64  ConfigBit;
 
+  if (!mIsCore) {
+    return FALSE;
+  }
+
   if (AllocateType == AllocateAddress) {
     return FALSE;
   }