OcAppleBootCompatLib: Update memory attribute table on memory free

Library/OcAppleBootCompatLib/BootCompatInternal.h

@@ -142,16 +142,26 @@ typedef struct UEFI_SERVICES_POINTERS_ {
   ///
   EFI_ALLOCATE_PAGES          AllocatePages;
   ///
-  /// Original pool allocator. We override it to fix memory
-  /// attributes table as it is updated after pool alloc.
+  /// Original page deallocator. We override it to fix memory
+  /// attributes table as it is updated after page dealloc.
   ///
-  EFI_ALLOCATE_POOL           AllocatePool;
+  EFI_FREE_PAGES              FreePages;
   ///
   /// Original memory map function. We override it to make
   /// memory map shrinking and CSM region protection.
   ///
   EFI_GET_MEMORY_MAP          GetMemoryMap;
   ///
+  /// Original pool allocator. We override it to fix memory
+  /// attributes table as it is updated after pool alloc.
+  ///
+  EFI_ALLOCATE_POOL           AllocatePool;
+  ///
+  /// Original pool deallocator. We override it to fix memory
+  /// attributes table as it is updated after pool dealloc.
+  ///
+  EFI_FREE_POOL               FreePool;
+  ///
   /// Original exit boot services function. We override it
   /// to ensure we always succeed exiting boot services.
   ///

Library/OcAppleBootCompatLib/ServiceOverrides.c

@@ -357,13 +357,16 @@ OcAllocatePages (
   return Status;
 }
 
+/**
+  UEFI Boot Services FreePages override.
+  Ensures synchronised memory attribute table.
+**/
 STATIC
 EFI_STATUS
 EFIAPI
-OcAllocatePool (
-  IN  EFI_MEMORY_TYPE              PoolType,
-  IN  UINTN                        Size,
-  OUT VOID                         **Buffer
+OcFreePages (
+  IN  EFI_PHYSICAL_ADDRESS         Memory,
+  IN  UINTN                        Pages
   )
 {
   EFI_STATUS              Status;
@@ -371,10 +374,9 @@ OcAllocatePool (
 
   BootCompat  = GetBootCompatContext ();
 
-  Status = BootCompat->ServicePtrs.AllocatePool (
-    PoolType,
-    Size,
-    Buffer
+  Status = BootCompat->ServicePtrs.FreePages (
+    Memory,
+    Pages
     );
 
   if (!EFI_ERROR (Status)) {
@@ -498,6 +500,64 @@ OcGetMemoryMap (
   return Status;
 }
 
+/**
+  UEFI Boot Services AllocatePool override.
+  Ensures synchronised memory attribute table.
+**/
+STATIC
+EFI_STATUS
+EFIAPI
+OcAllocatePool (
+  IN  EFI_MEMORY_TYPE              PoolType,
+  IN  UINTN                        Size,
+  OUT VOID                         **Buffer
+  )
+{
+  EFI_STATUS              Status;
+  BOOT_COMPAT_CONTEXT     *BootCompat;
+
+  BootCompat  = GetBootCompatContext ();
+
+  Status = BootCompat->ServicePtrs.AllocatePool (
+    PoolType,
+    Size,
+    Buffer
+    );
+
+  if (!EFI_ERROR (Status)) {
+    FixRuntimeAttributes (BootCompat);
+  }
+
+  return Status;
+}
+
+/**
+  UEFI Boot Services FreePool override.
+  Ensures synchronised memory attribute table.
+**/
+STATIC
+EFI_STATUS
+EFIAPI
+OcFreePool (
+  IN VOID                         *Buffer
+  )
+{
+  EFI_STATUS              Status;
+  BOOT_COMPAT_CONTEXT     *BootCompat;
+
+  BootCompat  = GetBootCompatContext ();
+
+  Status = BootCompat->ServicePtrs.FreePool (
+    Buffer
+    );
+
+  if (!EFI_ERROR (Status)) {
+    FixRuntimeAttributes (BootCompat);
+  }
+
+  return Status;
+}
+
 /**
   UEFI Boot Services StartImage override. Called to start an efi image.
   If this is boot.efi, then our overrides are enabled.
@@ -912,15 +972,19 @@ InstallServiceOverrides (
   OriginalTpl = gBS->RaiseTPL (TPL_HIGH_LEVEL);
 
   ServicePtrs->AllocatePages        = gBS->AllocatePages;
-  ServicePtrs->AllocatePool         = gBS->AllocatePool;
+  ServicePtrs->FreePages            = gBS->FreePages;
   ServicePtrs->GetMemoryMap         = gBS->GetMemoryMap;
+  ServicePtrs->AllocatePool         = gBS->AllocatePool;
+  ServicePtrs->FreePool             = gBS->FreePool;
   ServicePtrs->ExitBootServices     = gBS->ExitBootServices;
   ServicePtrs->StartImage           = gBS->StartImage;
   ServicePtrs->SetVirtualAddressMap = gRT->SetVirtualAddressMap;
 
   gBS->AllocatePages        = OcAllocatePages;
-  gBS->AllocatePool         = OcAllocatePool;
+  gBS->FreePages            = OcFreePages;
   gBS->GetMemoryMap         = OcGetMemoryMap;
+  gBS->AllocatePool         = OcAllocatePool;
+  gBS->FreePool             = OcFreePool;
   gBS->ExitBootServices     = OcExitBootServices;
   gBS->StartImage           = OcStartImage;
   gRT->SetVirtualAddressMap = OcSetVirtualAddressMap;