MdeModulePkg/Core: allow HeapGuard even before CpuArchProtocol installed
[mirror_edk2.git] / MdeModulePkg / Core / Dxe / Mem / HeapGuard.c
index ac043b5..f6068c4 100644 (file)
@@ -576,6 +576,10 @@ SetGuardPage (
   IN  EFI_PHYSICAL_ADDRESS      BaseAddress\r
   )\r
 {\r
+  if (gCpu == NULL) {\r
+    return;\r
+  }\r
+\r
   //\r
   // Set flag to make sure allocating memory without GUARD for page table\r
   // operation; otherwise infinite loops could be caused.\r
@@ -606,6 +610,10 @@ UnsetGuardPage (
 {\r
   UINT64          Attributes;\r
 \r
+  if (gCpu == NULL) {\r
+    return;\r
+  }\r
+\r
   //\r
   // Once the Guard page is unset, it will be freed back to memory pool. NX\r
   // memory protection must be restored for this page if NX is enabled for free\r
@@ -652,7 +660,7 @@ IsMemoryTypeToGuard (
   UINT64 ConfigBit;\r
   BOOLEAN     InSmm;\r
 \r
-  if (gCpu == NULL || AllocateType == AllocateAddress) {\r
+  if (AllocateType == AllocateAddress) {\r
     return FALSE;\r
   }\r
 \r
@@ -1164,6 +1172,128 @@ CoreConvertPagesWithGuard (
   return CoreConvertPages (Start, NumberOfPages, NewType);\r
 }\r
 \r
+/**\r
+  Set all Guard pages which cannot be set before CPU Arch Protocol installed.\r
+**/\r
+VOID\r
+SetAllGuardPages (\r
+  VOID\r
+  )\r
+{\r
+  UINTN     Entries[GUARDED_HEAP_MAP_TABLE_DEPTH];\r
+  UINTN     Shifts[GUARDED_HEAP_MAP_TABLE_DEPTH];\r
+  UINTN     Indices[GUARDED_HEAP_MAP_TABLE_DEPTH];\r
+  UINT64    Tables[GUARDED_HEAP_MAP_TABLE_DEPTH];\r
+  UINT64    Addresses[GUARDED_HEAP_MAP_TABLE_DEPTH];\r
+  UINT64    TableEntry;\r
+  UINT64    Address;\r
+  UINT64    GuardPage;\r
+  INTN      Level;\r
+  UINTN     Index;\r
+  BOOLEAN   OnGuarding;\r
+\r
+  if (mGuardedMemoryMap == 0 ||\r
+      mMapLevel == 0 ||\r
+      mMapLevel > GUARDED_HEAP_MAP_TABLE_DEPTH) {\r
+    return;\r
+  }\r
+\r
+  CopyMem (Entries, mLevelMask, sizeof (Entries));\r
+  CopyMem (Shifts, mLevelShift, sizeof (Shifts));\r
+\r
+  SetMem (Tables, sizeof(Tables), 0);\r
+  SetMem (Addresses, sizeof(Addresses), 0);\r
+  SetMem (Indices, sizeof(Indices), 0);\r
+\r
+  Level         = GUARDED_HEAP_MAP_TABLE_DEPTH - mMapLevel;\r
+  Tables[Level] = mGuardedMemoryMap;\r
+  Address       = 0;\r
+  OnGuarding    = FALSE;\r
+\r
+  DEBUG_CODE (\r
+    DumpGuardedMemoryBitmap ();\r
+  );\r
+\r
+  while (TRUE) {\r
+    if (Indices[Level] > Entries[Level]) {\r
+      Tables[Level] = 0;\r
+      Level        -= 1;\r
+    } else {\r
+\r
+      TableEntry  = ((UINT64 *)(UINTN)(Tables[Level]))[Indices[Level]];\r
+      Address     = Addresses[Level];\r
+\r
+      if (TableEntry == 0) {\r
+\r
+        OnGuarding = FALSE;\r
+\r
+      } else if (Level < GUARDED_HEAP_MAP_TABLE_DEPTH - 1) {\r
+\r
+        Level            += 1;\r
+        Tables[Level]     = TableEntry;\r
+        Addresses[Level]  = Address;\r
+        Indices[Level]    = 0;\r
+\r
+        continue;\r
+\r
+      } else {\r
+\r
+        Index = 0;\r
+        while (Index < GUARDED_HEAP_MAP_ENTRY_BITS) {\r
+          if ((TableEntry & 1) == 1) {\r
+            if (OnGuarding) {\r
+              GuardPage = 0;\r
+            } else {\r
+              GuardPage = Address - EFI_PAGE_SIZE;\r
+            }\r
+            OnGuarding = TRUE;\r
+          } else {\r
+            if (OnGuarding) {\r
+              GuardPage = Address;\r
+            } else {\r
+              GuardPage = 0;\r
+            }\r
+            OnGuarding = FALSE;\r
+          }\r
+\r
+          if (GuardPage != 0) {\r
+            SetGuardPage (GuardPage);\r
+          }\r
+\r
+          if (TableEntry == 0) {\r
+            break;\r
+          }\r
+\r
+          TableEntry = RShiftU64 (TableEntry, 1);\r
+          Address   += EFI_PAGE_SIZE;\r
+          Index     += 1;\r
+        }\r
+      }\r
+    }\r
+\r
+    if (Level < (GUARDED_HEAP_MAP_TABLE_DEPTH - (INTN)mMapLevel)) {\r
+      break;\r
+    }\r
+\r
+    Indices[Level] += 1;\r
+    Address = (Level == 0) ? 0 : Addresses[Level - 1];\r
+    Addresses[Level] = Address | LShiftU64(Indices[Level], Shifts[Level]);\r
+\r
+  }\r
+}\r
+\r
+/**\r
+  Notify function used to set all Guard pages before CPU Arch Protocol installed.\r
+**/\r
+VOID\r
+HeapGuardCpuArchProtocolNotify (\r
+  VOID\r
+  )\r
+{\r
+  ASSERT (gCpu != NULL);\r
+  SetAllGuardPages ();\r
+}\r
+\r
 /**\r
   Helper function to convert a UINT64 value in binary to a string.\r
 \r