[PATCH 09/10] test: lib: add tests for the PE image checks

[Ana Bobillo <anabobillo3@gmail.com>]

From: Ana Bobillo <anabobillo3@gmail.com>

Add unit tests that pass malformed PE images to efi_check_pe(),
LoadImage() and efi_image_parse() and check that they are rejected:

- an image ending after a 32bit NT header;
- a section whose end wraps around 4 GiB;
- a section whose raw data runs past the end of the image;
- headers that do not fit into the loaded image;
- a relocation table past the end of the image;
- an image that cannot be measured, when TCG2 is installed;
- headers, an optional header, a certificate table entry, a section or
  a certificate table that do not fit, a file alignment that is not a
  power of two, and sections sharing their raw data.

Check that PE32 and PE32+ images with 16, 6 or 5 data directories, an
empty section, a last section that is not padded and trailing data of an
odd size still load and parse.

Each image is built at the start of a larger zeroed buffer, so that
without the checks most tests read the bytes after it and get a wrong
status instead of faulting. Without the checks the tests of the wrapped
section and of the relocation table far past the image crash, and the
test of the headers writes past the loaded image.

Signed-off-by: Ana Bobillo <anabobillo3@gmail.com>
---
 test/lib/Makefile           |   2 +-
 test/lib/efi_image_loader.c | 442 ++++++++++++++++++++++++++++++++++++
 2 files changed, 443 insertions(+), 1 deletion(-)
 create mode 100644 test/lib/efi_image_loader.c

diff --git a/test/lib/Makefile b/test/lib/Makefile
index 9fc4a1a9b5d..a67e53ca9c5 100644
--- a/test/lib/Makefile
+++ b/test/lib/Makefile
@@ -9,7 +9,7 @@ ifeq ($(CONFIG_XPL_BUILD),)
 obj-y += abuf.o
 obj-y += alist.o
 obj-$(CONFIG_ECDSA_VERIFY_MBEDTLS) += ecdsa.o
-obj-$(CONFIG_EFI_LOADER) += efi_device_path.o efi_memory.o
+obj-$(CONFIG_EFI_LOADER) += efi_device_path.o efi_image_loader.o efi_memory.o
 obj-$(CONFIG_EFI_SECURE_BOOT) += efi_image_region.o
 ifdef CONFIG_RISCV
 obj-$(CONFIG_USE_PRIVATE_LIBGCC) += test_clz.o
