[PATCH v4 1/2] ArmPkg/ArmMmuLib AARCH64: rewrite page table code

["Ard Biesheuvel" <ard.biesheuvel@linaro.org>]

Replace the slightly overcomplicated page table management code with
a simplified, recursive implementation that should be far easier to
reason about.

Note that, as a side effect, this extends the per-entry cache invalidation
that we do on page table entries to block and page entries, whereas the
previous change inadvertently only affected the creation of table entries.

Signed-off-by: Ard Biesheuvel <ard.biesheuvel@linaro.org>
---
 ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c | 391 ++++++++------------
 1 file changed, 148 insertions(+), 243 deletions(-)

diff --git a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c
index 204e33c75f95..00a38bc31d0a 100644
--- a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c
+++ b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c
@@ -1,7 +1,7 @@
 /** @file
 *  File managing the MMU for ARMv8 architecture
 *
-*  Copyright (c) 2011-2014, ARM Limited. All rights reserved.
+*  Copyright (c) 2011-2020, ARM Limited. All rights reserved.
 *  Copyright (c) 2016, Linaro Limited. All rights reserved.
 *  Copyright (c) 2017, Intel Corporation. All rights reserved.<BR>
 *
@@ -121,19 +121,150 @@ GetRootTranslationTableInfo (
 
 STATIC
 VOID
-ReplaceLiveEntry (
+ReplaceTableEntry (
   IN  UINT64  *Entry,
   IN  UINT64  Value,
-  IN  UINT64  RegionStart
+  IN  UINT64  RegionStart,
+  IN  BOOLEAN IsLiveBlockMapping
   )
 {
-  if (!ArmMmuEnabled ()) {
+  if (!ArmMmuEnabled () || !IsLiveBlockMapping) {
     *Entry = Value;
+    ArmUpdateTranslationTableEntry (Entry, (VOID *)(UINTN)RegionStart);
   } else {
     ArmReplaceLiveTranslationEntry (Entry, Value, RegionStart);
   }
 }
 
+STATIC
+VOID
+FreePageTablesRecursive (
+  IN  UINT64  *TranslationTable
+  )
+{
+  UINTN   Index;
+
+  for (Index = 0; Index < TT_ENTRY_COUNT; Index++) {
+    if ((TranslationTable[Index] & TT_TYPE_MASK) == TT_TYPE_TABLE_ENTRY) {
+      FreePageTablesRecursive ((VOID *)(UINTN)(TranslationTable[Index] &
+                                               TT_ADDRESS_MASK_BLOCK_ENTRY));
+    }
+  }
+  FreePages (TranslationTable, 1);
+}
+
+STATIC
+EFI_STATUS
+UpdateRegionMappingRecursive (
+  IN  UINT64      RegionStart,
+  IN  UINT64      RegionEnd,
+  IN  UINT64      AttributeSetMask,
+  IN  UINT64      AttributeClearMask,
+  IN  UINT64      *PageTable,
+  IN  UINTN       Level
+  )
+{
+  UINTN           BlockShift;
+  UINT64          BlockMask;
+  UINT64          BlockEnd;
+  UINT64          *Entry;
+  UINT64          EntryValue;
+  VOID            *TranslationTable;
+  EFI_STATUS      Status;
+
+  ASSERT (((RegionStart | RegionEnd) & EFI_PAGE_MASK) == 0);
+
+  BlockShift = (Level + 1) * BITS_PER_LEVEL + MIN_T0SZ;
+  BlockMask = MAX_UINT64 >> BlockShift;
+
+  DEBUG ((DEBUG_VERBOSE, "%a(%d): %llx - %llx set %lx clr %lx\n", __FUNCTION__,
+    Level, RegionStart, RegionEnd, AttributeSetMask, AttributeClearMask));
+
+  for (; RegionStart < RegionEnd; RegionStart = BlockEnd) {
+    BlockEnd = MIN (RegionEnd, (RegionStart | BlockMask) + 1);
+    Entry = &PageTable[(RegionStart >> (64 - BlockShift)) & (TT_ENTRY_COUNT - 1)];
+
+    //
+    // If RegionStart or BlockEnd is not aligned to the block size at this
+    // level, we will have to create a table mapping in order to map less
+    // than a block, and recurse to create the block or page entries at
+    // the next level. No block mappings are allowed at all at level 0,
+    // so in that case, we have to recurse unconditionally.
+    //
+    if (Level == 0 || ((RegionStart | BlockEnd) & BlockMask) != 0) {
+      ASSERT (Level < 3);
+
+      if ((*Entry & TT_TYPE_MASK) != TT_TYPE_TABLE_ENTRY) {
+        //
+        // No table entry exists yet, so we need to allocate a page table
+        // for the next level.
+        //
+        TranslationTable = AllocatePages (1);
+        if (TranslationTable == NULL) {
+          return EFI_OUT_OF_RESOURCES;
+        }
+
+        if ((*Entry & TT_TYPE_MASK) == TT_TYPE_BLOCK_ENTRY) {
+          //
+          // We are splitting an existing block entry, so we have to populate
+          // the new table with the attributes of the block entry it replaces.
+          //
+          Status = UpdateRegionMappingRecursive (RegionStart & ~BlockMask,
+                     (RegionStart | BlockMask) + 1, *Entry & TT_ATTRIBUTES_MASK,
+                     0, TranslationTable, Level + 1);
+          if (EFI_ERROR (Status)) {
+            //
+            // The range we passed to UpdateRegionMappingRecursive () is block
+            // aligned, so it is guaranteed that no further pages were allocated
+            // by it, and so we only have to free the page we allocated here.
+            //
+            FreePages (TranslationTable, 1);
+            return Status;
+          }
+        } else {
+          ZeroMem (TranslationTable, EFI_PAGE_SIZE);
+        }
+      } else {
+        TranslationTable = (VOID *)(UINTN)(*Entry & TT_ADDRESS_MASK_BLOCK_ENTRY);
+      }
+
+      //
+      // Recurse to the next level
+      //
+      Status = UpdateRegionMappingRecursive (RegionStart, BlockEnd,
+                 AttributeSetMask, AttributeClearMask, TranslationTable,
+                 Level + 1);
+      if (EFI_ERROR (Status)) {
+        if ((*Entry & TT_TYPE_MASK) != TT_TYPE_TABLE_ENTRY) {
+          //
+          // We are creating a new table entry, so on failure, we can free all
+          // allocations we made recursively, given that the whole subhierarchy
+          // has not been wired into the live page tables yet. (This is not
+          // possible for existing table entries, since we cannot revert the
+          // modifications we made to the subhierarchy it represents.)
+          //
+          FreePageTablesRecursive (TranslationTable);
+        }
+        return Status;
+      }
+
+      if ((*Entry & TT_TYPE_MASK) != TT_TYPE_TABLE_ENTRY) {
+        EntryValue = (UINTN)TranslationTable | TT_TYPE_TABLE_ENTRY;
+        ReplaceTableEntry (Entry, EntryValue, RegionStart,
+                           (*Entry & TT_TYPE_MASK) == TT_TYPE_BLOCK_ENTRY);
+      }
+    } else {
+      EntryValue = (*Entry & AttributeClearMask) | AttributeSetMask;
+      EntryValue |= RegionStart;
+      EntryValue |= (Level == 3) ? TT_TYPE_BLOCK_ENTRY_LEVEL3
+                                 : TT_TYPE_BLOCK_ENTRY;
+
+      ReplaceTableEntry (Entry, EntryValue, RegionStart, FALSE);
+    }
+  }
+  return EFI_SUCCESS;
+}
+
 STATIC
 VOID
 LookupAddresstoRootTable (
@@ -164,230 +295,28 @@ LookupAddresstoRootTable (
   GetRootTranslationTableInfo (*T0SZ, NULL, TableEntryCount);
 }
 
-STATIC
-UINT64*
-GetBlockEntryListFromAddress (
-  IN  UINT64       *RootTable,
-  IN  UINT64        RegionStart,
-  OUT UINTN        *TableLevel,
-  IN OUT UINT64    *BlockEntrySize,
-  OUT UINT64      **LastBlockEntry
-  )
-{
-  UINTN   RootTableLevel;
-  UINTN   RootTableEntryCount;
-  UINT64 *TranslationTable;
-  UINT64 *BlockEntry;
-  UINT64 *SubTableBlockEntry;
-  UINT64  BlockEntryAddress;
-  UINTN   BaseAddressAlignment;
-  UINTN   PageLevel;
-  UINTN   Index;
-  UINTN   IndexLevel;
-  UINTN   T0SZ;
-  UINT64  Attributes;
-  UINT64  TableAttributes;
-
-  // Initialize variable
-  BlockEntry = NULL;
-
-  // Ensure the parameters are valid
-  if (!(TableLevel && BlockEntrySize && LastBlockEntry)) {
-    ASSERT_EFI_ERROR (EFI_INVALID_PARAMETER);
-    return NULL;
-  }
-
-  // Ensure the Region is aligned on 4KB boundary
-  if ((RegionStart & (SIZE_4KB - 1)) != 0) {
-    ASSERT_EFI_ERROR (EFI_INVALID_PARAMETER);
-    return NULL;
-  }
-
-  // Ensure the required size is aligned on 4KB boundary and not 0
-  if ((*BlockEntrySize & (SIZE_4KB - 1)) != 0 || *BlockEntrySize == 0) {
-    ASSERT_EFI_ERROR (EFI_INVALID_PARAMETER);
-    return NULL;
-  }
-
-  T0SZ = ArmGetTCR () & TCR_T0SZ_MASK;
-  // Get the Table info from T0SZ
-  GetRootTranslationTableInfo (T0SZ, &RootTableLevel, &RootTableEntryCount);
-
-  // If the start address is 0x0 then we use the size of the region to identify the alignment
-  if (RegionStart == 0) {
-    // Identify the highest possible alignment for the Region Size
-    BaseAddressAlignment = LowBitSet64 (*BlockEntrySize);
-  } else {
-    // Identify the highest possible alignment for the Base Address
-    BaseAddressAlignment = LowBitSet64 (RegionStart);
-  }
-
-  // Identify the Page Level the RegionStart must belong to. Note that PageLevel
-  // should be at least 1 since block translations are not supported at level 0
-  PageLevel = MAX (3 - ((BaseAddressAlignment - 12) / 9), 1);
-
-  // If the required size is smaller than the current block size then we need to go to the page below.
-  // The PageLevel was calculated on the Base Address alignment but did not take in account the alignment
-  // of the allocation size
-  while (*BlockEntrySize < TT_BLOCK_ENTRY_SIZE_AT_LEVEL (PageLevel)) {
-    // It does not fit so we need to go a page level above
-    PageLevel++;
-  }
-
-  //
-  // Get the Table Descriptor for the corresponding PageLevel. We need to decompose RegionStart to get appropriate entries
-  //
-
-  TranslationTable = RootTable;
-  for (IndexLevel = RootTableLevel; IndexLevel <= PageLevel; IndexLevel++) {
-    BlockEntry = (UINT64*)TT_GET_ENTRY_FOR_ADDRESS (TranslationTable, IndexLevel, RegionStart);
-
-    if ((IndexLevel != 3) && ((*BlockEntry & TT_TYPE_MASK) == TT_TYPE_TABLE_ENTRY)) {
-      // Go to the next table
-      TranslationTable = (UINT64*)(*BlockEntry & TT_ADDRESS_MASK_DESCRIPTION_TABLE);
-
-      // If we are at the last level then update the last level to next level
-      if (IndexLevel == PageLevel) {
-        // Enter the next level
-        PageLevel++;
-      }
-    } else if ((*BlockEntry & TT_TYPE_MASK) == TT_TYPE_BLOCK_ENTRY) {
-      // If we are not at the last level then we need to split this BlockEntry
-      if (IndexLevel != PageLevel) {
-        // Retrieve the attributes from the block entry
-        Attributes = *BlockEntry & TT_ATTRIBUTES_MASK;
-
-        // Convert the block entry attributes into Table descriptor attributes
-        TableAttributes = TT_TABLE_AP_NO_PERMISSION;
-        if (Attributes & TT_NS) {
-          TableAttributes = TT_TABLE_NS;
-        }
-
-        // Get the address corresponding at this entry
-        BlockEntryAddress = RegionStart;
-        BlockEntryAddress = BlockEntryAddress >> TT_ADDRESS_OFFSET_AT_LEVEL(IndexLevel);
-        // Shift back to right to set zero before the effective address
-        BlockEntryAddress = BlockEntryAddress << TT_ADDRESS_OFFSET_AT_LEVEL(IndexLevel);
-
-        // Set the correct entry type for the next page level
-        if ((IndexLevel + 1) == 3) {
-          Attributes |= TT_TYPE_BLOCK_ENTRY_LEVEL3;
-        } else {
-          Attributes |= TT_TYPE_BLOCK_ENTRY;
-        }
-
-        // Create a new translation table
-        TranslationTable = AllocatePages (1);
-        if (TranslationTable == NULL) {
-          return NULL;
-        }
-
-        // Populate the newly created lower level table
-        SubTableBlockEntry = TranslationTable;
-        for (Index = 0; Index < TT_ENTRY_COUNT; Index++) {
-          *SubTableBlockEntry = Attributes | (BlockEntryAddress + (Index << TT_ADDRESS_OFFSET_AT_LEVEL(IndexLevel + 1)));
-          SubTableBlockEntry++;
-        }
-
-        // Fill the BlockEntry with the new TranslationTable
-        ReplaceLiveEntry (BlockEntry,
-          (UINTN)TranslationTable | TableAttributes | TT_TYPE_TABLE_ENTRY,
-          RegionStart);
-      }
-    } else {
-      if (IndexLevel != PageLevel) {
-        //
-        // Case when we have an Invalid Entry and we are at a page level above of the one targetted.
-        //
-
-        // Create a new translation table
-        TranslationTable = AllocatePages (1);
-        if (TranslationTable == NULL) {
-          return NULL;
-        }
-
-        ZeroMem (TranslationTable, TT_ENTRY_COUNT * sizeof(UINT64));
-
-        // Fill the new BlockEntry with the TranslationTable
-        *BlockEntry = ((UINTN)TranslationTable & TT_ADDRESS_MASK_DESCRIPTION_TABLE) | TT_TYPE_TABLE_ENTRY;
-      }
-    }
-  }
-
-  // Expose the found PageLevel to the caller
-  *TableLevel = PageLevel;
-
-  // Now, we have the Table Level we can get the Block Size associated to this table
-  *BlockEntrySize = TT_BLOCK_ENTRY_SIZE_AT_LEVEL (PageLevel);
-
-  // The last block of the root table depends on the number of entry in this table,
-  // otherwise it is always the (TT_ENTRY_COUNT - 1)th entry in the table.
-  *LastBlockEntry = TT_LAST_BLOCK_ADDRESS(TranslationTable,
-      (PageLevel == RootTableLevel) ? RootTableEntryCount : TT_ENTRY_COUNT);
-
-  return BlockEntry;
-}
-
 STATIC
 EFI_STATUS
 UpdateRegionMapping (
-  IN  UINT64  *RootTable,
   IN  UINT64  RegionStart,
   IN  UINT64  RegionLength,
-  IN  UINT64  Attributes,
-  IN  UINT64  BlockEntryMask
+  IN  UINT64  AttributeSetMask,
+  IN  UINT64  AttributeClearMask
   )
 {
-  UINT32  Type;
-  UINT64  *BlockEntry;
-  UINT64  *LastBlockEntry;
-  UINT64  BlockEntrySize;
-  UINTN   TableLevel;
+  UINTN     RootTableLevel;
+  UINTN     T0SZ;
 
-  // Ensure the Length is aligned on 4KB boundary
-  if ((RegionLength == 0) || ((RegionLength & (SIZE_4KB - 1)) != 0)) {
-    ASSERT_EFI_ERROR (EFI_INVALID_PARAMETER);
+  if (((RegionStart | RegionLength) & EFI_PAGE_MASK)) {
     return EFI_INVALID_PARAMETER;
   }
 
-  do {
-    // Get the first Block Entry that matches the Virtual Address and also the information on the Table Descriptor
-    // such as the size of the Block Entry and the address of the last BlockEntry of the Table Descriptor
-    BlockEntrySize = RegionLength;
-    BlockEntry = GetBlockEntryListFromAddress (RootTable, RegionStart, &TableLevel, &BlockEntrySize, &LastBlockEntry);
-    if (BlockEntry == NULL) {
-      // GetBlockEntryListFromAddress() return NULL when it fails to allocate new pages from the Translation Tables
-      return EFI_OUT_OF_RESOURCES;
-    }
+  T0SZ = ArmGetTCR () & TCR_T0SZ_MASK;
+  GetRootTranslationTableInfo (T0SZ, &RootTableLevel, NULL);
 
-    if (TableLevel != 3) {
-      Type = TT_TYPE_BLOCK_ENTRY;
-    } else {
-      Type = TT_TYPE_BLOCK_ENTRY_LEVEL3;
-    }
-
-    do {
-      // Fill the Block Entry with attribute and output block address
-      *BlockEntry &= BlockEntryMask;
-      *BlockEntry |= (RegionStart & TT_ADDRESS_MASK_BLOCK_ENTRY) | Attributes | Type;
-
-      ArmUpdateTranslationTableEntry (BlockEntry, (VOID *)RegionStart);
-
-      // Go to the next BlockEntry
-      RegionStart += BlockEntrySize;
-      RegionLength -= BlockEntrySize;
-      BlockEntry++;
-
-      // Break the inner loop when next block is a table
-      // Rerun GetBlockEntryListFromAddress to avoid page table memory leak
-      if (TableLevel != 3 && BlockEntry <= LastBlockEntry &&
-          (*BlockEntry & TT_TYPE_MASK) == TT_TYPE_TABLE_ENTRY) {
-            break;
-      }
-    } while ((RegionLength >= BlockEntrySize) && (BlockEntry <= LastBlockEntry));
-  } while (RegionLength != 0);
-
-  return EFI_SUCCESS;
+  return UpdateRegionMappingRecursive (RegionStart, RegionStart + RegionLength,
+           AttributeSetMask, AttributeClearMask, ArmGetTTBR0BaseAddress (),
+           RootTableLevel);
 }
 
 STATIC
@@ -398,7 +327,6 @@ FillTranslationTable (
   )
 {
   return UpdateRegionMapping (
-           RootTable,
            MemoryRegion->VirtualBase,
            MemoryRegion->Length,
            ArmMemoryAttributeToPageAttribute (MemoryRegion->Attributes) | TT_AF,
@@ -455,8 +383,6 @@ ArmSetMemoryAttributes (
   IN UINT64                    Attributes
   )
 {
-  EFI_STATUS                   Status;
-  UINT64                      *TranslationTable;
   UINT64                       PageAttributes;
   UINT64                       PageAttributeMask;
 
@@ -473,19 +399,8 @@ ArmSetMemoryAttributes (
                           TT_PXN_MASK | TT_XN_MASK);
   }
 
-  TranslationTable = ArmGetTTBR0BaseAddress ();
-
-  Status = UpdateRegionMapping (
-             TranslationTable,
-             BaseAddress,
-             Length,
-             PageAttributes,
-             PageAttributeMask);
-  if (EFI_ERROR (Status)) {
-    return Status;
-  }
-
-  return EFI_SUCCESS;
+  return UpdateRegionMapping (BaseAddress, Length, PageAttributes,
+           PageAttributeMask);
 }
 
 STATIC
@@ -497,17 +412,7 @@ SetMemoryRegionAttribute (
   IN  UINT64                    BlockEntryMask
   )
 {
-  EFI_STATUS                   Status;
-  UINT64                       *RootTable;
-
-  RootTable = ArmGetTTBR0BaseAddress ();
-
-  Status = UpdateRegionMapping (RootTable, BaseAddress, Length, Attributes, BlockEntryMask);
-  if (EFI_ERROR (Status)) {
-    return Status;
-  }
-
-  return EFI_SUCCESS;
+  return UpdateRegionMapping (BaseAddress, Length, Attributes, BlockEntryMask);
 }
 
 EFI_STATUS
-- 
2.17.1