2 files changed, 71 insertions, 3 deletions

diff --git a/lib/efi_selftest/efi_selftest_debug_support.c b/lib/efi_selftest/efi_selftest_debug_support.c
index 9ca8b3f82f5..ac77caaf561 100644
--- a/lib/efi_selftest/efi_selftest_debug_support.c
+++ b/lib/efi_selftest/efi_selftest_debug_support.c
@@ -9,6 +9,7 @@
 
 #include <efi_loader.h>
 #include <efi_selftest.h>
+#include <linux/sizes.h>
 
 /**
  * efi_st_debug_support_execute() - execute test
@@ -21,6 +22,12 @@ static int efi_st_debug_support_execute(void)
 {
 	struct efi_debug_image_info_table_header *efi_st_debug_info_table_header = NULL;
 	efi_guid_t efi_debug_image_info_table_guid = EFI_DEBUG_IMAGE_INFO_TABLE_GUID;
+	struct efi_mem_desc *memory_map;
+	efi_uintn_t map_size = 0;
+	efi_uintn_t map_key;
+	efi_uintn_t desc_size;
+	u32 desc_version;
+	efi_status_t ret;
 
 	/* get EFI_DEBUG_IMAGE_INFO_TABLE */
 	efi_st_debug_info_table_header = efi_st_get_config_table(&efi_debug_image_info_table_guid);
@@ -30,7 +37,66 @@ static int efi_st_debug_support_execute(void)
 		return EFI_ST_FAILURE;
 	}
 
-	return EFI_ST_SUCCESS;
+	/* Load memory map */
+	ret = st_boottime->get_memory_map(&map_size, NULL, &map_key, &desc_size,
+					  &desc_version);
+	if (ret != EFI_BUFFER_TOO_SMALL) {
+		efi_st_error
+			("GetMemoryMap did not return EFI_BUFFER_TOO_SMALL\n");
+		return EFI_ST_FAILURE;
+	}
+	/* Allocate extra space for newly allocated memory */
+	map_size += sizeof(struct efi_mem_desc);
+	ret = st_boottime->allocate_pool(EFI_BOOT_SERVICES_DATA, map_size,
+					 (void **)&memory_map);
+	if (ret != EFI_SUCCESS) {
+		efi_st_error("AllocatePool failed\n");
+		return EFI_ST_FAILURE;
+	}
+	ret = st_boottime->get_memory_map(&map_size, memory_map, &map_key,
+					  &desc_size, &desc_version);
+	if (ret != EFI_SUCCESS) {
+		efi_st_error("GetMemoryMap failed\n");
+		return EFI_ST_FAILURE;
+	}
+	/* Find the system table pointer */
+	for (efi_uintn_t i = 0; map_size; ++i, map_size -= desc_size) {
+		struct efi_mem_desc *entry = &memory_map[i];
+		u64 end;
+
+		if (entry->type != EFI_RUNTIME_SERVICES_DATA)
+			continue;
+
+		end = entry->physical_start +
+		      (entry->num_pages << EFI_PAGE_SHIFT);
+		for (u64 pos = ALIGN(entry->physical_start, SZ_4M);
+		     pos <= end; pos += SZ_4M) {
+			struct efi_system_table_pointer *systab_pointer =
+				(void *)(uintptr_t)pos;
+
+			/* check for overflow */
+			if (pos < entry->physical_start)
+				break;
+			if (systab_pointer->signature ==
+			    EFI_SYSTEM_TABLE_SIGNATURE) {
+				if (systab_pointer->efi_system_table_base !=
+				    (uintptr_t)st_systable) {
+					efi_st_error("Wrong system table address\n");
+					ret = EFI_ST_FAILURE;
+					goto out;
+				}
+				ret = EFI_ST_SUCCESS;
+				goto out;
+			}
+		}
+	}
+	efi_st_error("System table pointer not found\n");
+	ret = EFI_ST_FAILURE;
+
+out:
+	st_boottime->free_pool(memory_map);
+
+	return ret;
 }
 
 EFI_UNIT_TEST(debug_support) = {
diff --git a/lib/efi_selftest/efi_selftest_esrt.c b/lib/efi_selftest/efi_selftest_esrt.c
index b7688deb496..7eadac90fbc 100644
--- a/lib/efi_selftest/efi_selftest_esrt.c
+++ b/lib/efi_selftest/efi_selftest_esrt.c
@@ -69,10 +69,12 @@ EFIAPI efi_test_fmp_get_image_info(struct efi_firmware_management_protocol *this
 	if (package_version_name)
 		*package_version_name = NULL;
 
-	if (*image_info_size < sizeof(*image_info)) {
-		*image_info_size = *descriptor_size * *descriptor_count;
+	if (*image_info_size < sizeof(*image_info) * TEST_ESRT_NUM_ENTRIES) {
+		*image_info_size = sizeof(*image_info) * TEST_ESRT_NUM_ENTRIES;
 		return EFI_BUFFER_TOO_SMALL;
 	}
+	if (!image_info)
+		return EFI_INVALID_PARAMETER;
 
 	for (int idx = 0; idx < TEST_ESRT_NUM_ENTRIES; idx++)
 		image_info[idx] = static_img_info[idx];