diff --git a/test/lib/efi_image_loader.c b/test/lib/efi_image_loader.c
new file mode 100644
index 00000000000..2cd4520a4e3
--- /dev/null
+++ b/test/lib/efi_image_loader.c
@@ -0,0 +1,442 @@
+// SPDX-License-Identifier: GPL-2.0+
+/*
+ * Test the checks of the PE image loader
+ *
+ * Copyright (c) 2026 Ana Bobillo <anabobillo3@gmail.com>
+ */
+
+#include <efi_loader.h>
+#include <pe.h>
+#include <test/lib.h>
+#include <test/test.h>
+#include <test/ut.h>
+
+#if defined(__aarch64__)
+#define TEST_MACHINE	IMAGE_FILE_MACHINE_ARM64
+#elif defined(__arm__)
+#define TEST_MACHINE	IMAGE_FILE_MACHINE_THUMB
+#elif defined(__x86_64__)
+#define TEST_MACHINE	IMAGE_FILE_MACHINE_AMD64
+#elif defined(__i386__)
+#define TEST_MACHINE	IMAGE_FILE_MACHINE_I386
+#elif defined(__riscv) && (__riscv_xlen == 64)
+#define TEST_MACHINE	IMAGE_FILE_MACHINE_RISCV64
+#elif defined(__riscv) && (__riscv_xlen == 32)
+#define TEST_MACHINE	IMAGE_FILE_MACHINE_RISCV32
+#else
+/* The loader supports no other machine: skip the tests that load images */
+#define TEST_MACHINE	0
+#endif
+
+#define TEST_FILE_ALIGN		0x200
+#define TEST_SECTION_ALIGN	0x1000
+/* Size of a valid image: headers and one section, each TEST_FILE_ALIGN */
+#define TEST_IMAGE_SIZE		(2 * TEST_FILE_ALIGN)
+
+/**
+ * struct test_image - a PE image at the start of buf
+ *
+ * @dos:	DOS header
+ * @nt:		NT headers, of a PE32 or a PE32+ image
+ * @opt32:	optional header of a PE32 image, else NULL
+ * @opt64:	optional header of a PE32+ image, else NULL
+ * @dirs:	data directories
+ * @sections:	section table
+ */
+struct test_image {
+	IMAGE_DOS_HEADER *dos;
+	IMAGE_NT_HEADERS32 *nt;
+	IMAGE_OPTIONAL_HEADER32 *opt32;
+	IMAGE_OPTIONAL_HEADER64 *opt64;
+	IMAGE_DATA_DIRECTORY *dirs;
+	IMAGE_SECTION_HEADER *sections;
+};
+
+/* Set a field of the optional header of a PE32 or a PE32+ image */
+#define opt_set(img, field, val)			\
+	do {						\
+		if ((img)->opt64)			\
+			(img)->opt64->field = (val);	\
+		else					\
+			(img)->opt32->field = (val);	\
+	} while (0)
+
+/*
+ * Images are built at the start of this buffer, which is larger than any of
+ * them. The bytes after an image are zero, so that they can be read without
+ * a fault if a size check is missing.
+ */
+static u8 buf[0x4000] __aligned(8);
+
+/**
+ * test_image_init() - set up a valid image at the start of buf
+ *
+ * The image is TEST_IMAGE_SIZE bytes long: the headers, padded to the file
+ * alignment, and one section of TEST_FILE_ALIGN bytes. Further sections are
+ * empty.
+ *
+ * @img:		image to set up
+ * @pe32plus:		true for a PE32+ image, false for a PE32 image
+ * @num_dirs:		number of data directories
+ * @num_sections:	number of section headers
+ */
+static void test_image_init(struct test_image *img, bool pe32plus,
+			    u32 num_dirs, int num_sections)
+{
+	size_t opt_size;
+
+	memset(buf, 0, sizeof(buf));
+	memset(img, 0, sizeof(*img));
+
+	img->dos = (void *)buf;
+	img->dos->e_magic = IMAGE_DOS_SIGNATURE;
+	img->dos->e_lfanew = sizeof(*img->dos);
+
+	img->nt = (void *)buf + img->dos->e_lfanew;
+	img->nt->Signature = IMAGE_NT_SIGNATURE;
+	img->nt->FileHeader.Machine = TEST_MACHINE;
+	img->nt->FileHeader.NumberOfSections = num_sections;
+
+	if (pe32plus) {
+		img->opt64 = (void *)&img->nt->OptionalHeader;
+		img->opt64->Magic = IMAGE_NT_OPTIONAL_HDR64_MAGIC;
+		img->dirs = img->opt64->DataDirectory;
+		opt_size = offsetof(IMAGE_OPTIONAL_HEADER64, DataDirectory);
+	} else {
+		img->opt32 = &img->nt->OptionalHeader;
+		img->opt32->Magic = IMAGE_NT_OPTIONAL_HDR32_MAGIC;
+		img->dirs = img->opt32->DataDirectory;
+		opt_size = offsetof(IMAGE_OPTIONAL_HEADER32, DataDirectory);
+	}
+	opt_size += num_dirs * sizeof(IMAGE_DATA_DIRECTORY);
+	img->nt->FileHeader.SizeOfOptionalHeader = opt_size;
+	img->sections = (void *)&img->nt->OptionalHeader + opt_size;
+
+	opt_set(img, AddressOfEntryPoint, TEST_SECTION_ALIGN);
+	opt_set(img, SectionAlignment, TEST_SECTION_ALIGN);
+	opt_set(img, FileAlignment, TEST_FILE_ALIGN);
+	opt_set(img, SizeOfImage, 2 * TEST_SECTION_ALIGN);
+	opt_set(img, SizeOfHeaders, TEST_FILE_ALIGN);
+	opt_set(img, Subsystem, IMAGE_SUBSYSTEM_EFI_APPLICATION);
+	opt_set(img, NumberOfRvaAndSizes, num_dirs);
+
+	memcpy(img->sections[0].Name, ".text", 5);
+	img->sections[0].Misc.VirtualSize = TEST_FILE_ALIGN;
+	img->sections[0].VirtualAddress = TEST_SECTION_ALIGN;
+	img->sections[0].SizeOfRawData = TEST_FILE_ALIGN;
+	img->sections[0].PointerToRawData = TEST_FILE_ALIGN;
+}
+
+/**
+ * load_image() - pass an image at the start of buf to LoadImage()
+ *
+ * If the image is loaded, unload it again.
+ *
+ * @size:	size of the image
+ * Return:	status code of LoadImage()
+ */
+static efi_status_t load_image(size_t size)
+{
+	efi_handle_t handle = NULL;
+	efi_status_t ret;
+
+	ret = EFI_CALL(efi_load_image(false, efi_root, NULL, buf, size,
+				      &handle));
+	if (handle)
+		EFI_CALL(efi_unload_image(handle));
+
+	return ret;
+}
+
+/**
+ * parse_image() - pass an image at the start of buf to efi_image_parse()
+ *
+ * @size:	size of the image
+ * Return:	result of efi_image_parse()
+ */
+static bool parse_image(size_t size)
+{
+	struct efi_image_regions *regs = NULL;
+	WIN_CERTIFICATE *auth = NULL;
+	size_t auth_len = 0;
+	bool ret;
+
+	ret = efi_image_parse(buf, size, &regs, &auth, &auth_len);
+	free(regs);
+
+	return ret;
+}
+
+/* The checks must let valid images through */
+static int lib_test_efi_image_valid(struct unit_test_state *uts)
+{
+	static const u32 num_dirs[] = { IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 6 };
+	struct test_image img;
+	int i, j;
+
+	if (!TEST_MACHINE)
+		return -EAGAIN;
+
+	ut_asserteq_64(EFI_SUCCESS, efi_init_obj_list());
+
+	for (i = 0; i < 2; i++) {
+		for (j = 0; j < ARRAY_SIZE(num_dirs); j++) {
+			test_image_init(&img, i, num_dirs[j], 1);
+			ut_asserteq(true, parse_image(TEST_IMAGE_SIZE));
+			ut_asserteq_64(EFI_SUCCESS, load_image(TEST_IMAGE_SIZE));
+		}
+	}
+
+	/* An empty section, and trailing data of an odd size */
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 2);
+	memcpy(img.sections[1].Name, ".bss", 4);
+	img.sections[1].Misc.VirtualSize = TEST_FILE_ALIGN;
+	img.sections[1].VirtualAddress = 2 * TEST_SECTION_ALIGN;
+	opt_set(&img, SizeOfImage, 3 * TEST_SECTION_ALIGN);
+	ut_asserteq(true, parse_image(TEST_IMAGE_SIZE + 3));
+	ut_asserteq_64(EFI_SUCCESS, load_image(TEST_IMAGE_SIZE + 3));
+
+	/* An image with five data directories has no relocation table */
+	test_image_init(&img, true, 5, 1);
+	ut_asserteq(true, parse_image(TEST_IMAGE_SIZE));
+	ut_asserteq_64(EFI_SUCCESS, load_image(TEST_IMAGE_SIZE));
+
+	/* The last section is not padded to the file alignment */
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 1);
+	img.sections[0].SizeOfRawData = TEST_FILE_ALIGN / 2;
+	ut_asserteq(true, parse_image(TEST_IMAGE_SIZE - TEST_FILE_ALIGN / 2));
+	ut_asserteq_64(EFI_SUCCESS,
+		       load_image(TEST_IMAGE_SIZE - TEST_FILE_ALIGN / 2));
+
+	return 0;
+}
+LIB_TEST(lib_test_efi_image_valid, 0);
+
+/* efi_check_pe() must check that the image holds a 64bit NT header */
+static int lib_test_efi_check_pe_nt_header(struct unit_test_state *uts)
+{
+	struct test_image img;
+	size_t size;
+
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 1);
+	ut_asserteq_64(EFI_SUCCESS, efi_check_pe(buf, TEST_IMAGE_SIZE, NULL));
+
+	/* The image ends after a 32bit NT header */
+	size = img.dos->e_lfanew + sizeof(IMAGE_NT_HEADERS32);
+	ut_asserteq_64(EFI_INVALID_PARAMETER, efi_check_pe(buf, size, NULL));
+
+	/* The NT header offset lies past the end (wraps on 32bit hosts) */
+	img.dos->e_lfanew = 0xfffffff0;
+	ut_asserteq_64(EFI_INVALID_PARAMETER,
+		       efi_check_pe(buf, TEST_IMAGE_SIZE, NULL));
+
+	return 0;
+}
+LIB_TEST(lib_test_efi_check_pe_nt_header, 0);
+
+/* efi_load_pe() must reject a section ending past 4 GiB */
+static int lib_test_efi_load_pe_section_wrap(struct unit_test_state *uts)
+{
+	struct test_image img;
+
+	if (!TEST_MACHINE)
+		return -EAGAIN;
+
+	ut_asserteq_64(EFI_SUCCESS, efi_init_obj_list());
+
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 1);
+	/* VirtualAddress + VirtualSize wraps around to TEST_SECTION_ALIGN */
+	img.sections[0].VirtualAddress = 0xfffff000;
+	img.sections[0].Misc.VirtualSize = 0x2000;
+	ut_asserteq_64(EFI_LOAD_ERROR, load_image(TEST_IMAGE_SIZE));
+
+	return 0;
+}
+LIB_TEST(lib_test_efi_load_pe_section_wrap, 0);
+
+/* efi_load_pe() must check the raw data of a section against the image */
+static int lib_test_efi_load_pe_section_raw(struct unit_test_state *uts)
+{
+	struct test_image img;
+
+	if (!TEST_MACHINE)
+		return -EAGAIN;
+
+	ut_asserteq_64(EFI_SUCCESS, efi_init_obj_list());
+
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 1);
+	/* The raw data of the section runs past the end of the image */
+	img.sections[0].SizeOfRawData = 2 * TEST_FILE_ALIGN;
+	img.sections[0].Misc.VirtualSize = 2 * TEST_FILE_ALIGN;
+	ut_asserteq_64(EFI_LOAD_ERROR, load_image(TEST_IMAGE_SIZE));
+
+	return 0;
+}
+LIB_TEST(lib_test_efi_load_pe_section_raw, 0);
+
+/* efi_load_pe() must check that the headers fit into the loaded image */
+static int lib_test_efi_load_pe_headers(struct unit_test_state *uts)
+{
+	struct test_image img;
+	int num_sections, i;
+	size_t hdr_size, size;
+
+	if (!TEST_MACHINE)
+		return -EAGAIN;
+
+	ut_asserteq_64(EFI_SUCCESS, efi_init_obj_list());
+
+	/*
+	 * Enough empty sections for the section table to end past the first
+	 * page, while the only section with data is TEST_FILE_ALIGN bytes
+	 * long at virtual address 0: the loaded image is a single page.
+	 */
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 1);
+	num_sections = (EFI_PAGE_SIZE - ((void *)img.sections - (void *)buf)) /
+		       sizeof(IMAGE_SECTION_HEADER) + 2;
+	hdr_size = ALIGN((void *)&img.sections[num_sections] - (void *)buf,
+			 TEST_FILE_ALIGN);
+	size = hdr_size + TEST_FILE_ALIGN;
+	ut_assert(size <= sizeof(buf));
+
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES,
+			num_sections);
+	opt_set(&img, SizeOfHeaders, hdr_size);
+	opt_set(&img, AddressOfEntryPoint, 0);
+	img.sections[0].VirtualAddress = 0;
+	img.sections[0].PointerToRawData = hdr_size;
+	for (i = 1; i < num_sections; i++)
+		memcpy(img.sections[i].Name, ".bss", 4);
+	ut_asserteq_64(EFI_LOAD_ERROR, load_image(size));
+
+	return 0;
+}
+LIB_TEST(lib_test_efi_load_pe_headers, 0);
+
+/* efi_loader_relocate() must check the relocation table against the image */
+static int lib_test_efi_load_pe_reloc_table(struct unit_test_state *uts)
+{
+	struct test_image img;
+	IMAGE_DATA_DIRECTORY *dir;
+
+	if (!TEST_MACHINE)
+		return -EAGAIN;
+
+	ut_asserteq_64(EFI_SUCCESS, efi_init_obj_list());
+
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 1);
+	dir = &img.dirs[IMAGE_DIRECTORY_ENTRY_BASERELOC];
+	/* The relocation table starts in the section and ends past it */
+	dir->VirtualAddress = TEST_SECTION_ALIGN;
+	dir->Size = TEST_SECTION_ALIGN;
+	ut_asserteq_64(EFI_LOAD_ERROR, load_image(TEST_IMAGE_SIZE));
+
+	/* The relocation table lies far past the end of the image */
+	dir->VirtualAddress = 0xfffff000;
+	dir->Size = sizeof(IMAGE_BASE_RELOCATION);
+	ut_asserteq_64(EFI_LOAD_ERROR, load_image(TEST_IMAGE_SIZE));
+
+	return 0;
+}
+LIB_TEST(lib_test_efi_load_pe_reloc_table, 0);
+
+/* efi_load_pe() must not load an image that cannot be measured */
+static int lib_test_efi_load_pe_measure(struct unit_test_state *uts)
+{
+	const efi_guid_t tcg2_guid = EFI_TCG2_PROTOCOL_GUID;
+	struct efi_handler *handler;
+	struct test_image img;
+
+	if (!TEST_MACHINE)
+		return -EAGAIN;
+
+	ut_asserteq_64(EFI_SUCCESS, efi_init_obj_list());
+	if (efi_search_protocol(efi_root, &tcg2_guid, &handler) != EFI_SUCCESS)
+		return -EAGAIN;
+
+	/*
+	 * LoadImage() takes a file alignment that is not a power of two, but
+	 * the image cannot be hashed for its measurement.
+	 */
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 1);
+	opt_set(&img, FileAlignment, 3 * TEST_FILE_ALIGN / 2);
+	ut_asserteq_64(EFI_ACCESS_DENIED, load_image(TEST_IMAGE_SIZE));
+
+	return 0;
+}
+LIB_TEST(lib_test_efi_load_pe_measure, 0);
+
+/* efi_image_parse() must check the PE headers against the image */
+static int lib_test_efi_image_parse_headers(struct unit_test_state *uts)
+{
+	struct test_image img;
+
+	/* The headers run past the end of the image */
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 1);
+	opt_set(&img, SizeOfHeaders, 2 * TEST_IMAGE_SIZE);
+	ut_asserteq(false, parse_image(TEST_IMAGE_SIZE));
+
+	/* The optional header is shorter than its fixed fields */
+	test_image_init(&img, true, 0, 1);
+	img.nt->FileHeader.SizeOfOptionalHeader = 0x40;
+	ut_asserteq(false, parse_image(TEST_IMAGE_SIZE));
+
+	/*
+	 * The certificate table entry lies past the optional header, in the
+	 * section table: clear the section name read as the entry.
+	 */
+	test_image_init(&img, false, IMAGE_DIRECTORY_ENTRY_SECURITY, 1);
+	opt_set(&img, NumberOfRvaAndSizes, IMAGE_NUMBEROF_DIRECTORY_ENTRIES);
+	memset(img.sections[0].Name, 0, sizeof(img.sections[0].Name));
+	ut_asserteq(false, parse_image(TEST_IMAGE_SIZE));
+
+	return 0;
+}
+LIB_TEST(lib_test_efi_image_parse_headers, 0);
+
+/* efi_image_parse() must check each section against the image */
+static int lib_test_efi_image_parse_section(struct unit_test_state *uts)
+{
+	struct test_image img;
+
+	/* The raw data of the section runs past the end of the image */
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 1);
+	img.sections[0].SizeOfRawData = 2 * TEST_FILE_ALIGN;
+	ut_asserteq(false, parse_image(TEST_IMAGE_SIZE));
+
+	/* The raw data of the section starts past the end of the image */
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 1);
+	img.sections[0].PointerToRawData = 0xfffffe00;
+	ut_asserteq(false, parse_image(TEST_IMAGE_SIZE));
+
+	/* The file alignment is not a power of two */
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 1);
+	opt_set(&img, FileAlignment, 3 * TEST_FILE_ALIGN / 2);
+	ut_asserteq(false, parse_image(2 * TEST_IMAGE_SIZE));
+
+	/* Two sections share their raw data */
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 2);
+	img.sections[1] = img.sections[0];
+	img.sections[1].VirtualAddress = 2 * TEST_SECTION_ALIGN;
+	ut_asserteq(false, parse_image(TEST_IMAGE_SIZE));
+
+	return 0;
+}
+LIB_TEST(lib_test_efi_image_parse_section, 0);
+
+/* efi_image_parse() must check the certificate table against the image */
+static int lib_test_efi_image_parse_cert_table(struct unit_test_state *uts)
+{
+	struct test_image img;
+	IMAGE_DATA_DIRECTORY *dir;
+
+	test_image_init(&img, true, IMAGE_NUMBEROF_DIRECTORY_ENTRIES, 1);
+	dir = &img.dirs[IMAGE_DIRECTORY_ENTRY_SECURITY];
+	/* The end of the certificate table wraps around to 0x10 */
+	dir->VirtualAddress = 0xfffffff0;
+	dir->Size = 0x20;
+	ut_asserteq(false, parse_image(TEST_IMAGE_SIZE));
+
+	return 0;
+}
+LIB_TEST(lib_test_efi_image_parse_cert_table, 0);