Fix memory manager crashing upon second reboot

src/core/kernel/memory-manager/VMManager.cpp

@@ -397,6 +397,8 @@ void VMManager::RestorePersistentMemory()
 void VMManager::SavePersistentMemory()
 {
 	PersistedMemory* persisted_mem;
+	size_t num_persisted_ptes;
+	std::vector<PMMPTE> cached_persisted_ptes;
 	HANDLE handle;
 	LPVOID addr;
 	PMMPTE PointerPte;
@@ -404,8 +406,27 @@ void VMManager::SavePersistentMemory()
 	int i;
 
 	Lock();
+
+	num_persisted_ptes = 0;
+	PointerPte = GetPteAddress(CONTIGUOUS_MEMORY_BASE);
+
+	if (m_MmLayoutRetail) {
+		EndingPte = GetPteAddress(CONTIGUOUS_MEMORY_BASE + XBOX_CONTIGUOUS_MEMORY_SIZE - 1);
+	}
+	else {
+		EndingPte = GetPteAddress(CONTIGUOUS_MEMORY_BASE + CHIHIRO_CONTIGUOUS_MEMORY_SIZE - 1);
+	}
+
+	while (PointerPte <= EndingPte)
+	{
+		if (PointerPte->Hardware.Valid != 0 && PointerPte->Hardware.Persist != 0) {
+			cached_persisted_ptes.push_back(PointerPte);
+			num_persisted_ptes++;
+		}
+		PointerPte++;
+	}
 
-	handle = CreateFileMapping(INVALID_HANDLE_VALUE, NULL, PAGE_READWRITE, 0, m_NumPersistentPtes * PAGE_SIZE + m_NumPersistentPtes * 4 * 2 + sizeof(PersistedMemory), "PersistentMemory");
+	handle = CreateFileMapping(INVALID_HANDLE_VALUE, NULL, PAGE_READWRITE, 0, num_persisted_ptes * PAGE_SIZE + num_persisted_ptes * 4 * 2 + sizeof(PersistedMemory), "PersistentMemory");
 	if (handle == NULL) {
 		CxbxKrnlCleanup("Couldn't persist memory! CreateFileMapping failed with error 0x%08X", GetLastError());
 		return;
@@ -417,7 +438,7 @@ void VMManager::SavePersistentMemory()
 	}
 
 	persisted_mem = (PersistedMemory*)addr;
-	persisted_mem->NumOfPtes = m_NumPersistentPtes;
+	persisted_mem->NumOfPtes = num_persisted_ptes;
 
 	if (xboxkrnl::LaunchDataPage != xbnullptr) {
 		persisted_mem->LaunchFrameAddresses[0] = (VAddr)xboxkrnl::LaunchDataPage;
@@ -430,27 +451,15 @@ void VMManager::SavePersistentMemory()
 	}
 
 	i = 0;
-	PointerPte = GetPteAddress(CONTIGUOUS_MEMORY_BASE);
 
-	if (m_MmLayoutRetail) {
+	for (const auto &pte : cached_persisted_ptes) {
-		EndingPte = GetPteAddress(CONTIGUOUS_MEMORY_BASE + XBOX_CONTIGUOUS_MEMORY_SIZE - 1);
+		persisted_mem->Data[i] = GetVAddrMappedByPte(pte);
-	}
+		persisted_mem->Data[num_persisted_ptes + i] = pte->Default;
-	else {
+		memcpy(&persisted_mem->Data[num_persisted_ptes * 2 + i * ONE_KB], (void *)(persisted_mem->Data[i]), PAGE_SIZE);
-		EndingPte = GetPteAddress(CONTIGUOUS_MEMORY_BASE + CHIHIRO_CONTIGUOUS_MEMORY_SIZE - 1);
+		i++;
-	}
-
-	while (PointerPte <= EndingPte)
-	{
-		if (PointerPte->Hardware.Valid != 0 && PointerPte->Hardware.Persist != 0) {
-			persisted_mem->Data[i] = GetVAddrMappedByPte(PointerPte);
-			persisted_mem->Data[m_NumPersistentPtes + i] = PointerPte->Default;
-			memcpy(&persisted_mem->Data[m_NumPersistentPtes * 2 + i * ONE_KB], (void*)(persisted_mem->Data[i]), PAGE_SIZE);
-			i++;
-		}
-		PointerPte++;
 	}
 
-	assert(i == m_NumPersistentPtes);
+	assert(i == num_persisted_ptes);
 
 	ipc_send_gui_update(IPC_UPDATE_GUI::VM_PERSIST_MEM, 1);
 
@@ -617,7 +626,6 @@ void VMManager::PersistMemory(VAddr addr, size_t Size, bool bPersist)
 		while (PointerPte <= EndingPte)
 		{
 			PointerPte->Hardware.Persist = 1;
-			m_NumPersistentPtes++;
 			PointerPte++;
 		}
 	}
@@ -625,7 +633,6 @@ void VMManager::PersistMemory(VAddr addr, size_t Size, bool bPersist)
 		while (PointerPte <= EndingPte)
 		{
 			PointerPte->Hardware.Persist = 0;
-			m_NumPersistentPtes--;
 			PointerPte++;
 		}
 	}

src/core/kernel/memory-manager/VMManager.h

@@ -84,7 +84,8 @@ typedef enum _MemoryRegionType
 	COUNTRegion,
 }MemoryRegionType;
 
-
+#pragma warning(push)
+#pragma warning(disable: 4200)
 /* struct used to save the persistent memory between reboots */
 typedef struct _PersistedMemory
 {
@@ -92,6 +93,7 @@ typedef struct _PersistedMemory
 	VAddr LaunchFrameAddresses[2];
 	uint32_t Data[];
 }PersistedMemory;
+#pragma warning(pop)
 
 
 /* VMManager class */
@@ -171,8 +173,6 @@ class VMManager : public PhysicalMemory
 		DWORD m_AllocationGranularity = 0;
 		// number of bytes reserved with XBOX_MEM_RESERVE by XbAllocateVirtualMemory
 		size_t m_VirtualMemoryBytesReserved = 0;
-		// number of persisted ptes between quick reboots
-		size_t m_NumPersistentPtes = 0;
 
 		// same as AllocateContiguousMemory, but it allows to allocate beyond m_MaxContiguousPfn
 		VAddr AllocateContiguousMemoryInternal(PFN_COUNT NumberOfPages, PFN LowestPfn, PFN HighestPfn, PFN PfnAlignment, DWORD Perms);