diff --git a/app/aboot/aboot.c b/app/aboot/aboot.c
index f8cc588..be347ce 100755
--- a/app/aboot/aboot.c
+++ b/app/aboot/aboot.c
@@ -105,6 +105,7 @@
 void write_device_info_flash(device_info *dev);
 static int aboot_save_boot_hash_mmc(uint32_t image_addr, uint32_t image_size);
 static int aboot_frp_unlock(char *pname, void *data, unsigned sz);
+static inline uint64_t validate_partition_size();
 bool pwr_key_is_pressed = false;
 
 static bool is_systemd_present=false;
@@ -3340,6 +3341,11 @@
 	unsigned extra = 0;
 	uint64_t partition_size = 0;
 
+	if((uintptr_t)data > (UINT_MAX - sz)) {
+		fastboot_fail("Cannot flash: image header corrupt");
+                return;
+        }
+
 	ptable = flash_get_ptable();
 	if (ptable == NULL) {
 		fastboot_fail("partition table doesn't exist");
@@ -3355,8 +3361,10 @@
 	}
 
 	if (!strcmp(ptn->name, "boot") || !strcmp(ptn->name, "recovery")) {
-		if (memcmp((void *)data, BOOT_MAGIC, BOOT_MAGIC_SIZE)) {
-			fastboot_fail("image is not a boot image");
+		if((sz > BOOT_MAGIC_SIZE) && (!memcmp((void *)data, BOOT_MAGIC, BOOT_MAGIC_SIZE))) {
+			dprintf(INFO, "Verified the BOOT_MAGIC in image header  \n");
+		} else {
+			fastboot_fail("Image is not a boot image");
 			return;
 		}
 	}
@@ -3367,22 +3375,23 @@
 		|| !strcmp(ptn->name, "recoveryfs")
 		|| !strcmp(ptn->name, "modem"))
 		extra = 1;
-	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;
+	else {
+		if (sz % page_size) {
+			fastboot_fail("Buffer size is not aligned to page_size");
+			return;
+		}
 	}
 
+	/*Checking partition_size for the possible integer overflow */
+	partition_size = validate_partition_size(ptn);
+
 	if (sz > partition_size) {
 		fastboot_fail("Image size too large");
 		return;
 	}
 
 	dprintf(INFO, "writing %d bytes to '%s'\n", sz, ptn->name);
-	if (!memcmp((void *)data, UBI_MAGIC, UBI_MAGIC_SIZE)) {
+	if ((sz > UBI_MAGIC_SIZE) && (!memcmp((void *)data, UBI_MAGIC, UBI_MAGIC_SIZE))) {
 		if (flash_ubi_img(ptn, data, sz)) {
 			fastboot_fail("flash write failure");
 			return;
@@ -3397,6 +3406,18 @@
 	fastboot_okay("");
 }
 
+
+static inline uint64_t validate_partition_size(struct ptentry *ptn)
+{
+	if (ptn->length && flash_num_pages_per_blk() && page_size) {
+		if ((ptn->length < ( UINT_MAX / flash_num_pages_per_blk())) && ((ptn->length * flash_num_pages_per_blk()) < ( UINT_MAX / page_size))) {
+			return ptn->length * flash_num_pages_per_blk() * page_size;
+		}
+        }
+	return 0;
+}
+
+
 void cmd_flash(const char *arg, void *data, unsigned sz)
 {
 	if(target_is_emmc_boot())