1 files changed, 18 insertions, 7 deletions

diff --git a/heimdall/source/FlashAction.cpp b/heimdall/source/FlashAction.cpp
index d79235e..1ce8efb 100644
--- a/heimdall/source/FlashAction.cpp
+++ b/heimdall/source/FlashAction.cpp
@@ -444,21 +444,32 @@ static PitData *getPitData(const map<string, FILE *>& argumentFileMap, BridgeMan
 		FILE *localPitFile = localPitFileIt->second;
 
 		// Load the local pit file into memory.
-		unsigned char *pitFileBuffer = new unsigned char[4096];
-		memset(pitFileBuffer, 0, 4096);
-
 		fseek(localPitFile, 0, SEEK_END);
 		long localPitFileSize = ftell(localPitFile);
 		rewind(localPitFile);
 
+		unsigned char *pitFileBuffer = new unsigned char[localPitFileSize];
+		memset(pitFileBuffer, 0, localPitFileSize);
+
 		// dataRead is discarded, it's here to remove warnings.
 		int dataRead = fread(pitFileBuffer, 1, localPitFileSize, localPitFile);
-		rewind(localPitFile);
 
-		localPitData = new PitData();
-		localPitData->Unpack(pitFileBuffer);
+		if (dataRead > 0)
+		{
+			rewind(localPitFile);
 
-		delete [] pitFileBuffer;
+			localPitData = new PitData();
+			localPitData->Unpack(pitFileBuffer);
+
+			delete [] pitFileBuffer;
+		}
+		else
+		{
+			Interface::PrintError("Failed to read PIT file.\n");
+
+			delete [] pitFileBuffer;
+			return (nullptr);
+		}
 	}
 
 	if (repartition)