From 43483580fbf1221fc217a5b72ce2880f1668619b Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Mon, 11 Mar 2024 17:00:20 +0300
Subject: [PATCH] Ring3: Fixed Ring3LocateHandleBuffer().

---
 MdeModulePkg/Core/Dxe/DxeRing3/Ring3.h        | 36 ---------------
 .../Core/Dxe/DxeRing3/Ring3UefiBootServices.c | 45 +++++++++++++++----
 MdeModulePkg/Core/Dxe/SysCall/BootServices.c  |  2 +-
 3 files changed, 38 insertions(+), 45 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/DxeRing3/Ring3.h b/MdeModulePkg/Core/Dxe/DxeRing3/Ring3.h
index fd5a605505..235af1eddb 100644
--- a/MdeModulePkg/Core/Dxe/DxeRing3/Ring3.h
+++ b/MdeModulePkg/Core/Dxe/DxeRing3/Ring3.h
@@ -131,42 +131,6 @@ Ring3GetMemoryMap (
   OUT UINT32                    *DescriptorVersion
   );
 
-/**
-  Allocate pool of a particular type.
-
-  @param  PoolType               Type of pool to allocate
-  @param  Size                   The amount of pool to allocate
-  @param  Buffer                 The address to return a pointer to the allocated
-                                 pool
-
-  @retval EFI_INVALID_PARAMETER  PoolType not valid or Buffer is NULL
-  @retval EFI_OUT_OF_RESOURCES   Size exceeds max pool size or allocation failed.
-  @retval EFI_SUCCESS            Pool successfully allocated.
-
-**/
-EFI_STATUS
-EFIAPI
-Ring3AllocatePool (
-  IN EFI_MEMORY_TYPE  PoolType,
-  IN UINTN            Size,
-  OUT VOID            **Buffer
-  );
-
-/**
-  Frees pool.
-
-  @param  Buffer                 The allocated pool entry to free
-
-  @retval EFI_INVALID_PARAMETER  Buffer is not a valid value.
-  @retval EFI_SUCCESS            Pool successfully freed.
-
-**/
-EFI_STATUS
-EFIAPI
-Ring3FreePool (
-  IN VOID  *Buffer
-  );
-
 /**
   Creates an event.
 
diff --git a/MdeModulePkg/Core/Dxe/DxeRing3/Ring3UefiBootServices.c b/MdeModulePkg/Core/Dxe/DxeRing3/Ring3UefiBootServices.c
index cda6f373be..6a23b4ba79 100644
--- a/MdeModulePkg/Core/Dxe/DxeRing3/Ring3UefiBootServices.c
+++ b/MdeModulePkg/Core/Dxe/DxeRing3/Ring3UefiBootServices.c
@@ -12,6 +12,7 @@
 
 #include <Library/BaseMemoryLib.h>
 #include <Library/DebugLib.h>
+#include <Library/MemoryPoolLib.h>
 
 #include "Ring3.h"
 
@@ -579,14 +580,42 @@ Ring3LocateHandleBuffer (
   OUT EFI_HANDLE             **Buffer
   )
 {
-  return SysCall (
-           SysCallLocateHandleBuffer,
-           SearchType,
-           Protocol,
-           SearchKey,
-           NumberHandles,
-           Buffer
-           );
+  EFI_STATUS  Status;
+  EFI_STATUS  StatusBS;
+  VOID        *Pool;
+  UINTN       PoolSize;
+
+  StatusBS = SysCall (
+               SysCallLocateHandleBuffer,
+               SearchType,
+               Protocol,
+               SearchKey,
+               NumberHandles,
+               Buffer
+               );
+
+  if ((NumberHandles != NULL) && (Buffer != NULL) && (*Buffer != NULL)) {
+    PoolSize = *NumberHandles * sizeof (EFI_HANDLE *);
+
+    Status = CoreAllocatePool (EfiRing3MemoryType, PoolSize, &Pool);
+    if (EFI_ERROR (Status)) {
+      return Status;
+    }
+
+    CopyMem (Pool, *Buffer, PoolSize);
+
+    Status = Ring3FreePages (
+               (EFI_PHYSICAL_ADDRESS)*Buffer,
+               EFI_SIZE_TO_PAGES (PoolSize)
+               );
+    if (EFI_ERROR (Status)) {
+      return Status;
+    }
+
+    *Buffer = Pool;
+  }
+
+  return StatusBS;
 }
 
 EFI_STATUS
diff --git a/MdeModulePkg/Core/Dxe/SysCall/BootServices.c b/MdeModulePkg/Core/Dxe/SysCall/BootServices.c
index 212204e5cb..320a4b4235 100644
--- a/MdeModulePkg/Core/Dxe/SysCall/BootServices.c
+++ b/MdeModulePkg/Core/Dxe/SysCall/BootServices.c
@@ -576,7 +576,7 @@ CallBootService (
         *(UINTN *)UserRsp->Arguments[4] = Argument4;
       }
 
-      if ((UINTN *)UserRsp->Arguments[5] != NULL) {
+      if ((EFI_HANDLE **)UserRsp->Arguments[5] != NULL) {
         gCpu->GetMemoryAttributes (gCpu, (EFI_PHYSICAL_ADDRESS)UserRsp->Arguments[5], &Attributes);
         ASSERT ((Attributes & EFI_MEMORY_USER) != 0);
         gCpu->GetMemoryAttributes (gCpu, (EFI_PHYSICAL_ADDRESS)(UserRsp->Arguments[5] + sizeof (EFI_HANDLE *) - 1), &Attributes);