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>

3 files changed, 45 insertions, 22 deletions

diff --git a/drivers/partition/partition.c b/drivers/partition/partition.c
index 1881c91b3..7b6551111 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 void dump_entries(int num)
 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 @@ static int load_mbr_header(uintptr_t image_handle, mbr_entry_t *mbr_entry)
 	}
 	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 @@ static int load_mbr_header(uintptr_t image_handle, mbr_entry_t *mbr_entry)
 	    (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 @@ static int load_mbr_header(uintptr_t image_handle, mbr_entry_t *mbr_entry)
  * 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 @@ static int load_gpt_header(uintptr_t image_handle)
 	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 @@ static int load_gpt_header(uintptr_t image_handle)
 	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 @@ int load_partition_table(unsigned int image_id)
 	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 @@ int load_partition_table(unsigned int image_id)
 		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);
diff --git a/include/drivers/partition/gpt.h b/include/drivers/partition/gpt.h
index c2a229e33..383c17de6 100644
--- a/include/drivers/partition/gpt.h
+++ b/include/drivers/partition/gpt.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2016-2023, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -12,10 +12,6 @@
 #include <tools_share/uuid.h>
 
 #define PARTITION_TYPE_GPT		0xee
-#define GPT_HEADER_OFFSET		PLAT_PARTITION_BLOCK_SIZE
-#define GPT_ENTRY_OFFSET		(GPT_HEADER_OFFSET +		\
-					 PLAT_PARTITION_BLOCK_SIZE)
-
 #define GPT_SIGNATURE			"EFI PART"
 
 typedef struct gpt_entry {
@@ -45,7 +41,7 @@ typedef struct gpt_header {
 	/* size of a single partition entry (usually 128) */
 	unsigned int		part_size;
 	unsigned int		part_crc;
-} gpt_header_t;
+} __packed gpt_header_t;
 
 int parse_gpt_entry(gpt_entry_t *gpt_entry, partition_entry_t *entry);
 
diff --git a/include/drivers/partition/partition.h b/include/drivers/partition/partition.h
index 6cb59c391..ab48e121e 100644
--- a/include/drivers/partition/partition.h
+++ b/include/drivers/partition/partition.h
@@ -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
  */
@@ -29,8 +29,6 @@ CASSERT((PLAT_PARTITION_BLOCK_SIZE == 512) ||
 
 #define LEGACY_PARTITION_BLOCK_SIZE	512
 
-#define DEFAULT_GPT_HEADER_SIZE 	92
-
 typedef struct partition_entry {
 	uint64_t		start;
 	uint64_t		length;