pe: Add IS_PAGE_ALIGNED macro

This makes some checks in `get_mem_attrs` and `update_mem_attrs` clearer. Also add `test-pe-util.c` with a test for the new macro. The file is named that way instead of `test-pe.c` to avoid having to get `pe.c` building in the unit test environment. Signed-off-by: Nicholas Bishop <[email protected]>
pbatard · Jan 27, 2023 · e4f40ae · e4f40ae
1 parent c7b3051
commit e4f40ae
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 2 deletions.
diff --git a/include/peimage.h b/include/peimage.h
@@ -29,6 +29,9 @@
 #define ALIGN_VALUE(Value, Alignment) ((Value) + (((Alignment) - (Value)) & ((Alignment) - 1)))
 #define ALIGN_POINTER(Pointer, Alignment) ((VOID *) (ALIGN_VALUE ((UINTN)(Pointer), (Alignment))))
 
+// Check if `val` is evenly aligned to the page size.
+#define IS_PAGE_ALIGNED(val) (!((val) & EFI_PAGE_MASK))
+
 //
 // PE32+ Subsystem type for EFI images
 //

diff --git a/pe.c b/pe.c
@@ -937,7 +937,7 @@ get_mem_attrs (uintptr_t addr, size_t size, uint64_t *attrs)
 	if (EFI_ERROR(efi_status) || !proto)
 		return efi_status;
 
-	if (physaddr & 0xfff || size & 0xfff || size == 0 || attrs == NULL) {
+	if (!IS_PAGE_ALIGNED(physaddr) || !IS_PAGE_ALIGNED(size) || size == 0 || attrs == NULL) {
 		dprint(L"%a called on 0x%llx-0x%llx and attrs 0x%llx\n",
 		       __func__, (unsigned long long)physaddr,
 		       (unsigned long long)(physaddr+size-1),
@@ -971,7 +971,7 @@ update_mem_attrs(uintptr_t addr, uint64_t size,
 		       (unsigned long long)addr, (unsigned long long)size,
 		       &before, efi_status);
 
-	if (physaddr & 0xfff || size & 0xfff || size == 0) {
+	if (!IS_PAGE_ALIGNED(physaddr) || !IS_PAGE_ALIGNED(size) || size == 0) {
 		dprint(L"%a called on 0x%llx-0x%llx (size 0x%llx) +%a%a%a -%a%a%a\n",
 		       __func__, (unsigned long long)physaddr,
 		       (unsigned long long)(physaddr + size - 1),

diff --git a/test-pe-util.c b/test-pe-util.c
@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: BSD-2-Clause-Patent
+/*
+ * test-pe-util.c - test PE utilities
+ */
+
+#ifndef SHIM_UNIT_TEST
+#define SHIM_UNIT_TEST
+#endif
+#include "shim.h"
+
+static int
+test_is_page_aligned(void)
+{
+	assert_true_return(IS_PAGE_ALIGNED(0), -1, "\n");
+	assert_false_return(IS_PAGE_ALIGNED(1), -1, "\n");
+	assert_false_return(IS_PAGE_ALIGNED(4095), -1, "\n");
+	assert_true_return(IS_PAGE_ALIGNED(4096), -1, "\n");
+	assert_false_return(IS_PAGE_ALIGNED(4097), -1, "\n");
+
+	return 0;
+}
+
+int
+main(void)
+{
+	int status = 0;
+	test(test_is_page_aligned);
+
+	return status;
+}