refactor(partition): get GPT header location from MBR

GPT header is located in first LBA after MBR entry and mbr header has
details of beginning of first entry, so use mbr header entry first_lba
data to locate GPT header rather than GPT_HEADER_OFFSET.

GPT header size is available in gpt_header, so use that
rather than using DEFAULT_GPT_HEADER_SIZE.

The location of GPT entries is available once we parse gpt_header
and is available as partitiona_lba use that to load gpt_entries rather
than GPT_ENTRY_OFFSET.

Change-Id: I3c11f8cc9d4b0b1778a37fe342fb845ea4a4eff1
Signed-off-by: Govindraj Raja <govindraj.raja@arm.com>
diff --git a/drivers/partition/partition.c b/drivers/partition/partition.c
index 1881c91..7b65511 100644
--- a/drivers/partition/partition.c
+++ b/drivers/partition/partition.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2022, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2016-2023, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -49,8 +49,8 @@
 static int load_mbr_header(uintptr_t image_handle, mbr_entry_t *mbr_entry)
 {
 	size_t bytes_read;
-	uintptr_t offset;
 	int result;
+	mbr_entry_t *tmp;
 
 	assert(mbr_entry != NULL);
 	/* MBR partition table is in LBA0. */
@@ -61,7 +61,7 @@
 	}
 	result = io_read(image_handle, (uintptr_t)&mbr_sector,
 			 PLAT_PARTITION_BLOCK_SIZE, &bytes_read);
-	if (result != 0) {
+	if ((result != 0) || (bytes_read != PLAT_PARTITION_BLOCK_SIZE)) {
 		WARN("Failed to read data (%i)\n", result);
 		return result;
 	}
@@ -71,8 +71,20 @@
 	    (mbr_sector[LEGACY_PARTITION_BLOCK_SIZE - 1] != MBR_SIGNATURE_SECOND)) {
 		return -ENOENT;
 	}
-	offset = (uintptr_t)&mbr_sector + MBR_PRIMARY_ENTRY_OFFSET;
-	memcpy(mbr_entry, (void *)offset, sizeof(mbr_entry_t));
+
+	tmp = (mbr_entry_t *)(&mbr_sector[MBR_PRIMARY_ENTRY_OFFSET]);
+
+	if (tmp->first_lba != 1) {
+		WARN("MBR header may have an invalid first LBA\n");
+		return -EINVAL;
+	}
+
+	if ((tmp->sector_nums == 0) || (tmp->sector_nums == UINT32_MAX)) {
+		WARN("MBR header entry has an invalid number of sectors\n");
+		return -EINVAL;
+	}
+
+	memcpy(mbr_entry, tmp, sizeof(mbr_entry_t));
 	return 0;
 }
 
@@ -80,14 +92,15 @@
  * Load GPT header and check the GPT signature and header CRC.
  * If partition numbers could be found, check & update it.
  */
-static int load_gpt_header(uintptr_t image_handle)
+static int load_gpt_header(uintptr_t image_handle, size_t header_offset,
+			   unsigned long long *part_lba)
 {
 	gpt_header_t header;
 	size_t bytes_read;
 	int result;
 	uint32_t header_crc, calc_crc;
 
-	result = io_seek(image_handle, IO_SEEK_SET, GPT_HEADER_OFFSET);
+	result = io_seek(image_handle, IO_SEEK_SET, header_offset);
 	if (result != 0) {
 		return result;
 	}
@@ -109,7 +122,7 @@
 	header_crc = header.header_crc;
 	header.header_crc = 0U;
 
-	calc_crc = tf_crc32(0U, (uint8_t *)&header, DEFAULT_GPT_HEADER_SIZE);
+	calc_crc = tf_crc32(0U, (uint8_t *)&header, sizeof(gpt_header_t));
 	if (header_crc != calc_crc) {
 		ERROR("Invalid GPT Header CRC: Expected 0x%x but got 0x%x.\n",
 		      header_crc, calc_crc);
@@ -123,11 +136,13 @@
 	if (list.entry_count > PLAT_PARTITION_MAX_ENTRIES) {
 		list.entry_count = PLAT_PARTITION_MAX_ENTRIES;
 	}
+
+	*part_lba = header.part_lba;
 	return 0;
 }
 
 static int load_mbr_entry(uintptr_t image_handle, mbr_entry_t *mbr_entry,
-			int part_number)
+			  int part_number)
 {
 	size_t bytes_read;
 	uintptr_t offset;
@@ -221,6 +236,8 @@
 	uintptr_t dev_handle, image_handle, image_spec = 0;
 	mbr_entry_t mbr_entry;
 	int result;
+	size_t gpt_header_offset;
+	unsigned long long part_lba = 0, gpt_entry_offset = 0;
 
 	result = plat_get_image_source(image_id, &dev_handle, &image_spec);
 	if (result != 0) {
@@ -241,10 +258,22 @@
 		return result;
 	}
 	if (mbr_entry.type == PARTITION_TYPE_GPT) {
-		result = load_gpt_header(image_handle);
-		assert(result == 0);
-		result = io_seek(image_handle, IO_SEEK_SET, GPT_ENTRY_OFFSET);
-		assert(result == 0);
+		/* Try to load GPT header from LBA-1 */
+		gpt_header_offset = mbr_entry.first_lba * PLAT_PARTITION_BLOCK_SIZE;
+		result = load_gpt_header(image_handle, gpt_header_offset, &part_lba);
+		if (result != 0 || part_lba == 0) {
+			WARN("Failed to retrieve Primary GPT header\n");
+			return result;
+		}
+
+		gpt_entry_offset = part_lba * PLAT_PARTITION_BLOCK_SIZE;
+		result = io_seek(image_handle, IO_SEEK_SET, gpt_entry_offset);
+		if (result != 0) {
+			WARN("Failed to seek (%i),"
+			     "Failed loading GPT partition table entries\n",
+			     result);
+			return result;
+		}
 		result = verify_partition_gpt(image_handle);
 	} else {
 		result = load_mbr_entries(image_handle);