mmc protect: Handled invalid GPT

Change-Id: Ie68e278230cb3ccfbd60cc03a7936ff1e226d47a
Signed-off-by: Gokul Sriram Palanisamy <gokulsri@codeaurora.org>
diff --git a/board/qca/arm/common/board_init.c b/board/qca/arm/common/board_init.c
index 19f1bf6..2677660 100644
--- a/board/qca/arm/common/board_init.c
+++ b/board/qca/arm/common/board_init.c
@@ -240,30 +240,32 @@
 #ifdef CONFIG_FLASH_PROTECT
 void board_flash_protect(void)
 {
-	unsigned int num_part;
+	int num_part;
 	int i;
+	int ret;
 #ifdef CONFIG_QCA_MMC
 	block_dev_desc_t *mmc_dev;
 	disk_partition_t info;
 
 	mmc_dev = mmc_get_dev(mmc_host.dev_num);
 	if (mmc_dev != NULL && mmc_dev->type != DEV_TYPE_UNKNOWN) {
-		ALLOC_CACHE_ALIGN_BUFFER_PAD(gpt_header,
-					     gpt_head, 1, mmc_dev->blksz);
-		if (mmc_dev->block_read(mmc_dev->dev,
-		    (lbaint_t)GPT_PRIMARY_PARTITION_TABLE_LBA,
-		    1, gpt_head) == 1) {
-			num_part = le32_to_cpu(gpt_head->num_partition_entries);
-			for (i = 1; i <= num_part; i++) {
-				if (!get_partition_info_efi(mmc_dev, i, &info)
-				    && info.readonly
-				    && !mmc_write_protect(mmc_host.mmc,
-							  info.start,
-							  info.size, 1))
-					printf("\"%s\""
-						"-protected MMC partition\n",
-						info.name);
-			}
+		num_part = get_partition_count_efi(mmc_dev);
+		if (num_part < 0) {
+			printf("Both primary & backup GPT are invalid, skipping mmc write protection.\n");
+			return;
+		}
+
+		for (i = 1; i <= num_part; i++) {
+			ret = get_partition_info_efi(mmc_dev, i, &info);
+			if (ret == -1)
+				return;
+			if (!ret && info.readonly
+			    && !mmc_write_protect(mmc_host.mmc,
+						  info.start,
+						  info.size, 1))
+				printf("\"%s\""
+					"-protected MMC partition\n",
+					info.name);
 		}
 	}
 #endif
diff --git a/disk/part_efi.c b/disk/part_efi.c
index 57840b0..6d8eacf 100644
--- a/disk/part_efi.c
+++ b/disk/part_efi.c
@@ -340,6 +340,37 @@
 	return 0;
 }
 
+int get_partition_count_efi(block_dev_desc_t * dev_desc)
+{
+	ALLOC_CACHE_ALIGN_BUFFER_PAD(gpt_header, gpt_head, 1, dev_desc->blksz);
+	gpt_entry *gpt_pte = NULL;
+
+	if (!dev_desc) {
+		printf("%s: Invalid Argument(s)\n", __func__);
+		return -1;
+	}
+
+	/* This function validates AND fills in the GPT header and PTE */
+	if (is_gpt_valid(dev_desc, GPT_PRIMARY_PARTITION_TABLE_LBA,
+			gpt_head, &gpt_pte) != 1) {
+		printf("%s: *** ERROR: Invalid GPT ***\n", __func__);
+		if (is_gpt_valid(dev_desc, (dev_desc->lba - 1),
+				 gpt_head, &gpt_pte) != 1) {
+			printf("%s: *** ERROR: Invalid Backup GPT ***\n",
+			       __func__);
+			if (gpt_pte != NULL)
+				free(gpt_pte);
+			return -1;
+		} else {
+			printf("%s: ***        Using Backup GPT ***\n",
+			       __func__);
+		}
+	}
+
+	free(gpt_pte);
+	return le32_to_cpu(gpt_head->num_partition_entries);
+}
+
 /**
  * set_protective_mbr(): Set the EFI protective MBR
  * @param dev_desc - block device descriptor
diff --git a/include/part.h b/include/part.h
index ce8b497..647a0e9 100644
--- a/include/part.h
+++ b/include/part.h
@@ -197,6 +197,7 @@
 	const char *name, disk_partition_t *info);
 void print_part_efi (block_dev_desc_t *dev_desc);
 int   test_part_efi (block_dev_desc_t *dev_desc);
+int get_partition_count_efi(block_dev_desc_t * dev_desc);
 
 /**
  * write_gpt_table() - Write the GUID Partition Table to disk