OcMemoryLib: Add memory map splitting routines

Include/Library/OcMemoryLib.h

@@ -16,6 +16,7 @@
 #define OC_MEMORY_LIB_H
 
 #include <Uefi.h>
+#include <Guid/MemoryAttributesTable.h>
 #include <IndustryStandard/VirtualMemory.h>
 
 /**
@@ -245,6 +246,19 @@ OcUpdateDescriptors (
   IN UINT64                 DropAttributes
   );
 
+/**
+  Obtain memory attributes table.
+
+  @param[out] MemoryAttributesEntry  memory descriptor pointer, optional.
+
+  @retval pointer to memory attributes table.
+  @retval NULL if memory attributes table is unsupported.
+**/
+EFI_MEMORY_ATTRIBUTES_TABLE *
+OcGetMemoryAttributes (
+  OUT EFI_MEMORY_DESCRIPTOR  **MemoryAttributesEntry  OPTIONAL
+  );
+
 /**
   Refresh memory attributes entry containing the specified address.
 
@@ -275,6 +289,28 @@ OcCountSplitDescritptors (
   VOID
   );
 
+/**
+  Split memory map by memory attributes if available.
+
+  @param[in]     OriginalMemoryMapSize   Upper memory map size bound for growth.
+  @param[in,out] MemoryMapSize           Current memory map size, updated on return.
+  @param[in,out] MemoryMap               Memory map to split.
+  @param[in]     DescriptorSize          Memory map descriptor size.
+
+  Note, the function is guaranteed to return valid memory map, though not necessarily split.
+
+  @retval EFI_SUCCESS on success.
+  @retval EFI_UNSUPPORTED memory attributes are not supported by the platform.
+  @retval EFI_OUT_OF_RESOURCES new memory map did not fit.
+**/
+EFI_STATUS
+OcSplitMemoryMapByAttributes (
+  IN     UINTN                  OriginalMemoryMapSize,
+  IN OUT UINTN                  *MemoryMapSize,
+  IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,
+  IN     UINTN                  DescriptorSize
+  );
+
 /**
   Return pointer to PML4 table in PageTable and PWT and PCD flags in Flags.
 

Library/OcMemoryLib/MemoryMap.c

@@ -510,6 +510,27 @@ OcUpdateDescriptors (
   return EFI_NOT_FOUND;
 }
 
+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
 OcUpdateAttributes (
   IN EFI_PHYSICAL_ADDRESS  Address,
@@ -518,27 +539,23 @@ OcUpdateAttributes (
   IN UINT64                DropAttributes
   )
 {
-  UINTN                             Index;
   CONST EFI_MEMORY_ATTRIBUTES_TABLE *MemoryAttributesTable;
   EFI_MEMORY_DESCRIPTOR             *MemoryAttributesEntry;
 
-  for (Index = 0; Index < gST->NumberOfTableEntries; ++Index) {
-    if (CompareGuid (&gST->ConfigurationTable[Index].VendorGuid, &gEfiMemoryAttributesTableGuid)) {
-      MemoryAttributesTable = (CONST EFI_MEMORY_ATTRIBUTES_TABLE *) gST->ConfigurationTable[Index].VendorTable;
-      MemoryAttributesEntry = (EFI_MEMORY_DESCRIPTOR *) (MemoryAttributesTable + 1);
-      return OcUpdateDescriptors (
-        MemoryAttributesTable->NumberOfEntries * MemoryAttributesTable->DescriptorSize,
-        MemoryAttributesEntry,
-        MemoryAttributesTable->DescriptorSize,
-        Address,
-        Type,
-        SetAttributes,
-        DropAttributes
-        );
-    }
+  MemoryAttributesTable = OcGetMemoryAttributes (&MemoryAttributesEntry);
+  if (MemoryAttributesTable == NULL) {
+    return EFI_UNSUPPORTED;
   }
 
-  return EFI_UNSUPPORTED;
+  return OcUpdateDescriptors (
+    MemoryAttributesTable->NumberOfEntries * MemoryAttributesTable->DescriptorSize,
+    MemoryAttributesEntry,
+    MemoryAttributesTable->DescriptorSize,
+    Address,
+    Type,
+    SetAttributes,
+    DropAttributes
+    );
 }
 
 UINTN
@@ -551,16 +568,222 @@ OcCountSplitDescritptors (
   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;
+    }
 
-  for (Index = 0; Index < gST->NumberOfTableEntries; ++Index) {
-    if (CompareGuid (&gST->ConfigurationTable[Index].VendorGuid, &gEfiMemoryAttributesTableGuid)) {
-      MemoryAttributesTable = (CONST EFI_MEMORY_ATTRIBUTES_TABLE *) gST->ConfigurationTable[Index].VendorTable;
-      MemoryAttributesEntry = (EFI_MEMORY_DESCRIPTOR *) (MemoryAttributesTable + 1);
-      for (Index = 0; Index < MemoryAttributesTable->NumberOfEntries; ++Index) {
+    MemoryAttributesEntry = NEXT_MEMORY_DESCRIPTOR (
+      MemoryAttributesEntry,
+      MemoryAttributesTable->DescriptorSize
+      );
+  }
+
+  return DescriptorCount;
+}
+
+/**
+  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  *MemoryAttribte,
+  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 (MemoryAttribte->PhysicalStart > MemoryMapEntry->PhysicalStart) {
+    if (*CurrentEntryCount == TotalEntryCount) {
+      return EFI_OUT_OF_RESOURCES;
+    }
+
+    NewMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);
+    DiffPages         = EFI_SIZE_TO_PAGES (MemoryAttribte->PhysicalStart - MemoryMapEntry->PhysicalStart);
+    CopyMem (
+      NewMemoryMapEntry,
+      MemoryMapEntry,
+      DescriptorSize * (*CurrentEntryCount - *CurrentEntryIndex)
+      );
+    MemoryMapEntry->NumberOfPages     = DiffPages;
+    NewMemoryMapEntry->PhysicalStart  = MemoryAttribte->PhysicalStart;
+    NewMemoryMapEntry->NumberOfPages -= DiffPages;
+
+    MemoryMapEntry = NewMemoryMapEntry;
+
+    //
+    // Current processed entry is now the one we inserted.
+    //
+    ++(*CurrentEntryIndex);
+    ++(*CurrentEntryCount);
+  }
+
+  ASSERT (MemoryAttribte->PhysicalStart == MemoryMapEntry->PhysicalStart);
+
+  //
+  // Memory attribute matches our descriptor.
+  // Simply update its protection.
+  // [DESC1] -> [DESC1*]
+  //
+  if (MemoryMapEntry->NumberOfPages == MemoryAttribte->NumberOfPages) {
+    MemoryMapEntry->Type = MemoryAttribte->Type;
+    *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);
+  DiffPages         = MemoryMapEntry->NumberOfPages - MemoryAttribte->NumberOfPages;
+  CopyMem (
+    NewMemoryMapEntry,
+    MemoryMapEntry,
+    DescriptorSize * (*CurrentEntryCount - *CurrentEntryIndex)
+    );
+  MemoryMapEntry->Type              = MemoryAttribte->Type;
+  MemoryMapEntry->NumberOfPages     = DiffPages;
+  NewMemoryMapEntry->PhysicalStart += EFI_PAGES_TO_SIZE (DiffPages);
+  NewMemoryMapEntry->NumberOfPages -= DiffPages;
+
+  //
+  // Current processed entry is now the one we need to process.
+  //
+  ++(*CurrentEntryIndex);
+  ++(*CurrentEntryCount);
+
+  *RetMemoryMapEntry = NewMemoryMapEntry;
+
+  return EFI_SUCCESS;
+}
+
+EFI_STATUS
+OcSplitMemoryMapByAttributes (
+  IN     UINTN                  OriginalMemoryMapSize,
+  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 (OriginalMemoryMapSize >= *MemoryMapSize);
+
+  MemoryAttributesTable = OcGetMemoryAttributes (&MemoryAttributesEntry);
+  if (MemoryAttributesTable == NULL) {
+    return EFI_UNSUPPORTED;
+  }
+
+  LastAttributeEntry = MemoryAttributesEntry;
+  LastAttributeIndex = 0;
+  MemoryMapEntry     = MemoryMap;
+  CurrentEntryCount  = *MemoryMapSize / DescriptorSize;
+  TotalEntryCount    = OriginalMemoryMapSize / DescriptorSize;
+  AttributeCount     = MemoryAttributesTable->NumberOfEntries;
+
+  //
+  // We assume that the memory map and attribute table are sorted.
+  // If this is not the case we can add a sort call here.
+  //
+  for (Index = 0; Index < CurrentEntryCount; ++Index) {
+    //
+    // 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) {
-          ++DescriptorCount;
+          //
+          // 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.
+            // FIXME: Do we need this to inspect attributes?
+            //
+            if (MemoryAttributesEntry->Type != 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 (
@@ -568,8 +791,37 @@ OcCountSplitDescritptors (
           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;
+        }
+      } else {
+        //
+        // Did not find a suitable attribute or processed all the attributes.
+        //
+        CanSplit = FALSE;
+      }
     }
+
+    MemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (
+      MemoryMapEntry,
+      DescriptorSize
+      );
   }
 
-  return DescriptorCount;
+  *MemoryMapSize = CurrentEntryCount * DescriptorSize;
+  return EFI_SUCCESS;
 }