From 6d0c49e84cdc7477e5052dde09278c8ec1e2c8db Mon Sep 17 00:00:00 2001
From: vit9696 <vit9696@users.noreply.github.com>
Date: Sun, 8 Nov 2020 20:11:43 +0300
Subject: [PATCH] OcMemoryLib: Allow overriding page allocator for allocate
 from top

---
 Include/Acidanthera/Library/OcMemoryLib.h    | 6 ++++--
 Library/OcAfterBootCompatLib/KernelSupport.c | 1 +
 Library/OcMemoryLib/MemoryAlloc.c            | 7 ++++---
 Library/OcMemoryLib/MemoryMap.c              | 1 +
 Library/OcMemoryLib/VirtualMemory.c          | 1 +
 5 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/Include/Acidanthera/Library/OcMemoryLib.h b/Include/Acidanthera/Library/OcMemoryLib.h
index f51913f0..dbadb2f1 100644
--- a/Include/Acidanthera/Library/OcMemoryLib.h
+++ b/Include/Acidanthera/Library/OcMemoryLib.h
@@ -221,6 +221,7 @@ VOID
   @param[in]      Pages            Amount of pages to allocate.
   @param[in,out]  Memory           Top address for input, allocated address for output.
   @param[in]      GetMemoryMap     Custom GetMemoryMap implementation to use, optional.
+  @param[in]      AllocatePages    Custom AllocatePages implementation to use, optional.
   @param[in]      CheckRange       Handler allowing to not allocate select ranges, optional.
 
   @retval EFI_SUCCESS on successful allocation.
@@ -230,8 +231,9 @@ OcAllocatePagesFromTop (
   IN     EFI_MEMORY_TYPE         MemoryType,
   IN     UINTN                   Pages,
   IN OUT EFI_PHYSICAL_ADDRESS    *Memory,
-  IN     EFI_GET_MEMORY_MAP      GetMemoryMap  OPTIONAL,
-  IN     CHECK_ALLOCATION_RANGE  CheckRange  OPTIONAL
+  IN     EFI_GET_MEMORY_MAP      GetMemoryMap   OPTIONAL,
+  IN     EFI_ALLOCATE_PAGES      AllocatePages  OPTIONAL,
+  IN     CHECK_ALLOCATION_RANGE  CheckRange     OPTIONAL
   );
 
 /**
diff --git a/Library/OcAfterBootCompatLib/KernelSupport.c b/Library/OcAfterBootCompatLib/KernelSupport.c
index b9a5acb1..5a49c6b1 100644
--- a/Library/OcAfterBootCompatLib/KernelSupport.c
+++ b/Library/OcAfterBootCompatLib/KernelSupport.c
@@ -536,6 +536,7 @@ AppleMapPrepareBooterState (
         EFI_SIZE_TO_PAGES (gST->Hdr.HeaderSize),
         &BootCompat->KernelState.SysTableRtArea,
         GetMemoryMap,
+        NULL,
         NULL
         );
       if (EFI_ERROR (Status)) {
diff --git a/Library/OcMemoryLib/MemoryAlloc.c b/Library/OcMemoryLib/MemoryAlloc.c
index ce8b6dd4..0440dec5 100644
--- a/Library/OcMemoryLib/MemoryAlloc.c
+++ b/Library/OcMemoryLib/MemoryAlloc.c
@@ -28,8 +28,9 @@ OcAllocatePagesFromTop (
   IN     EFI_MEMORY_TYPE         MemoryType,
   IN     UINTN                   Pages,
   IN OUT EFI_PHYSICAL_ADDRESS    *Memory,
-  IN     EFI_GET_MEMORY_MAP      GetMemoryMap,
-  IN     CHECK_ALLOCATION_RANGE  CheckRange  OPTIONAL
+  IN     EFI_GET_MEMORY_MAP      GetMemoryMap   OPTIONAL,
+  IN     EFI_ALLOCATE_PAGES      AllocatePages  OPTIONAL,
+  IN     CHECK_ALLOCATION_RANGE  CheckRange     OPTIONAL
   )
 {
   EFI_STATUS              Status;
@@ -90,7 +91,7 @@ OcAllocatePagesFromTop (
         continue;
       }
 
-      Status = gBS->AllocatePages (
+      Status = (AllocatePages != NULL ? AllocatePages : gBS->AllocatePages) (
         AllocateAddress,
         MemoryType,
         Pages,
diff --git a/Library/OcMemoryLib/MemoryMap.c b/Library/OcMemoryLib/MemoryMap.c
index 54bfb9f3..374cfbbd 100644
--- a/Library/OcMemoryLib/MemoryMap.c
+++ b/Library/OcMemoryLib/MemoryMap.c
@@ -249,6 +249,7 @@ OcGetCurrentMemoryMapAlloc (
         (UINTN) *TopMemory,
         &MemoryMapAlloc,
         GetMemoryMap,
+        NULL,
         NULL
         );
 
diff --git a/Library/OcMemoryLib/VirtualMemory.c b/Library/OcMemoryLib/VirtualMemory.c
index 60484c12..c97738bd 100755
--- a/Library/OcMemoryLib/VirtualMemory.c
+++ b/Library/OcMemoryLib/VirtualMemory.c
@@ -157,6 +157,7 @@ VmAllocateMemoryPool (
     NumPages,
     &Addr,
     GetMemoryMap,
+    NULL,
     NULL
     );