From d4f5ae14aa1d51321db0a736b60e6c34fd2527ab Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Thu, 29 Feb 2024 20:53:02 +0300
Subject: [PATCH] Ring3: Refactored out AllocateRing3Copy().

---
 MdeModulePkg/Core/Dxe/DxeMain.h               |  8 ---
 MdeModulePkg/Core/Dxe/Image/Image.c           | 24 -------
 MdeModulePkg/Core/Dxe/SysCall/BootServices.c  | 12 ++--
 .../Core/Dxe/SysCall/SupportedProtocols.c     | 67 ++++++++++---------
 4 files changed, 41 insertions(+), 70 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/DxeMain.h b/MdeModulePkg/Core/Dxe/DxeMain.h
index 973b2f2a45..328890495f 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain.h
+++ b/MdeModulePkg/Core/Dxe/DxeMain.h
@@ -1171,14 +1171,6 @@ CoreAllocatePages (
   IN OUT EFI_PHYSICAL_ADDRESS  *Memory
   );
 
-VOID *
-EFIAPI
-AllocateRing3Copy (
-  IN VOID    *Source,
-  IN UINT32  AllocationSize,
-  IN UINT32  CopySize
-  );
-
 /**
   Frees previous allocated pages.
 
diff --git a/MdeModulePkg/Core/Dxe/Image/Image.c b/MdeModulePkg/Core/Dxe/Image/Image.c
index 069ffefa0c..51c8edc820 100644
--- a/MdeModulePkg/Core/Dxe/Image/Image.c
+++ b/MdeModulePkg/Core/Dxe/Image/Image.c
@@ -1570,30 +1570,6 @@ CoreLoadImage (
   return Status;
 }
 
-VOID *
-EFIAPI
-AllocateRing3Copy (
-  IN VOID    *Source,
-  IN UINT32  AllocationSize,
-  IN UINT32  CopySize
-  )
-{
-  EFI_STATUS Status;
-  VOID       *MemoryRing3;
-
-  Status = CoreAllocatePool (EfiRing3MemoryType, AllocationSize, &MemoryRing3);
-  if (EFI_ERROR (Status)) {
-    DEBUG ((DEBUG_ERROR, "Core: Failed to allocate %d bytes for Ring3.\n", AllocationSize));
-    return NULL;
-  }
-
-  ASSERT (CopySize <= AllocationSize);
-
-  CopyMem (MemoryRing3, Source, CopySize);
-
-  return MemoryRing3;
-}
-
 EFI_STATUS
 EFIAPI
 InitializeRing3 (
diff --git a/MdeModulePkg/Core/Dxe/SysCall/BootServices.c b/MdeModulePkg/Core/Dxe/SysCall/BootServices.c
index b17f748676..d494ae7622 100644
--- a/MdeModulePkg/Core/Dxe/SysCall/BootServices.c
+++ b/MdeModulePkg/Core/Dxe/SysCall/BootServices.c
@@ -103,10 +103,7 @@ PrepareRing3Interface (
 
   Ring3Limit = (UINTN)gRing3Interfaces + EFI_PAGES_TO_SIZE (RING3_INTERFACES_PAGES);
 
-  ASSERT ((mRing3InterfacePointer + sizeof (EFI_GUID) + CoreSize) <= Ring3Limit);
-
-  CopyMem ((VOID *)mRing3InterfacePointer, (VOID *)Guid, sizeof (EFI_GUID));
-  mRing3InterfacePointer += sizeof (EFI_GUID);
+  ASSERT ((mRing3InterfacePointer + CoreSize) <= Ring3Limit);
 
   Ring3Interface = (VOID *)mRing3InterfacePointer;
 
@@ -265,7 +262,12 @@ CallBootService (
         Status = FindGuid ((EFI_GUID *)UserArgList[Index], (EFI_GUID **)&CoreArgList[Index], &MemoryCoreSize);
         if (EFI_ERROR (Status)) {
           EnableSMAP ();
-          //TODO: Free CoreArgList.
+
+          while (Index > 0) {
+            FreePool (CoreArgList[Index - 1]);
+            Index -= 2;
+          }
+
           return Status;
         }
 
diff --git a/MdeModulePkg/Core/Dxe/SysCall/SupportedProtocols.c b/MdeModulePkg/Core/Dxe/SysCall/SupportedProtocols.c
index edca2fd49f..4639fe6e11 100644
--- a/MdeModulePkg/Core/Dxe/SysCall/SupportedProtocols.c
+++ b/MdeModulePkg/Core/Dxe/SysCall/SupportedProtocols.c
@@ -77,6 +77,34 @@ GoToRing3 (
   return Status;
 }
 
+STATIC
+EFIAPI
+VOID *
+Ring3Copy (
+  IN VOID    *Core,
+  IN UINT32  Size
+  )
+{
+  EFI_STATUS  Status;
+  VOID        *Ring3;
+
+  Status = CoreAllocatePages (
+             AllocateAnyPages,
+             EfiRing3MemoryType,
+             1,
+             (EFI_PHYSICAL_ADDRESS *)&Ring3
+             );
+  if (EFI_ERROR (Status)) {
+    return NULL;
+  }
+
+  DisableSMAP ();
+  CopyMem (Ring3, Core, Size);
+  EnableSMAP ();
+
+  return Ring3;
+}
+
 EFI_STATUS
 EFIAPI
 CoreDriverBindingSupported (
@@ -87,17 +115,10 @@ CoreDriverBindingSupported (
 {
   EFI_STATUS  Status;
 
-  DisableSMAP ();
-  This = AllocateRing3Copy (
-           This,
-           sizeof (EFI_DRIVER_BINDING_PROTOCOL),
-           sizeof (EFI_DRIVER_BINDING_PROTOCOL)
-           );
+  This = Ring3Copy (This, sizeof (EFI_DRIVER_BINDING_PROTOCOL));
   if (This == NULL) {
-    EnableSMAP ();
     return EFI_OUT_OF_RESOURCES;
   }
-  EnableSMAP ();
 
   Status = GoToRing3 (
              3,
@@ -107,9 +128,7 @@ CoreDriverBindingSupported (
              RemainingDevicePath
              );
 
-  DisableSMAP ();
-  FreePool (This);
-  EnableSMAP ();
+  CoreFreePages ((EFI_PHYSICAL_ADDRESS)This, 1);
 
   return Status;
 }
@@ -124,17 +143,10 @@ CoreDriverBindingStart (
 {
   EFI_STATUS  Status;
 
-  DisableSMAP ();
-  This = AllocateRing3Copy (
-           This,
-           sizeof (EFI_DRIVER_BINDING_PROTOCOL),
-           sizeof (EFI_DRIVER_BINDING_PROTOCOL)
-           );
+  This = Ring3Copy (This, sizeof (EFI_DRIVER_BINDING_PROTOCOL));
   if (This == NULL) {
-    EnableSMAP ();
     return EFI_OUT_OF_RESOURCES;
   }
-  EnableSMAP ();
 
   Status = GoToRing3 (
              3,
@@ -144,9 +156,7 @@ CoreDriverBindingStart (
              RemainingDevicePath
              );
 
-  DisableSMAP ();
-  FreePool (This);
-  EnableSMAP ();
+  CoreFreePages ((EFI_PHYSICAL_ADDRESS)This, 1);
 
   return Status;
 }
@@ -162,17 +172,10 @@ CoreDriverBindingStop (
 {
   EFI_STATUS  Status;
 
-  DisableSMAP ();
-  This = AllocateRing3Copy (
-           This,
-           sizeof (EFI_DRIVER_BINDING_PROTOCOL),
-           sizeof (EFI_DRIVER_BINDING_PROTOCOL)
-           );
+  This = Ring3Copy (This, sizeof (EFI_DRIVER_BINDING_PROTOCOL));
   if (This == NULL) {
-    EnableSMAP ();
     return EFI_OUT_OF_RESOURCES;
   }
-  EnableSMAP ();
 
   Status = GoToRing3 (
              4,
@@ -183,9 +186,7 @@ CoreDriverBindingStop (
              ChildHandleBuffer
              );
 
-  DisableSMAP ();
-  FreePool (This);
-  EnableSMAP ();
+  CoreFreePages ((EFI_PHYSICAL_ADDRESS)This, 1);
 
   return Status;
 }