app: aboot: add size check when flashing on NAND

Added size check during flashing on NAND memory so that
nothing is written beyond the intended partition.

Change-Id: Ie1d3307fb8f7631d2b33a372c0683484f82c1cd2

FPIIM-1224

Change-Id: I0a28b488406f0595cd8ae9034dee8b3372d0e920
diff --git a/app/aboot/aboot.c b/app/aboot/aboot.c
index e975f5b..183f9a4 100755
--- a/app/aboot/aboot.c
+++ b/app/aboot/aboot.c
@@ -2185,6 +2185,7 @@
 	struct ptentry *ptn;
 	struct ptable *ptable;
 	unsigned extra = 0;
+	uint64_t partition_size = 0;
 
 	ptable = flash_get_ptable();
 	if (ptable == NULL) {
@@ -2217,6 +2218,17 @@
 	} else
 		sz = ROUND_TO_PAGE(sz, page_mask);
 
+	partition_size = (uint64_t)ptn->length * (uint64_t)flash_num_pages_per_blk() *  (uint64_t)flash_page_size();
+	if (partition_size > UINT_MAX) {
+		fastboot_fail("Invalid partition size");
+		return;
+	}
+
+	if (sz > partition_size) {
+		fastboot_fail("Image size too large");
+		return;
+	}
+
 	dprintf(INFO, "writing %d bytes to '%s'\n", sz, ptn->name);
 	if (flash_write(ptn, extra, data, sz)) {
 		fastboot_fail("flash write failure");
diff --git a/include/dev/flash.h b/include/dev/flash.h
index 9fa10bc..d2cea0e 100644
--- a/include/dev/flash.h
+++ b/include/dev/flash.h
@@ -76,6 +76,8 @@
 {
 	return flash_read_ext(ptn, 0, offset, data, bytes);
 }
+
+unsigned flash_num_pages_per_blk(void);
 unsigned flash_page_size(void);
 int flash_ecc_bch_enabled(void);
 
diff --git a/platform/msm_shared/qpic_nand.c b/platform/msm_shared/qpic_nand.c
index df8480d..43bbc78 100644
--- a/platform/msm_shared/qpic_nand.c
+++ b/platform/msm_shared/qpic_nand.c
@@ -1296,6 +1296,12 @@
 }
 
 unsigned
+flash_num_pages_per_blk(void)
+{
+	return flash.num_pages_per_blk;
+}
+
+unsigned
 flash_page_size(void)
 {
 	return flash.page_size;