diff --git a/Library/OcAppleBootCompatLib/BootCompatInternal.h b/Library/OcAppleBootCompatLib/BootCompatInternal.h index 646f5588..79ee780f 100644 --- a/Library/OcAppleBootCompatLib/BootCompatInternal.h +++ b/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. /// diff --git a/Library/OcAppleBootCompatLib/ServiceOverrides.c b/Library/OcAppleBootCompatLib/ServiceOverrides.c index 89c4206d..5bbd7a25 100644 --- a/Library/OcAppleBootCompatLib/ServiceOverrides.c +++ b/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;