From 60ab9bbaffe31f5072195f0dcf352d6e1037bc85 Mon Sep 17 00:00:00 2001
From: Henrik Grimler <henrik@grimler.se>
Date: Thu, 9 Dec 2021 22:14:42 +0100
Subject: [PATCH] FlashAction: Make sure file fit partition before flashing

Partition size is calculated as block count times a block size of 512
for DeviceType 2 and 4096 for DeviceType 8.

Flashing stock android on some devices fail due to some partitions
being too small for the corresponding file.  This is the case on
klimtlte (and other exynos5420 devices), where sboot.bin does not fit
into the BOOTLOADER partition.

Tests done:

* Flash stock android (A500FXXS1CSB2) to a5lte, works fine
* Flash stock android (I9300XXUGPE1) to i9300, works fine
* Flash stock android (I9505XXUPQG1) to i9505, works fine
* Flash stock android (G930FXXU8ETI2) to herolte (with block
  size 4096), works fine
* Flash stock android (T705XXU1CPL1) to klimtlte, fails due to
  BOOTLOADER partition being smaller than sboot.bin, the file has a
  size of 1148160, and the partition 2046*512=1047552
---
 heimdall/source/FlashAction.cpp | 32 +++++++++++++++++++++++++++-----
 1 file changed, 27 insertions(+), 5 deletions(-)

diff --git a/heimdall/source/FlashAction.cpp b/heimdall/source/FlashAction.cpp
index a4eca07..64b8d0e 100644
--- a/heimdall/source/FlashAction.cpp
+++ b/heimdall/source/FlashAction.cpp
@@ -62,11 +62,13 @@ struct PartitionFile
 {
 	const char *argumentName;
 	FILE *file;
+	unsigned long fileSize;
 
-	PartitionFile(const char *argumentName, FILE *file)
+	PartitionFile(const char *argumentName, FILE *file, unsigned long fileSize)
 	{
 		this->argumentName = argumentName;
 		this->file = file;
+		this->fileSize = fileSize;
 	}
 };
 
@@ -113,6 +115,9 @@ static bool openFiles(Arguments& arguments, vector<PartitionFile>& partitionFile
 		{
 			const StringArgument *stringArgument = static_cast<const StringArgument *>(*it);
 			FILE *file = FileOpen(stringArgument->GetValue().c_str(), "rb");
+			FileSeek(file, 0, SEEK_END);
+			unsigned long fileSize = (unsigned long)FileTell(file);
+			FileRewind(file);
 
 			if (!file)
 			{
@@ -120,7 +125,7 @@ static bool openFiles(Arguments& arguments, vector<PartitionFile>& partitionFile
 				return (false);
 			}
 
-			partitionFiles.push_back(PartitionFile(argumentName.c_str(), file));
+			partitionFiles.push_back(PartitionFile(argumentName.c_str(), file, fileSize));
 		}
 	}
 
@@ -151,9 +156,7 @@ static bool sendTotalTransferSize(BridgeManager *bridgeManager, const vector<Par
 
 	for (vector<PartitionFile>::const_iterator it = partitionFiles.begin(); it != partitionFiles.end(); it++)
 	{
-		FileSeek(it->file, 0, SEEK_END);
-		totalBytes += (unsigned long)FileTell(it->file);
-		FileRewind(it->file);
+		totalBytes += it->fileSize;
 	}
 
 	if (repartition)
@@ -292,6 +295,25 @@ static bool flashPartitions(BridgeManager *bridgeManager, const vector<Partition
 	if (!setupPartitionFlashInfo(partitionFiles, pitData, partitionFlashInfos))
 		return (false);
 
+	/* Verify that the files we want to flash fit in partitions */
+	for (vector<PartitionFile>::const_iterator it = partitionFiles.begin(); it != partitionFiles.end(); it++)
+	{
+		const PitEntry *part = pitData->FindEntry(it->argumentName);
+		if (part->GetDeviceType() != PitEntry::kDeviceTypeMMC &&
+		    part->GetDeviceType() != PitEntry::kDeviceTypeMMC4096)
+			continue;
+		unsigned long partitionSize = part->GetBlockCount();
+		unsigned int blockSize = 512;
+		if (part->GetDeviceType() == PitEntry::kDeviceTypeMMC4096)
+			blockSize = 4096;
+		if (partitionSize > 0 && it->fileSize > partitionSize*blockSize)
+		{
+			Interface::PrintError("%s partition is too small for specified file\n",
+					      it->argumentName);
+			return (false);
+		}
+	}
+
 	// If we're repartitioning then we need to flash the PIT file first (if it is listed in the PIT file).
 	if (repartition)
 	{