OcMemoryLib: Implement MAT deduplication and watch out on recursion

Application/MmapDump/MmapDump.c

@@ -84,9 +84,11 @@ UefiMain (
     OcSortMemoryMap (MemoryMapSize, MemoryMap, DescriptorSize);
     DEBUG ((DEBUG_INFO, "MMDD: Dumping the original memory map\n"));
     OcPrintMemoryMap (MemoryMapSize, MemoryMap, DescriptorSize);
-    OcUpdateAttributes (Address, EfiRuntimeServicesCode, EFI_MEMORY_RO, EFI_MEMORY_XP);
+    DEBUG ((DEBUG_INFO, "MMDD: Dumping patched attributes\n"));
+    OcRebuildAttributes (Address, NULL);
     Status = OcSplitMemoryMapByAttributes (OriginalSize, &MemoryMapSize, MemoryMap, DescriptorSize);
     if (!EFI_ERROR (Status)) {
+      OcPrintMemoryAttributesTable ();
       DEBUG ((DEBUG_INFO, "MMDD: Dumping patched memory map\n"));
       OcPrintMemoryMap (MemoryMapSize, MemoryMap, DescriptorSize);
     } else {

Include/Library/OcMemoryLib.h

@@ -103,12 +103,12 @@ OcGetCurrentMemoryMap (
   @param[out]    DescriptorSize     Resulting memory map descriptor size in bytes.
   @param[out]    DescriptorVersion  Memory map descriptor version.
   @param[in]     GetMemoryMap       Custom GetMemoryMap implementation to use, optional.
-  @param[in,out] TopMemory          Base top address for AllocatePagesFromTop allocation, number of pages after return.
+  @param[in,out] TopMemory          Base top address for OcAllocatePagesFromTop allocation, number of pages after return.
 
   @retval EFI_SUCCESS on success.
 **/
 EFI_STATUS
-GetCurrentMemoryMapAlloc (
+OcGetCurrentMemoryMapAlloc (
      OUT UINTN                  *MemoryMapSize,
      OUT EFI_MEMORY_DESCRIPTOR  **MemoryMap,
      OUT UINTN                  *MapKey,
@@ -163,7 +163,7 @@ OcShrinkMemoryMap (
 **/
 EFI_STATUS
 OcDeduplicateDescriptors (
-  IN OUT UINTN                  *EntryCount,
+  IN OUT UINT32                 *EntryCount,
   IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,
   IN     UINTN                  DescriptorSize
   );
@@ -183,7 +183,6 @@ BOOLEAN
   IN UINTN                 Size
   );
 
-
 /**
   Filter memory map entries.
 
@@ -214,7 +213,7 @@ VOID
   @retval EFI_SUCCESS on successful allocation.
 **/
 EFI_STATUS
-AllocatePagesFromTop (
+OcAllocatePagesFromTop (
   IN     EFI_MEMORY_TYPE         MemoryType,
   IN     UINTN                   Pages,
   IN OUT EFI_PHYSICAL_ADDRESS    *Memory,
@@ -233,7 +232,7 @@ AllocatePagesFromTop (
   @retval Number of runtime pages.
 **/
 UINT64
-CountRuntimePages (
+OcCountRuntimePages (
   IN  UINTN                  MemoryMapSize,
   IN  EFI_MEMORY_DESCRIPTOR  *MemoryMap,
   IN  UINTN                  DescriptorSize,
@@ -248,7 +247,7 @@ CountRuntimePages (
   @retval Number of free pages.
 **/
 UINTN
-CountFreePages (
+OcCountFreePages (
   OUT UINTN                  *LowerMemory  OPTIONAL
   );
 
@@ -317,20 +316,16 @@ OcGetMemoryAttributes (
   Refresh memory attributes entry containing the specified address.
 
   @param[in]  Address         Address contained in the updated entry.
-  @param[in]  Type            Memory type to assign to the entry.
-  @param[in]  SetAttribues    Attributes to set.
-  @param[in]  DropAttributes  Attributes to remove.
+  @param[in]  GetMemoryMap    
 
   @retval EFI_SUCCESS on success.
   @retval EFI_NOT_FOUND no entry contains the specified address.
   @retval EFI_UNSUPPORTED memory attributes are not supported by the platform.
 **/
 EFI_STATUS
-OcUpdateAttributes (
+OcRebuildAttributes (
   IN EFI_PHYSICAL_ADDRESS  Address,
-  IN EFI_MEMORY_TYPE       Type,
-  IN UINT64                SetAttributes,
-  IN UINT64                DropAttributes
+  IN EFI_GET_MEMORY_MAP    GetMemoryMap  OPTIONAL
   );
 
 /**
@@ -374,7 +369,7 @@ OcSplitMemoryMapByAttributes (
   @retval Current page table address.
 **/
 PAGE_MAP_AND_DIRECTORY_POINTER  *
-GetCurrentPageTable (
+OcGetCurrentPageTable (
   OUT UINTN                           *Flags  OPTIONAL
   );
 
@@ -388,7 +383,7 @@ GetCurrentPageTable (
   @retval EFI_SUCCESS on successful lookup.
 **/
 EFI_STATUS
-GetPhysicalAddress (
+OcGetPhysicalAddress (
   IN  PAGE_MAP_AND_DIRECTORY_POINTER   *PageTable  OPTIONAL,
   IN  EFI_VIRTUAL_ADDRESS              VirtualAddr,
   OUT EFI_PHYSICAL_ADDRESS             *PhysicalAddr

Library/OcAppleBootCompatLib/CustomSlide.c

@@ -261,7 +261,7 @@ ShouldUseCustomSlideOffset (
   }
 
   AllocatedMapPages = BASE_4GB;
-  Status = GetCurrentMemoryMapAlloc (
+  Status = OcGetCurrentMemoryMapAlloc (
     &MemoryMapSize,
     &MemoryMap,
     &MapKey,
@@ -285,7 +285,7 @@ ShouldUseCustomSlideOffset (
                                 CpuGeneration == OcCpuGenerationIvyBridge;
 
   SlideSupport->EstimatedKernelArea = (UINTN) EFI_PAGES_TO_SIZE (
-    CountRuntimePages (MemoryMapSize, MemoryMap, DescriptorSize, NULL)
+    OcCountRuntimePages (MemoryMapSize, MemoryMap, DescriptorSize, NULL)
     ) + ESTIMATED_KERNEL_SIZE;
 
   //

Library/OcAppleBootCompatLib/KernelSupport.c

@@ -181,7 +181,7 @@ PerformRtMemoryVirtualMapping (
   //
   // Get current VM page table.
   //
-  PageTable = GetCurrentPageTable (NULL);
+  PageTable = OcGetCurrentPageTable (NULL);
 
   for (Index = 0; Index < NumEntries; ++Index) {
     //
@@ -520,7 +520,7 @@ AppleMapPrepareBooterState (
       //
       BootCompat->KernelState.SysTableRtArea     = BASE_4GB;
       BootCompat->KernelState.SysTableRtAreaSize = gST->Hdr.HeaderSize;
-      Status = AllocatePagesFromTop (
+      Status = OcAllocatePagesFromTop (
         EfiRuntimeServicesData,
         EFI_SIZE_TO_PAGES (gST->Hdr.HeaderSize),
         &BootCompat->KernelState.SysTableRtArea,

Library/OcAppleBootCompatLib/OcAppleBootCompatLib.c

@@ -136,7 +136,7 @@ OcAbcInitialize (
     ));
 
   DEBUG_CODE_BEGIN ();
-  TotalMemory = CountFreePages (&LowMemory);
+  TotalMemory = OcCountFreePages (&LowMemory);
   DEBUG ((
     DEBUG_INFO,
     "OCABC: Firmware has %Lu free pages (%Lu in lower 4 GB)\n",

Library/OcAppleBootCompatLib/ServiceOverrides.c

@@ -55,6 +55,10 @@ FixRuntimeAttributes (
   }
 
   if (BootCompat->Settings.SyncRuntimePermissions && BootCompat->ServiceState.FwRuntime != NULL) {
+    //
+    // Be very careful of recursion here, who knows what the firmware can call.
+    //
+    BootCompat->Settings.SyncRuntimePermissions = FALSE;
     //
     // Some firmwares do not update MAT after loading runtime drivers after EndOfDxe.
     // Since the memory used to allocate runtime driver resides in BINs, MAT has whatever
@@ -65,8 +69,13 @@ FixRuntimeAttributes (
     Status = BootCompat->ServiceState.FwRuntime->GetExecArea (&Address, &Pages);
 
     if (!EFI_ERROR (Status)) {
-      OcUpdateAttributes (Address, EfiRuntimeServicesCode, EFI_MEMORY_RO, EFI_MEMORY_XP);
+      OcRebuildAttributes (Address, BootCompat->ServicePtrs.GetMemoryMap);
     }
+
+    //
+    // Permit syncing runtime permissions again.
+    //
+    BootCompat->Settings.SyncRuntimePermissions = TRUE;
   }
 }
 
@@ -115,7 +124,7 @@ ForceExitBootServices (
     // It is technically forbidden to allocate pool memory here, but we should not hit this code
     // in the first place, and for older firmwares, where it was necessary (?), it worked just fine.
     //
-    Status = GetCurrentMemoryMapAlloc (
+    Status = OcGetCurrentMemoryMapAlloc (
       &MemoryMapSize,
       &MemoryMap,
       &MapKey,

Library/OcMemoryLib/MemoryAlloc.c

+/** @file
+  Copyright (C) 2019, vit9696. All rights reserved.
+
+  All rights reserved.
+
+  This program and the accompanying materials
+  are licensed and made available under the terms and conditions of the BSD License
+  which accompanies this distribution.  The full text of the license may be found at
+  http://opensource.org/licenses/bsd-license.php
+
+  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+**/
+
+#include <Uefi.h>
+
+#include <Guid/MemoryAttributesTable.h>
+#include <Library/BaseMemoryLib.h>
+#include <Library/DebugLib.h>
+#include <Library/MemoryAllocationLib.h>
+#include <Library/OcGuardLib.h>
+#include <Library/OcMemoryLib.h>
+#include <Library/UefiBootServicesTableLib.h>
+#include <Library/UefiLib.h>
+
+EFI_STATUS
+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
+  )
+{
+  EFI_STATUS              Status;
+  UINTN                   MemoryMapSize;
+  EFI_MEMORY_DESCRIPTOR   *MemoryMap;
+  UINTN                   MapKey;
+  UINTN                   DescriptorSize;
+  UINT32                  DescriptorVersion;
+  EFI_MEMORY_DESCRIPTOR   *MemoryMapEnd;
+  EFI_MEMORY_DESCRIPTOR   *Desc;
+
+  Status = OcGetCurrentMemoryMapAlloc (
+    &MemoryMapSize,
+    &MemoryMap,
+    &MapKey,
+    &DescriptorSize,
+    &DescriptorVersion,
+    GetMemoryMap,
+    NULL
+    );
+
+  if (EFI_ERROR (Status)) {
+    return Status;
+  }
+
+  Status = EFI_NOT_FOUND;
+
+  MemoryMapEnd = NEXT_MEMORY_DESCRIPTOR (MemoryMap, MemoryMapSize);
+  Desc = PREV_MEMORY_DESCRIPTOR (MemoryMapEnd, DescriptorSize);
+
+  for ( ; Desc >= MemoryMap; Desc = PREV_MEMORY_DESCRIPTOR (Desc, DescriptorSize)) {
+    //
+    // We are looking for some free memory descriptor that contains enough
+    // space below the specified memory.
+    //
+    if (Desc->Type == EfiConventionalMemory && Pages <= Desc->NumberOfPages &&
+      Desc->PhysicalStart + EFI_PAGES_TO_SIZE (Pages) <= *Memory) {
+
+      //
+      // Free block found
+      //
+      if (Desc->PhysicalStart + EFI_PAGES_TO_SIZE (Desc->NumberOfPages) <= *Memory) {
+        //
+        // The whole block is under Memory: allocate from the top of the block
+        //
+        *Memory = Desc->PhysicalStart + EFI_PAGES_TO_SIZE (Desc->NumberOfPages - Pages);
+      } else {
+        //
+        // The block contains enough pages under Memory, but spans above it - allocate below Memory
+        //
+        *Memory = *Memory - EFI_PAGES_TO_SIZE (Pages);
+      }
+
+      //
+      // Ensure that the found block does not overlap with the restricted area.
+      //
+      if (CheckRange != NULL && CheckRange (*Memory, EFI_PAGES_TO_SIZE (Pages))) {
+        continue;
+      }
+
+      Status = gBS->AllocatePages (
+        AllocateAddress,
+        MemoryType,
+        Pages,
+        Memory
+        );
+
+      break;
+    }
+  }
+
+  FreePool (MemoryMap);
+
+  return Status;
+}
+
+UINT64
+OcCountRuntimePages (
+  IN  UINTN                  MemoryMapSize,
+  IN  EFI_MEMORY_DESCRIPTOR  *MemoryMap,
+  IN  UINTN                  DescriptorSize,
+  OUT UINTN                  *DescriptorCount OPTIONAL
+  )
+{
+  UINTN                  DescNum;
+  UINT64                 PageNum;
+  UINTN                  NumEntries;
+  UINTN                  Index;
+  EFI_MEMORY_DESCRIPTOR  *Desc;
+
+  DescNum    = 0;
+  PageNum    = 0;
+  NumEntries = MemoryMapSize / DescriptorSize;
+  Desc       = MemoryMap;
+
+  for (Index = 0; Index < NumEntries; ++Index) {
+    if (Desc->Type != EfiReservedMemoryType
+      && (Desc->Attribute & EFI_MEMORY_RUNTIME) != 0) {
+      ++DescNum;
+      PageNum += Desc->NumberOfPages;
+    }
+
+    Desc = NEXT_MEMORY_DESCRIPTOR (Desc, DescriptorSize);
+  }
+
+  if (DescriptorCount != NULL) {
+    *DescriptorCount = DescNum;
+  }
+
+  return PageNum;
+}
+
+UINTN
+OcCountFreePages (
+  OUT UINTN                  *LowerMemory  OPTIONAL
+  )
+{
+  UINTN                        MemoryMapSize;
+  UINTN                        DescriptorSize;
+  EFI_MEMORY_DESCRIPTOR        *MemoryMap;
+  EFI_MEMORY_DESCRIPTOR        *EntryWalker;
+  UINTN                        FreePages;
+
+  FreePages = 0;
+  if (LowerMemory != NULL) {
+    *LowerMemory = 0;
+  }
+
+  MemoryMap = OcGetCurrentMemoryMap (&MemoryMapSize, &DescriptorSize, NULL, NULL, NULL, FALSE);
+  if (MemoryMap == NULL) {
+    return 0;
+  }
+
+  for (
+    EntryWalker = MemoryMap;
+    (UINT8 *) EntryWalker < ((UINT8 *) MemoryMap + MemoryMapSize);
+    EntryWalker = NEXT_MEMORY_DESCRIPTOR (EntryWalker, DescriptorSize)) {
+
+    if (EntryWalker->Type != EfiConventionalMemory) {
+      continue;
+    }
+
+    //
+    // This cannot overflow even on 32-bit systems unless they have > 16 TB of RAM,
+    // just assert to ensure that we have valid MemoryMap.
+    //
+    ASSERT (EntryWalker->NumberOfPages <= MAX_UINTN);
+    ASSERT (MAX_UINTN - EntryWalker->NumberOfPages >= FreePages);
+    FreePages += (UINTN) EntryWalker->NumberOfPages;
+
+    if (LowerMemory == NULL || EntryWalker->PhysicalStart >= BASE_4GB) {
+      continue;
+    }
+
+    if (EntryWalker->PhysicalStart + EFI_PAGES_TO_SIZE (EntryWalker->NumberOfPages) > BASE_4GB) {
+      *LowerMemory += (UINTN) EFI_SIZE_TO_PAGES (BASE_4GB - EntryWalker->PhysicalStart);
+    } else {
+      *LowerMemory += (UINTN) EntryWalker->NumberOfPages;
+    }
+  }
+
+  FreePool (MemoryMap);
+
+  return FreePages;
+}

Library/OcMemoryLib/MemoryAttributes.c

+/** @file
+  Copyright (C) 2019, vit9696. All rights reserved.
+
+  All rights reserved.
+
+  This program and the accompanying materials
+  are licensed and made available under the terms and conditions of the BSD License
+  which accompanies this distribution.  The full text of the license may be found at
+  http://opensource.org/licenses/bsd-license.php
+
+  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+**/
+
+#include <Uefi.h>
+
+#include <Guid/MemoryAttributesTable.h>
+#include <Library/BaseMemoryLib.h>
+#include <Library/DebugLib.h>
+#include <Library/MemoryAllocationLib.h>
+#include <Library/OcGuardLib.h>
+#include <Library/OcMemoryLib.h>
+#include <Library/UefiBootServicesTableLib.h>
+#include <Library/UefiLib.h>
+
+/**
+  Determine actual memory type from the attribute.
+
+  @param[in]  MemoryAttribute  Attribute to inspect.
+**/
+STATIC
+UINT32
+OcRealMemoryType (
+  IN EFI_MEMORY_DESCRIPTOR  *MemoryAttribte
+  )
+{
+  ASSERT (MemoryAttribte->Type == EfiRuntimeServicesCode
+    || MemoryAttribte->Type == EfiRuntimeServicesData);
+
+  //
+  // Use code for write-protected areas.
+  //
+  if ((MemoryAttribte->Attribute & EFI_MEMORY_RO) != 0) {
+    return EfiRuntimeServicesCode;
+  }
+
+  //
+  // Use data for executable-protected areas.
+  //
+  if ((MemoryAttribte->Attribute & EFI_MEMORY_XP) != 0) {
+    return EfiRuntimeServicesData;
+  }
+
+  //
+  // Use whatever is set.
+  //
+  return MemoryAttribte->Type;
+}
+
+/**
+  Split memory map descriptor by attribute.
+
+  @param[in,out] RetMemoryMapEntry    Pointer to descriptor in the memory map, updated to next proccessed.
+  @param[in,out] CurrentEntryIndex    Current index of the descriptor in the memory map, updated on increase.
+  @param[in,out] CurrentEntryCount    Number of descriptors in the memory map, updated on increase.
+  @param[in]     TotalEntryCount      Max number of descriptors in the memory map.
+  @param[in]     MemoryAttribute      Memory attribute used for splitting.
+  @param[in]     DescriptorSize       Memory map descriptor size.
+
+  @retval EFI_SUCCESS on success.
+  @retval EFI_OUT_OF_RESOURCES when there are not enough free descriptor slots.
+**/
+STATIC
+EFI_STATUS
+OcSplitMemoryEntryByAttribute (
+  IN OUT EFI_MEMORY_DESCRIPTOR  **RetMemoryMapEntry,
+  IN OUT UINTN                  *CurrentEntryIndex,
+  IN OUT UINTN                  *CurrentEntryCount,
+  IN     UINTN                  TotalEntryCount,
+  IN     EFI_MEMORY_DESCRIPTOR  *MemoryAttribute,
+  IN     UINTN                  DescriptorSize
+
+  )
+{
+  EFI_MEMORY_DESCRIPTOR  *MemoryMapEntry;
+  EFI_MEMORY_DESCRIPTOR  *NewMemoryMapEntry;
+  UINTN                  DiffPages;
+
+  MemoryMapEntry = *RetMemoryMapEntry;
+
+  //
+  // Memory attribute starts after our descriptor.
+  // Shorten the existing descriptor and insert the new one after it.
+  // [DESC1] -> [DESC1][DESC2]
+  //
+  if (MemoryAttribute->PhysicalStart > MemoryMapEntry->PhysicalStart) {
+    if (*CurrentEntryCount == TotalEntryCount) {
+      return EFI_OUT_OF_RESOURCES;
+    }
+
+    NewMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);
+    DiffPages         = (UINTN) EFI_SIZE_TO_PAGES (MemoryAttribute->PhysicalStart - MemoryMapEntry->PhysicalStart);
+    CopyMem (
+      NewMemoryMapEntry,
+      MemoryMapEntry,
+      DescriptorSize * (*CurrentEntryCount - *CurrentEntryIndex)
+      );
+    MemoryMapEntry->NumberOfPages     = DiffPages;
+    NewMemoryMapEntry->PhysicalStart  = MemoryAttribute->PhysicalStart;
+    NewMemoryMapEntry->NumberOfPages -= DiffPages;
+
+    MemoryMapEntry = NewMemoryMapEntry;
+
+    //
+    // Current processed entry is now the one we inserted.
+    //
+    ++(*CurrentEntryIndex);
+    ++(*CurrentEntryCount);
+  }
+
+  ASSERT (MemoryAttribute->PhysicalStart == MemoryMapEntry->PhysicalStart);
+
+  //
+  // Memory attribute matches our descriptor.
+  // Simply update its protection.
+  // [DESC1] -> [DESC1*]
+  //
+  if (MemoryMapEntry->NumberOfPages == MemoryAttribute->NumberOfPages) {
+    MemoryMapEntry->Type = OcRealMemoryType (MemoryAttribute);
+    *RetMemoryMapEntry = MemoryMapEntry;
+    return EFI_SUCCESS;
+  }
+
+  //
+  // Memory attribute is shorter than our descriptor.
+  // Shorten current descriptor, update its type, and inseret the new one after it.
+  // [DESC1] -> [DESC1*][DESC2]
+  //
+  if (*CurrentEntryCount == TotalEntryCount) {
+    return EFI_OUT_OF_RESOURCES;
+  }
+
+  NewMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);
+  CopyMem (
+    NewMemoryMapEntry,
+    MemoryMapEntry,
+    DescriptorSize * (*CurrentEntryCount - *CurrentEntryIndex)
+    );
+  MemoryMapEntry->Type              = OcRealMemoryType (MemoryAttribute);
+  MemoryMapEntry->NumberOfPages     = MemoryAttribute->NumberOfPages;
+  NewMemoryMapEntry->PhysicalStart += EFI_PAGES_TO_SIZE (MemoryAttribute->NumberOfPages);
+  NewMemoryMapEntry->NumberOfPages -= MemoryAttribute->NumberOfPages;
+
+  //
+  // Current processed entry is now the one we need to process.
+  //
+  ++(*CurrentEntryIndex);
+  ++(*CurrentEntryCount);
+
+  *RetMemoryMapEntry = NewMemoryMapEntry;
+
+  return EFI_SUCCESS;
+}
+
+EFI_MEMORY_ATTRIBUTES_TABLE *
+OcGetMemoryAttributes (
+  OUT EFI_MEMORY_DESCRIPTOR  **MemoryAttributesEntry  OPTIONAL
+  )
+{
+  EFI_MEMORY_ATTRIBUTES_TABLE  *MemoryAttributesTable; 
+  UINTN                        Index;
+
+  for (Index = 0; Index < gST->NumberOfTableEntries; ++Index) {
+    if (CompareGuid (&gST->ConfigurationTable[Index].VendorGuid, &gEfiMemoryAttributesTableGuid)) {
+      MemoryAttributesTable = gST->ConfigurationTable[Index].VendorTable;
+      if (MemoryAttributesEntry != NULL) {
+        *MemoryAttributesEntry = (EFI_MEMORY_DESCRIPTOR *) (MemoryAttributesTable + 1);
+      }
+      return MemoryAttributesTable;
+    }
+  }
+
+  return NULL;
+}
+
+EFI_STATUS
+OcRebuildAttributes (
+  IN EFI_PHYSICAL_ADDRESS  Address,
+  IN EFI_GET_MEMORY_MAP    GetMemoryMap  OPTIONAL
+  )
+{
+  EFI_STATUS                         Status;
+  EFI_MEMORY_ATTRIBUTES_TABLE        *MemoryAttributesTable;
+  EFI_MEMORY_DESCRIPTOR              *MemoryAttributesEntry;
+  UINTN                              MaxDescriptors;
+
+  MemoryAttributesTable = OcGetMemoryAttributes (&MemoryAttributesEntry);
+  if (MemoryAttributesTable == NULL) {
+    return EFI_UNSUPPORTED;
+  }
+
+  //
+  // Some boards create entry duplicates and lose all non-PE entries
+  // after loading runtime drivers after EndOfDxe.
+  // REF: https://github.com/acidanthera/bugtracker/issues/491#issuecomment-609014334
+  //
+  MaxDescriptors = MemoryAttributesTable->NumberOfEntries;
+  Status = OcDeduplicateDescriptors (
+    &MemoryAttributesTable->NumberOfEntries,
+    MemoryAttributesEntry,
+    MemoryAttributesTable->DescriptorSize
+    );
+  if (!EFI_ERROR (Status)) {
+    //
+    // Assume effected and add missing entries.
+    //
+    if (GetMemoryMap == NULL) {
+      GetMemoryMap = gBS->GetMemoryMap;
+    }
+
+    //
+    // TODO: Implement
+    //
+  }
+
+  Status = OcUpdateDescriptors (
+    MemoryAttributesTable->NumberOfEntries * MemoryAttributesTable->DescriptorSize,
+    MemoryAttributesEntry,
+    MemoryAttributesTable->DescriptorSize,
+    Address,
+    EfiRuntimeServicesCode,
+    EFI_MEMORY_RO,
+    EFI_MEMORY_XP
+    );
+
+  return Status;
+}
+
+UINTN
+OcCountSplitDescritptors (
+  VOID
+  )
+{
+  UINTN                             Index;
+  UINTN                             DescriptorCount;
+  CONST EFI_MEMORY_ATTRIBUTES_TABLE *MemoryAttributesTable;
+  EFI_MEMORY_DESCRIPTOR             *MemoryAttributesEntry;
+
+  MemoryAttributesTable = OcGetMemoryAttributes (&MemoryAttributesEntry);
+  if (MemoryAttributesTable == NULL) {
+    return 0;
+  }
+
+  DescriptorCount = 0;
+  for (Index = 0; Index < MemoryAttributesTable->NumberOfEntries; ++Index) {
+    if (MemoryAttributesEntry->Type == EfiRuntimeServicesCode
+      || MemoryAttributesEntry->Type == EfiRuntimeServicesData) {
+      ++DescriptorCount;
+    }
+
+    MemoryAttributesEntry = NEXT_MEMORY_DESCRIPTOR (
+      MemoryAttributesEntry,
+      MemoryAttributesTable->DescriptorSize
+      );
+  }
+
+  return DescriptorCount;
+}
+
+EFI_STATUS
+OcSplitMemoryMapByAttributes (
+  IN     UINTN                  MaxMemoryMapSize,
+  IN OUT UINTN                  *MemoryMapSize,
+  IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,
+  IN     UINTN                  DescriptorSize
+  )
+{
+  EFI_STATUS                         Status;
+  CONST EFI_MEMORY_ATTRIBUTES_TABLE  *MemoryAttributesTable;
+  EFI_MEMORY_DESCRIPTOR              *MemoryAttributesEntry;
+  EFI_MEMORY_DESCRIPTOR              *MemoryMapEntry;
+  EFI_MEMORY_DESCRIPTOR              *LastAttributeEntry;
+  UINTN                              LastAttributeIndex;
+  UINTN                              Index;
+  UINTN                              Index2;
+  UINTN                              CurrentEntryCount;
+  UINTN                              TotalEntryCount;
+  UINTN                              AttributeCount;
+  BOOLEAN                            CanSplit;
+  BOOLEAN                            InDescAttrs;
+
+  ASSERT (MaxMemoryMapSize >= *MemoryMapSize);
+
+  MemoryAttributesTable = OcGetMemoryAttributes (&MemoryAttributesEntry);
+  if (MemoryAttributesTable == NULL) {
+    return EFI_UNSUPPORTED;
+  }
+
+  LastAttributeEntry = MemoryAttributesEntry;
+  LastAttributeIndex = 0;
+  MemoryMapEntry     = MemoryMap;
+  CurrentEntryCount  = *MemoryMapSize / DescriptorSize;
+  TotalEntryCount    = MaxMemoryMapSize / DescriptorSize;
+  AttributeCount     = MemoryAttributesTable->NumberOfEntries;
+
+  //
+  // We assume that the memory map and attribute table are sorted.
+  //
+  Index = 0;
+  while (Index < CurrentEntryCount) {
+    //
+    // Split entry by as many attributes as possible.
+    //
+    CanSplit = TRUE;
+    while ((MemoryMapEntry->Type == EfiRuntimeServicesCode
+      || MemoryMapEntry->Type == EfiRuntimeServicesData) && CanSplit) {
+      //
+      // Find corresponding memory attribute.
+      //
+      InDescAttrs = FALSE;
+      MemoryAttributesEntry = LastAttributeEntry;
+      for (Index2 = LastAttributeIndex; Index2 < AttributeCount; ++Index2) {
+        if (MemoryAttributesEntry->Type == EfiRuntimeServicesCode
+          || MemoryAttributesEntry->Type == EfiRuntimeServicesData) {
+          //
+          // UEFI spec says attribute entries are fully within memory map entries.
+          // Find first one of a different type.
+          //
+          if (AREA_WITHIN_DESCRIPTOR (
+            MemoryMapEntry,
+            MemoryAttributesEntry->PhysicalStart,
+            EFI_PAGES_TO_SIZE (MemoryAttributesEntry->NumberOfPages))) {
+            //
+            // We are within descriptor attribute sequence.
+            //
+            InDescAttrs = TRUE;
+            //
+            // No need to process the attribute of the same type.
+            //
+            if (OcRealMemoryType (MemoryAttributesEntry) != MemoryMapEntry->Type) {
+              //
+              // Start with the next attribute on the second iteration.
+              //
+              LastAttributeEntry = NEXT_MEMORY_DESCRIPTOR (
+                MemoryAttributesEntry,
+                MemoryAttributesTable->DescriptorSize
+                );
+              LastAttributeIndex = Index2 + 1;
+              break;
+            }
+          } else if (InDescAttrs) {
+            //
+            // Reached the end of descriptor attribute sequence, abort.
+            //
+            InDescAttrs = FALSE;
+            break;
+          }
+        }
+
+        MemoryAttributesEntry = NEXT_MEMORY_DESCRIPTOR (
+          MemoryAttributesEntry,
+          MemoryAttributesTable->DescriptorSize
+          );
+      }
+
+      if (Index2 < AttributeCount && InDescAttrs) {
+        //
+        // Split current memory map entry.
+        //
+        Status = OcSplitMemoryEntryByAttribute (
+          &MemoryMapEntry,
+          &Index,
+          &CurrentEntryCount,
+          TotalEntryCount,
+          MemoryAttributesEntry,
+          DescriptorSize
+          );
+        if (EFI_ERROR (Status)) {
+          *MemoryMapSize = CurrentEntryCount * DescriptorSize;
+          return Status;
+        }
+        continue;
+      } else {
+        //
+        // Did not find a suitable attribute or processed all the attributes.
+        //
+        CanSplit = FALSE;
+      }
+    }
+
+    MemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (
+      MemoryMapEntry,
+      DescriptorSize
+      );
+    ++Index;
+  }
+
+  *MemoryMapSize = CurrentEntryCount * DescriptorSize;
+  return EFI_SUCCESS;
+}

Library/OcMemoryLib/MemoryDebug.c

+/** @file
+  Copyright (C) 2019, vit9696. All rights reserved.
+
+  All rights reserved.
+
+  This program and the accompanying materials
+  are licensed and made available under the terms and conditions of the BSD License
+  which accompanies this distribution.  The full text of the license may be found at
+  http://opensource.org/licenses/bsd-license.php
+
+  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+**/
+
+#include <Uefi.h>
+
+#include <Guid/MemoryAttributesTable.h>
+#include <Library/BaseMemoryLib.h>
+#include <Library/DebugLib.h>
+#include <Library/MemoryAllocationLib.h>
+#include <Library/OcGuardLib.h>
+#include <Library/OcMemoryLib.h>
+#include <Library/UefiBootServicesTableLib.h>
+#include <Library/UefiLib.h>
+
+STATIC CONST CHAR8 *mEfiMemoryTypeDesc[EfiMaxMemoryType] = {
+  "Reserved ",
+  "LDR Code ",
+  "LDR Data ",
+  "BS Code  ",
+  "BS Data  ",
+  "RT Code  ",
+  "RT Data  ",
+  "Available",
+  "Unusable ",
+  "ACPI RECL",
+  "ACPI NVS ",
+  "MemMapIO ",
+  "MemPortIO",
+  "PAL Code ",
+  "Persist  "
+};
+
+STATIC
+VOID
+OcPrintMemoryDescritptor (
+  IN EFI_MEMORY_DESCRIPTOR  *Desc
+  )
+{
+  CONST CHAR8  *Type;
+  CONST CHAR8  *SizeType;
+  UINT64       SizeValue;
+
+  if (Desc->Type < ARRAY_SIZE (mEfiMemoryTypeDesc)) {
+    Type = mEfiMemoryTypeDesc[Desc->Type];
+  } else {
+    Type = "Invalid  ";
+  }
+
+  SizeValue = EFI_PAGES_TO_SIZE (Desc->NumberOfPages);
+  if (SizeValue >= BASE_1MB) {
+    SizeValue /= BASE_1MB;
+    SizeType   = "MB";
+  } else {
+    SizeValue /= BASE_1KB;
+    SizeType   = "KB";
+  }
+
+  DEBUG ((
+    DEBUG_INFO,
+    "OCMM: %a [%a|%a|%a|%a|%a|%a|%a|%a|%a|%a|%a|%a|%a|%a] 0x%016LX-0x%016LX -> 0x%016X (%Lu %a)\n",
+    Type,
+    (Desc->Attribute & EFI_MEMORY_RUNTIME)        != 0 ? "RUN" : "   ",
+    (Desc->Attribute & EFI_MEMORY_CPU_CRYPTO)     != 0 ? "CRY" : "   ",
+    (Desc->Attribute & EFI_MEMORY_SP)             != 0 ? "SP"  : "  ",
+    (Desc->Attribute & EFI_MEMORY_RO)             != 0 ? "RO"  : "  ",
+    (Desc->Attribute & EFI_MEMORY_MORE_RELIABLE)  != 0 ? "MR"  : "  ",
+    (Desc->Attribute & EFI_MEMORY_NV)             != 0 ? "NV"  : "  ",
+    (Desc->Attribute & EFI_MEMORY_XP)             != 0 ? "XP"  : "  ",
+    (Desc->Attribute & EFI_MEMORY_RP)             != 0 ? "RP"  : "  ",
+    (Desc->Attribute & EFI_MEMORY_WP)             != 0 ? "WP"  : "  ",
+    (Desc->Attribute & EFI_MEMORY_UCE)            != 0 ? "UCE" : "   ",
+    (Desc->Attribute & EFI_MEMORY_WB)             != 0 ? "WB"  : "  ",
+    (Desc->Attribute & EFI_MEMORY_WT)             != 0 ? "WT"  : "  ",
+    (Desc->Attribute & EFI_MEMORY_WC)             != 0 ? "WC"  : "  ",
+    (Desc->Attribute & EFI_MEMORY_UC)             != 0 ? "UC"  : "  ",
+    Desc->PhysicalStart,
+    Desc->PhysicalStart + (EFI_PAGES_TO_SIZE (Desc->NumberOfPages) - 1),
+    Desc->VirtualStart,
+    SizeValue,
+    SizeType
+    ));
+}
+
+VOID
+OcPrintMemoryAttributesTable (
+  VOID
+  )
+{
+  UINTN                             Index;
+  UINTN                             RealSize;
+  CONST EFI_MEMORY_ATTRIBUTES_TABLE *MemoryAttributesTable;
+  EFI_MEMORY_DESCRIPTOR             *MemoryAttributesEntry;
+
+  MemoryAttributesTable = OcGetMemoryAttributes (NULL);
+  if (MemoryAttributesTable == NULL) {
+    DEBUG ((DEBUG_INFO, "OCMM: MemoryAttributesTable is not present!\n"));
+    return;
+  }
+
+  //
+  // Printing may reallocate, so we create a copy of the memory attributes.
+  //
+  STATIC UINT8 mMemoryAttributesTable[EFI_PAGE_SIZE*2];
+  RealSize = (UINTN) (sizeof (EFI_MEMORY_ATTRIBUTES_TABLE)
+    + MemoryAttributesTable->NumberOfEntries * MemoryAttributesTable->DescriptorSize);
+
+  if (RealSize > sizeof(mMemoryAttributesTable)) {
+    DEBUG ((DEBUG_INFO, "OCMM: MemoryAttributesTable has too large size %u!\n", (UINT32) RealSize));
+    return;
+  }
+
+  CopyMem (mMemoryAttributesTable, MemoryAttributesTable, RealSize);
+
+  MemoryAttributesTable = (EFI_MEMORY_ATTRIBUTES_TABLE *) mMemoryAttributesTable;
+  MemoryAttributesEntry = (EFI_MEMORY_DESCRIPTOR *) (MemoryAttributesTable + 1);
+
+  DEBUG ((DEBUG_INFO, "OCMM: MemoryAttributesTable:\n"));
+  DEBUG ((DEBUG_INFO, "OCMM:   Version              - 0x%08x\n", MemoryAttributesTable->Version));
+  DEBUG ((DEBUG_INFO, "OCMM:   NumberOfEntries      - 0x%08x\n", MemoryAttributesTable->NumberOfEntries));
+  DEBUG ((DEBUG_INFO, "OCMM:   DescriptorSize       - 0x%08x\n", MemoryAttributesTable->DescriptorSize));
+
+  for (Index = 0; Index < MemoryAttributesTable->NumberOfEntries; ++Index) {
+    OcPrintMemoryDescritptor (MemoryAttributesEntry);
+    MemoryAttributesEntry = NEXT_MEMORY_DESCRIPTOR (
+      MemoryAttributesEntry,
+      MemoryAttributesTable->DescriptorSize
+      );
+  }
+}
+
+VOID
+OcPrintMemoryMap (
+  IN UINTN                  MemoryMapSize,
+  IN EFI_MEMORY_DESCRIPTOR  *MemoryMap,
+  IN UINTN                  DescriptorSize
+  )
+{
+  UINTN     Index;
+  UINT32    NumberOfEntries;
+
+  NumberOfEntries = (UINT32) (MemoryMapSize / DescriptorSize);
+
+  DEBUG ((DEBUG_INFO, "OCMM: MemoryMap:\n"));
+  DEBUG ((DEBUG_INFO, "OCMM:   Size                 - 0x%08x\n", MemoryMapSize));
+  DEBUG ((DEBUG_INFO, "OCMM:   NumberOfEntries      - 0x%08x\n", NumberOfEntries));
+  DEBUG ((DEBUG_INFO, "OCMM:   DescriptorSize       - 0x%08x\n", DescriptorSize));
+
+  for (Index = 0; Index < NumberOfEntries; ++Index) {
+    OcPrintMemoryDescritptor (MemoryMap);
+    MemoryMap = NEXT_MEMORY_DESCRIPTOR (
+      MemoryMap,
+      DescriptorSize
+      );
+  }
+}

Library/OcMemoryLib/OcMemoryLib.inf

@@ -45,6 +45,9 @@
   OcStringLib
 
 [Sources]
+  MemoryAlloc.c
+  MemoryAttributes.c
+  MemoryDebug.c
   MemoryMap.c
   LegacyRegionLock.c
   LegacyRegionUnLock.c

Library/OcMemoryLib/VirtualMemory.c

@@ -21,7 +21,7 @@
 #include <Library/OcMemoryLib.h>
 
 PAGE_MAP_AND_DIRECTORY_POINTER  *
-GetCurrentPageTable (
+OcGetCurrentPageTable (
   OUT UINTN                           *Flags  OPTIONAL
   )
 {
@@ -40,7 +40,7 @@ GetCurrentPageTable (
 }
 
 EFI_STATUS
-GetPhysicalAddress (
+OcGetPhysicalAddress (
   IN  PAGE_MAP_AND_DIRECTORY_POINTER   *PageTable  OPTIONAL,
   IN  EFI_VIRTUAL_ADDRESS              VirtualAddr,
   OUT EFI_PHYSICAL_ADDRESS             *PhysicalAddr
@@ -58,7 +58,7 @@ GetPhysicalAddress (
   PAGE_TABLE_1G_ENTRY             *PTE1G;
 
   if (PageTable == NULL) {
-    PageTable = GetCurrentPageTable (NULL);
+    PageTable = OcGetCurrentPageTable (NULL);
   }
 
   VA.Uint64 = (UINT64) VirtualAddr;
@@ -152,7 +152,7 @@ VmAllocateMemoryPool (
   EFI_PHYSICAL_ADDRESS    Addr;
 
   Addr = BASE_4GB;
-  Status = AllocatePagesFromTop (
+  Status = OcAllocatePagesFromTop (
     EfiBootServicesData,
     NumPages,
     &Addr,
@@ -209,7 +209,7 @@ VmMapVirtualPage (
   UINTN                           Index;
 
   if (PageTable == NULL) {
-    PageTable = GetCurrentPageTable (NULL);
+    PageTable = OcGetCurrentPageTable (NULL);
   }
 
   VA.Uint64 = (UINT64) VirtualAddr;
@@ -384,7 +384,7 @@ VmMapVirtualPages (
   EFI_STATUS  Status;
 
   if (PageTable == NULL) {
-    PageTable = GetCurrentPageTable (NULL);
+    PageTable = OcGetCurrentPageTable (NULL);
   }
 
   Status = EFI_SUCCESS;