MdeModulePkg PiSmmCore: Enhance SMM FreePool to catch buffer overflow

[mirror_edk2.git] / MdeModulePkg / Core / PiSmmCore / Pool.c

diff --git a/MdeModulePkg/Core/PiSmmCore/Pool.c b/MdeModulePkg/Core/PiSmmCore/Pool.c
index 43ce869d1e1a4d49efb4ff7d6b169b837b33af35..ebb9f8c49ee97e86ba7ebab07f7614e93c935590 100644 (file)
--- a/MdeModulePkg/Core/PiSmmCore/Pool.c
+++ b/MdeModulePkg/Core/PiSmmCore/Pool.c
@@ -133,6 +133,7 @@ InternalAllocPoolByIndex (
 {\r
   EFI_STATUS            Status;\r
   FREE_POOL_HEADER      *Hdr;\r
+  POOL_TAIL             *Tail;\r
   EFI_PHYSICAL_ADDRESS  Address;\r
   SMM_POOL_TYPE         SmmPoolType;\r
 \r
@@ -154,18 +155,26 @@ InternalAllocPoolByIndex (
   } else {\r
     Status = InternalAllocPoolByIndex (PoolType, PoolIndex + 1, &Hdr);\r
     if (!EFI_ERROR (Status)) {\r
+      Hdr->Header.Signature = 0;\r
       Hdr->Header.Size >>= 1;\r
       Hdr->Header.Available = TRUE;\r
-      Hdr->Header.Type = PoolType;\r
+      Hdr->Header.Type = 0;\r
+      Tail = HEAD_TO_TAIL(&Hdr->Header);\r
+      Tail->Signature = 0;\r
+      Tail->Size = 0;\r
       InsertHeadList (&mSmmPoolLists[SmmPoolType][PoolIndex], &Hdr->Link);\r
       Hdr = (FREE_POOL_HEADER*)((UINT8*)Hdr + Hdr->Header.Size);\r
     }\r
   }\r
 \r
   if (!EFI_ERROR (Status)) {\r
+    Hdr->Header.Signature = POOL_HEAD_SIGNATURE;\r
     Hdr->Header.Size = MIN_POOL_SIZE << PoolIndex;\r
     Hdr->Header.Available = FALSE;\r
     Hdr->Header.Type = PoolType;\r
+    Tail = HEAD_TO_TAIL(&Hdr->Header);\r
+    Tail->Signature = POOL_TAIL_SIGNATURE;\r
+    Tail->Size = Hdr->Header.Size;\r
   }\r
 \r
   *FreePoolHdr = Hdr;\r
@@ -187,6 +196,7 @@ InternalFreePoolByIndex (
 {\r
   UINTN                 PoolIndex;\r
   SMM_POOL_TYPE         SmmPoolType;\r
+  POOL_TAIL             *PoolTail;\r
 \r
   ASSERT ((FreePoolHdr->Header.Size & (FreePoolHdr->Header.Size - 1)) == 0);\r
   ASSERT (((UINTN)FreePoolHdr & (FreePoolHdr->Header.Size - 1)) == 0);\r
@@ -195,7 +205,12 @@ InternalFreePoolByIndex (
   SmmPoolType = UefiMemoryTypeToSmmPoolType(FreePoolHdr->Header.Type);\r
 \r
   PoolIndex = (UINTN) (HighBitSet32 ((UINT32)FreePoolHdr->Header.Size) - MIN_POOL_SHIFT);\r
+  FreePoolHdr->Header.Signature = 0;\r
   FreePoolHdr->Header.Available = TRUE;\r
+  FreePoolHdr->Header.Type = 0;\r
+  PoolTail = HEAD_TO_TAIL(&FreePoolHdr->Header);\r
+  PoolTail->Signature = 0;\r
+  PoolTail->Size = 0;\r
   ASSERT (PoolIndex < MAX_POOL_INDEX);\r
   InsertHeadList (&mSmmPoolLists[SmmPoolType][PoolIndex], &FreePoolHdr->Link);\r
   return EFI_SUCCESS;\r
@@ -223,6 +238,7 @@ SmmInternalAllocatePool (
   )\r
 {\r
   POOL_HEADER           *PoolHdr;\r
+  POOL_TAIL             *PoolTail;\r
   FREE_POOL_HEADER      *FreePoolHdr;\r
   EFI_STATUS            Status;\r
   EFI_PHYSICAL_ADDRESS  Address;\r
@@ -235,7 +251,10 @@ SmmInternalAllocatePool (
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r
-  Size += sizeof (*PoolHdr);\r
+  //\r
+  // Adjust the size by the pool header & tail overhead\r
+  //\r
+  Size += POOL_OVERHEAD;\r
   if (Size > MAX_POOL_SIZE) {\r
     Size = EFI_SIZE_TO_PAGES (Size);\r
     Status = SmmInternalAllocatePages (AllocateAnyPages, PoolType, Size, &Address);\r
@@ -244,9 +263,13 @@ SmmInternalAllocatePool (
     }\r
 \r
     PoolHdr = (POOL_HEADER*)(UINTN)Address;\r
+    PoolHdr->Signature = POOL_HEAD_SIGNATURE;\r
     PoolHdr->Size = EFI_PAGES_TO_SIZE (Size);\r
     PoolHdr->Available = FALSE;\r
     PoolHdr->Type = PoolType;\r
+    PoolTail = HEAD_TO_TAIL(PoolHdr);\r
+    PoolTail->Signature = POOL_TAIL_SIGNATURE;\r
+    PoolTail->Size = PoolHdr->Size;\r
     *Buffer = PoolHdr + 1;\r
     return Status;\r
   }\r
@@ -317,13 +340,30 @@ SmmInternalFreePool (
   )\r
 {\r
   FREE_POOL_HEADER  *FreePoolHdr;\r
+  POOL_TAIL         *PoolTail;\r
 \r
   if (Buffer == NULL) {\r
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r
   FreePoolHdr = (FREE_POOL_HEADER*)((POOL_HEADER*)Buffer - 1);\r
+  ASSERT (FreePoolHdr->Header.Signature == POOL_HEAD_SIGNATURE);\r
   ASSERT (!FreePoolHdr->Header.Available);\r
+  PoolTail = HEAD_TO_TAIL(&FreePoolHdr->Header);\r
+  ASSERT (PoolTail->Signature == POOL_TAIL_SIGNATURE);\r
+  ASSERT (FreePoolHdr->Header.Size == PoolTail->Size);\r
+\r
+  if (FreePoolHdr->Header.Signature != POOL_HEAD_SIGNATURE) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  if (PoolTail->Signature != POOL_TAIL_SIGNATURE) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  if (FreePoolHdr->Header.Size != PoolTail->Size) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
 \r
   if (FreePoolHdr->Header.Size > MAX_POOL_SIZE) {\r
     ASSERT (((UINTN)FreePoolHdr & EFI_PAGE_MASK) == 0);\r