Update the heap migration logic to avoid memory overflow.

Signed-off-by: Liming Gao <liming.gao@intel.com>
Reviewed-by: Star Zeng <star.zeng@intel.com>
git-svn-id: https://svn.code.sf.net/p/edk2/code/trunk/edk2@15036 6f19259b-4bc3-4df7-8a09-765794883524

diff --git a/MdeModulePkg/Core/Pei/Dispatcher/Dispatcher.c b/MdeModulePkg/Core/Pei/Dispatcher/Dispatcher.c
index b5e2a824553f22253938542817db1c10cabd7e94..635c70f7533aa4bad1bd15426af48c23b3b51fa8 100644 (file)
--- a/MdeModulePkg/Core/Pei/Dispatcher/Dispatcher.c
+++ b/MdeModulePkg/Core/Pei/Dispatcher/Dispatcher.c
@@ -661,6 +661,7 @@ PeiDispatcher (
   PEIM_FILE_HANDLE_EXTENDED_DATA      ExtendedData;\r
   EFI_PEI_TEMPORARY_RAM_SUPPORT_PPI   *TemporaryRamSupportPpi;\r
   UINT64                              NewStackSize;\r
+  UINTN                               HeapTemporaryRamSize;\r
   EFI_PHYSICAL_ADDRESS                BaseOfNewHeap;\r
   EFI_PHYSICAL_ADDRESS                TopOfNewStack;\r
   EFI_PHYSICAL_ADDRESS                TopOfOldStack;\r
@@ -1018,7 +1019,9 @@ PeiDispatcher (
                 //\r
                 // Migrate Heap\r
                 //\r
-                CopyMem ((UINT8 *) (UINTN) BaseOfNewHeap, (UINT8 *) PeiTemporaryRamBase, PeiTemporaryRamSize);\r
+                HeapTemporaryRamSize = (UINTN) (Private->HobList.HandoffInformationTable->EfiFreeMemoryBottom - Private->HobList.HandoffInformationTable->EfiMemoryBottom);\r
+                ASSERT (BaseOfNewHeap + HeapTemporaryRamSize <= Private->FreePhysicalMemoryTop);\r
+                CopyMem ((UINT8 *) (UINTN) BaseOfNewHeap, (UINT8 *) PeiTemporaryRamBase, HeapTemporaryRamSize);\r
                 \r
                 //\r
                 // Migrate Stack\r