From 2e86c435d5ffd6475f14b3d3cb1e6f2aee5dd350 Mon Sep 17 00:00:00 2001
From: vit9696 <vit9696@users.noreply.github.com>
Date: Sat, 26 Oct 2019 14:26:29 +0300
Subject: [PATCH] OpenCoreMisc: Implement ballooning stubs

---
 Include/OpenCore.h               | 15 +++++++++++++++
 Platform/OpenCore/OpenCoreMisc.c |  1 +
 Platform/OpenCore/OpenCoreUefi.c | 23 ++++++++++++-----------
 3 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/Include/OpenCore.h b/Include/OpenCore.h
index b068aee1..5dc1594d 100644
--- a/Include/OpenCore.h
+++ b/Include/OpenCore.h
@@ -15,6 +15,7 @@
 #ifndef OPEN_CORE_H
 #define OPEN_CORE_H
 
+#include <Library/OcBootManagementLib.h>
 #include <Library/OcConfigurationLib.h>
 #include <Library/OcCpuLib.h>
 #include <Library/OcCryptoLib.h>
@@ -175,6 +176,8 @@ OcLoadUefiSupport (
   upon changing screen resolution.
 
   @param[in]  Config    OpenCore configuration.
+
+  @retval TRUE when required.
 **/
 BOOLEAN
 OcShouldReconnectConsoleOnResolutionChange (
@@ -191,6 +194,18 @@ OcMiscGetVersionString (
   VOID
   );
 
+/**
+  Get ballooning handler for memory allocation protections.
+
+  @param[in]  Config    OpenCore configuration.
+
+  @retval Handler address or NULL.
+**/
+OC_BALLOON_ALLOC
+OcGetBallooningHandler (
+  IN  OC_GLOBAL_CONFIG  *Config
+  );
+
 /**
   Load early miscellaneous support like configuration.
 
diff --git a/Platform/OpenCore/OpenCoreMisc.c b/Platform/OpenCore/OpenCoreMisc.c
index b15a8464..b2c36f69 100644
--- a/Platform/OpenCore/OpenCoreMisc.c
+++ b/Platform/OpenCore/OpenCoreMisc.c
@@ -525,6 +525,7 @@ OcMiscBoot (
   Context->CustomRead         = OcToolLoadEntry;
   Context->PrivilegeContext   = Privilege;
   Context->RequestPrivilege   = OcShowSimplePasswordRequest;
+  Context->BalloonAllocator   = OcGetBallooningHandler (Config);
 
   if ((Config->Misc.Security.ExposeSensitiveData & OCS_EXPOSE_VERSION_UI) != 0) {
     Context->TitleSuffix      = OcMiscGetVersionString ();
diff --git a/Platform/OpenCore/OpenCoreUefi.c b/Platform/OpenCore/OpenCoreUefi.c
index c5f39a46..71f6b661 100644
--- a/Platform/OpenCore/OpenCoreUefi.c
+++ b/Platform/OpenCore/OpenCoreUefi.c
@@ -412,6 +412,18 @@ OcShouldReconnectConsoleOnResolutionChange (
   return Config->Uefi.Quirks.ReconnectOnResChange;
 }
 
+OC_BALLOON_ALLOC
+OcGetBallooningHandler (
+  IN  OC_GLOBAL_CONFIG  *Config
+  )
+{
+  if (Config->Uefi.Quirks.AvoidHighAlloc) {
+    return OcHandleKernelProtectionZone;
+  }
+
+  return NULL;
+}
+
 VOID
 OcLoadBooterUefiSupport (
   IN OC_GLOBAL_CONFIG  *Config
@@ -506,17 +518,6 @@ OcLoadUefiSupport (
     &Config->Uefi.Quirks.RequestBootVarRouting
     );
 
-  //
-  // Inform allocations when we want to use lower memory only.
-  //
-  gRT->SetVariable (
-    OC_AVOID_HIGH_ALLOC_VARIABLE_NAME,
-    &gOcVendorVariableGuid,
-    OPEN_CORE_INT_NVRAM_ATTR,
-    sizeof (Config->Uefi.Quirks.AvoidHighAlloc),
-    &Config->Uefi.Quirks.AvoidHighAlloc
-    );
-
   if (Config->Uefi.Quirks.ReleaseUsbOwnership
     || Config->Uefi.Quirks.ExitBootServicesDelay > 0
     || AgiExitBs) {