From 2ef8e235c6d19da1238233267f5baa764b6664fa Mon Sep 17 00:00:00 2001
From: Sour <mesenemu@gmail.com>
Date: Sun, 20 Jan 2019 15:04:23 -0500
Subject: [PATCH] Fixed infinite loop when trying to load invalid rom files
 (when file size does not match header)

---
 Core/NESHeader.cpp  | 27 ---------------------------
 Core/NESHeader.h    |  1 -
 Core/iNesLoader.cpp | 18 +++++++++++++-----
 3 files changed, 13 insertions(+), 33 deletions(-)

diff --git a/Core/NESHeader.cpp b/Core/NESHeader.cpp
index 5139cfce..c755317f 100644
--- a/Core/NESHeader.cpp
+++ b/Core/NESHeader.cpp
@@ -248,30 +248,3 @@ PpuModel NESHeader::GetVsSystemPpuModel()
 	}
 	return PpuModel::Ppu2C03;
 }
-
-void NESHeader::SanitizeHeader(size_t romLength)
-{
-	uint32_t originalPrgSize = GetPrgSize();
-	uint32_t originalChrSize = GetChrSize();
-
-	size_t calculatedLength = sizeof(NESHeader) + GetPrgSize();
-	while(calculatedLength > romLength) {
-		Byte9 = 0;
-		PrgCount--;
-		calculatedLength = sizeof(NESHeader) + GetPrgSize();
-	}
-
-	calculatedLength = sizeof(NESHeader) + GetPrgSize() + GetChrSize();
-	while(calculatedLength > romLength) {
-		Byte9 = 0;
-		ChrCount--;
-		calculatedLength = sizeof(NESHeader) + GetPrgSize() + GetChrSize();
-	}
-
-	if(originalPrgSize != GetPrgSize()) {
-		MessageManager::Log("[iNes] Invalid ROM file length - PRG data has been truncated.");
-	}
-	if(originalChrSize != GetChrSize()) {
-		MessageManager::Log("[iNes] Invalid ROM file length - CHR data has been truncated.");
-	}
-}
diff --git a/Core/NESHeader.h b/Core/NESHeader.h
index bbdc05ce..bf1df39b 100644
--- a/Core/NESHeader.h
+++ b/Core/NESHeader.h
@@ -48,5 +48,4 @@ struct NESHeader
 	GameInputType GetInputType();
 	VsSystemType GetVsSystemType();
 	PpuModel GetVsSystemPpuModel();
-	void SanitizeHeader(size_t romLength);
 };
diff --git a/Core/iNesLoader.cpp b/Core/iNesLoader.cpp
index 4a93d9a5..3e800f0c 100644
--- a/Core/iNesLoader.cpp
+++ b/Core/iNesLoader.cpp
@@ -15,11 +15,9 @@ RomData iNesLoader::LoadRom(vector<uint8_t>& romFile, NESHeader *preloadedHeader
 	uint32_t dataSize = (uint32_t)romFile.size();
 	if(preloadedHeader) {
 		header = *preloadedHeader;
-		header.SanitizeHeader(romFile.size() + sizeof(NESHeader));
 	} else {
 		memcpy((char*)&header, buffer, sizeof(NESHeader));
 		buffer += sizeof(NESHeader);
-		header.SanitizeHeader(romFile.size());
 		dataSize -= sizeof(NESHeader);
 	}
 
@@ -46,9 +44,16 @@ RomData iNesLoader::LoadRom(vector<uint8_t>& romFile, NESHeader *preloadedHeader
 	romData.SaveRamSize = header.GetSaveRamSize();
 
 	if(romData.Info.HasTrainer) {
-		//512-byte trainer at $7000-$71FF (stored before PRG data)
-		romData.TrainerData.insert(romData.TrainerData.end(), buffer, buffer + 512);
-		buffer += 512;
+		if(dataSize >= 512) {
+			//512-byte trainer at $7000-$71FF (stored before PRG data)
+			romData.TrainerData.insert(romData.TrainerData.end(), buffer, buffer + 512);
+			buffer += 512;
+			dataSize -= 512;
+		} else {
+			romData.Error = true;
+			MessageManager::Log("[iNes] Invalid file (file length does not match header information) - load operation cancelled.");
+			return romData;
+		}
 	}
 
 	size_t bytesRead = buffer - romFile.data();
@@ -68,8 +73,11 @@ RomData iNesLoader::LoadRom(vector<uint8_t>& romFile, NESHeader *preloadedHeader
 
 	if(prgSize + chrSize > dataSize) {
 		//Invalid rom file
+		MessageManager::Log("[iNes] Invalid file (file length does not match header information) - load operation cancelled.");
 		romData.Error = true;
 		return romData;
+	} else if(prgSize + chrSize < dataSize) {
+		MessageManager::Log("[iNes] Warning: File is larger than excepted (based on the file header).");
 	}
 
 	romData.PrgRom.insert(romData.PrgRom.end(), buffer, buffer + prgSize);