MdeModulePkg/Core: fix bits operation error on a boundary condition

[mirror_edk2.git] / MdeModulePkg / Core / Dxe / Mem / HeapGuard.c

diff --git a/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c b/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
index 19245049c28b847c66665c5ea08a4a967632dd0b..fd6aeee8da9d70c5e7523d41064cdc2b85e5c6df 100644 (file)
--- a/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
+++ b/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
@@ -70,7 +70,7 @@ SetBits (
   StartBit  = (UINTN)GUARDED_HEAP_MAP_ENTRY_BIT_INDEX (Address);\r
   EndBit    = (StartBit + BitNumber - 1) % GUARDED_HEAP_MAP_ENTRY_BITS;\r
 \r
-  if ((StartBit + BitNumber) > GUARDED_HEAP_MAP_ENTRY_BITS) {\r
+  if ((StartBit + BitNumber) >= GUARDED_HEAP_MAP_ENTRY_BITS) {\r
     Msbs    = (GUARDED_HEAP_MAP_ENTRY_BITS - StartBit) %\r
               GUARDED_HEAP_MAP_ENTRY_BITS;\r
     Lsbs    = (EndBit + 1) % GUARDED_HEAP_MAP_ENTRY_BITS;\r
@@ -123,7 +123,7 @@ ClearBits (
   StartBit  = (UINTN)GUARDED_HEAP_MAP_ENTRY_BIT_INDEX (Address);\r
   EndBit    = (StartBit + BitNumber - 1) % GUARDED_HEAP_MAP_ENTRY_BITS;\r
 \r
-  if ((StartBit + BitNumber) > GUARDED_HEAP_MAP_ENTRY_BITS) {\r
+  if ((StartBit + BitNumber) >= GUARDED_HEAP_MAP_ENTRY_BITS) {\r
     Msbs    = (GUARDED_HEAP_MAP_ENTRY_BITS - StartBit) %\r
               GUARDED_HEAP_MAP_ENTRY_BITS;\r
     Lsbs    = (EndBit + 1) % GUARDED_HEAP_MAP_ENTRY_BITS;\r
@@ -188,10 +188,14 @@ GetBits (
     Lsbs = 0;\r
   }\r
 \r
-  Result    = RShiftU64 ((*BitMap), StartBit) & (LShiftU64 (1, Msbs) - 1);\r
-  if (Lsbs > 0) {\r
-    BitMap  += 1;\r
-    Result  |= LShiftU64 ((*BitMap) & (LShiftU64 (1, Lsbs) - 1), Msbs);\r
+  if (StartBit == 0 && BitNumber == GUARDED_HEAP_MAP_ENTRY_BITS) {\r
+    Result = *BitMap;\r
+  } else {\r
+    Result    = RShiftU64((*BitMap), StartBit) & (LShiftU64(1, Msbs) - 1);\r
+    if (Lsbs > 0) {\r
+      BitMap  += 1;\r
+      Result  |= LShiftU64 ((*BitMap) & (LShiftU64 (1, Lsbs) - 1), Msbs);\r
+    }\r
   }\r
 \r
   return Result;\r
@@ -225,8 +229,8 @@ FindGuardedMemoryMap (
   //\r
   // Adjust current map table depth according to the address to access\r
   //\r
-  while (mMapLevel < GUARDED_HEAP_MAP_TABLE_DEPTH\r
-         &&\r
+  while (AllocMapUnit &&\r
+         mMapLevel < GUARDED_HEAP_MAP_TABLE_DEPTH &&\r
          RShiftU64 (\r
            Address,\r
            mLevelShift[GUARDED_HEAP_MAP_TABLE_DEPTH - mMapLevel - 1]\r
@@ -576,6 +580,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 +614,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 +664,7 @@ IsMemoryTypeToGuard (
   UINT64 ConfigBit;\r
   BOOLEAN     InSmm;\r
 \r
-  if (gCpu == NULL || AllocateType == AllocateAddress) {\r
+  if (AllocateType == AllocateAddress) {\r
     return FALSE;\r
   }\r
 \r
@@ -904,6 +916,10 @@ AdjustMemoryS (
   }\r
 \r
   Target = Start + Size - SizeRequested;\r
+  ASSERT (Target >= Start);\r
+  if (Target == 0) {\r
+    return 0;\r
+  }\r
 \r
   if (!IsGuardPage (Start + Size)) {\r
     // No Guard at tail to share. One more page is needed.\r
@@ -1160,6 +1176,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