From c779a246e961499c66a7ba1b63a6c26bdb21deec Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Tue, 12 Mar 2024 16:02:30 +0300
Subject: [PATCH] Ring3: Properly freed Ring3 resources.

---
 MdeModulePkg/Core/Dxe/DxeMain.h               |  2 +
 MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c       | 41 +++++++++++++------
 .../Core/Dxe/SysCall/Initialization.c         | 19 +++++----
 3 files changed, 41 insertions(+), 21 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/DxeMain.h b/MdeModulePkg/Core/Dxe/DxeMain.h
index f8fd68f866..f4082f184a 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain.h
+++ b/MdeModulePkg/Core/Dxe/DxeMain.h
@@ -273,6 +273,8 @@ extern LOADED_IMAGE_PRIVATE_DATA  *                mCurrentImage;
 
 extern RING3_DATA                        *gRing3Data;
 extern VOID                              *gRing3Interfaces;
+extern VOID                              *gCoreSysCallStackBase;
+extern VOID                              *gRing3CallStackBase;
 
 //
 // Service Initialization Functions
diff --git a/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c b/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
index 2eaec6f4d5..61cb28e6b9 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
+++ b/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
@@ -763,6 +763,35 @@ CoreExitBootServices (
 {
   EFI_STATUS  Status;
 
+  //
+  // Free resources allocated for Ring3.
+  //
+  if (gRing3Data != NULL) {
+    CoreFreePages (
+      (EFI_PHYSICAL_ADDRESS)gRing3Data,
+      EFI_SIZE_TO_PAGES (sizeof (RING3_DATA))
+    );
+    
+    gRing3Data = NULL;
+
+    CoreFreePages (
+      (EFI_PHYSICAL_ADDRESS)gRing3Interfaces,
+      RING3_INTERFACES_PAGES
+    );
+
+    CoreFreePages (
+      (EFI_PHYSICAL_ADDRESS)gCoreSysCallStackBase,
+      EFI_SIZE_TO_PAGES (USER_STACK_SIZE)
+    );
+
+    CoreFreePages (
+      (EFI_PHYSICAL_ADDRESS)gRing3CallStackBase,
+      EFI_SIZE_TO_PAGES (USER_STACK_SIZE)
+    );
+
+    FreeProtocolsList ();
+  }
+
   //
   // Notify other drivers of their last chance to use boot services
   // before the memory map is terminated.
@@ -840,18 +869,6 @@ CoreExitBootServices (
   //
   gRuntime->AtRuntime = TRUE;
 
-  CoreFreePages (
-    (EFI_PHYSICAL_ADDRESS)gRing3Data,
-    EFI_SIZE_TO_PAGES (sizeof (RING3_DATA))
-    );
-
-  CoreFreePages (
-    (EFI_PHYSICAL_ADDRESS)gRing3Interfaces,
-    RING3_INTERFACES_PAGES
-    );
-
-  FreeProtocolsList ();
-  
   return Status;
 }
 
diff --git a/MdeModulePkg/Core/Dxe/SysCall/Initialization.c b/MdeModulePkg/Core/Dxe/SysCall/Initialization.c
index 7ad0424a30..ec1b645acb 100644
--- a/MdeModulePkg/Core/Dxe/SysCall/Initialization.c
+++ b/MdeModulePkg/Core/Dxe/SysCall/Initialization.c
@@ -10,7 +10,9 @@
 #include <Register/Intel/ArchitecturalMsr.h>
 
 VOID        *gCoreSysCallStackTop;
+VOID        *gCoreSysCallStackBase;
 VOID        *gRing3CallStackTop;
+VOID        *gRing3CallStackBase;
 VOID        *gRing3EntryPoint;
 RING3_DATA  *gRing3Data;
 VOID        *gRing3Interfaces;
@@ -23,7 +25,6 @@ InitializeRing3 (
   )
 {
   EFI_STATUS              Status;
-  VOID                    *BaseOfStack;
   VOID                    *TopOfStack;
   UINTN                   SizeOfStack;
   UINT64                  Msr;
@@ -105,35 +106,35 @@ InitializeRing3 (
   //
   // Allocate 128KB for the Core SysCall Stack.
   //
-  BaseOfStack = AllocatePages (EFI_SIZE_TO_PAGES (USER_STACK_SIZE));
-  ASSERT (BaseOfStack != NULL);
+  gCoreSysCallStackBase = AllocatePages (EFI_SIZE_TO_PAGES (USER_STACK_SIZE));
+  ASSERT (gCoreSysCallStackBase != NULL);
 
   //
   // Compute the top of the allocated stack. Pre-allocate a UINTN for safety.
   //
-  TopOfStack = (VOID *)((UINTN)BaseOfStack + SizeOfStack - CPU_STACK_ALIGNMENT);
+  TopOfStack = (VOID *)((UINTN)gCoreSysCallStackBase + SizeOfStack - CPU_STACK_ALIGNMENT);
   TopOfStack = ALIGN_POINTER (TopOfStack, CPU_STACK_ALIGNMENT);
 
   gCoreSysCallStackTop = TopOfStack;
 
-  SetUefiImageMemoryAttributes ((UINTN)BaseOfStack, SizeOfStack, EFI_MEMORY_XP);
+  SetUefiImageMemoryAttributes ((UINTN)gCoreSysCallStackBase, SizeOfStack, EFI_MEMORY_XP);
   DEBUG ((DEBUG_ERROR, "Core: gCoreSysCallStackTop = %p\n", gCoreSysCallStackTop));
 
   //
   // Allocate 128KB for the User Stack.
   //
-  BaseOfStack = AllocatePages (EFI_SIZE_TO_PAGES (USER_STACK_SIZE));
-  ASSERT (BaseOfStack != NULL);
+  gRing3CallStackBase = AllocatePages (EFI_SIZE_TO_PAGES (USER_STACK_SIZE));
+  ASSERT (gRing3CallStackBase != NULL);
 
   //
   // Compute the top of the allocated stack. Pre-allocate a UINTN for safety.
   //
-  TopOfStack = (VOID *)((UINTN)BaseOfStack + SizeOfStack - CPU_STACK_ALIGNMENT);
+  TopOfStack = (VOID *)((UINTN)gRing3CallStackBase + SizeOfStack - CPU_STACK_ALIGNMENT);
   TopOfStack = ALIGN_POINTER (TopOfStack, CPU_STACK_ALIGNMENT);
 
   gRing3CallStackTop = TopOfStack;
 
-  SetUefiImageMemoryAttributes ((UINTN)BaseOfStack, SizeOfStack, EFI_MEMORY_XP | EFI_MEMORY_USER);
+  SetUefiImageMemoryAttributes ((UINTN)gRing3CallStackBase, SizeOfStack, EFI_MEMORY_XP | EFI_MEMORY_USER);
   DEBUG ((DEBUG_ERROR, "Core: gRing3CallStackTop = %p\n", gRing3CallStackTop));
 
   //