[PATCH 09/16] efi/libstub: Have efi_kaslr_relocate_kernel() handle extra_size

[Jason Gunthorpe <jgg@nvidia.com>]

There are three possibilities here, each a little different:
- relocating for kaslr, just request extra_size like the zboot flow does
  when allocating from EFI
- in place, attempt to allocate from EFI precisely after Image to
  confirm nothing weird is there, failing that
- like any other error with the in place EFI location allocate
  non-randomly with extra_size included

Remove the callers' assignments of reserve_size since
efi_kaslr_relocate_kernel() now does it.

Signed-off-by: Jason Gunthorpe <jgg@nvidia.com>
---
 drivers/firmware/efi/libstub/arm64-stub.c |  1 -
 drivers/firmware/efi/libstub/kaslr.c      | 72 +++++++++++++++++++----
 drivers/firmware/efi/libstub/riscv-stub.c |  1 -
 3 files changed, 62 insertions(+), 12 deletions(-)

diff --git a/drivers/firmware/efi/libstub/arm64-stub.c b/drivers/firmware/efi/libstub/arm64-stub.c
index 1d3fc8ac13efa7..1aa169464e529e 100644
--- a/drivers/firmware/efi/libstub/arm64-stub.c
+++ b/drivers/firmware/efi/libstub/arm64-stub.c
@@ -38,7 +38,6 @@ efi_status_t handle_kernel_image(unsigned long *image_addr,
 	info = efi_get_image_info((unsigned long)_text);
 	kernel_codesize = le64_to_cpu(info->code_size);
 	kernel_memsize = kernel_size + (_end - _edata);
-	*reserve_size = kernel_memsize;
 	*image_addr = (unsigned long)_text;
 
 	return efi_kaslr_relocate_kernel(image_addr, reserve_addr, reserve_size,
diff --git a/drivers/firmware/efi/libstub/kaslr.c b/drivers/firmware/efi/libstub/kaslr.c
index 4bc963e999eb97..e419f3800af517 100644
--- a/drivers/firmware/efi/libstub/kaslr.c
+++ b/drivers/firmware/efi/libstub/kaslr.c
@@ -83,11 +83,47 @@ static bool check_image_region(u64 base, u64 size)
 	return ret;
 }
 
+/*
+ * The PE loader owns only kernel_memsize of the image.  Try to own the
+ * requested tail separately at the exact adjacent address instead of
+ * relocating the complete region.
+ */
+static efi_status_t allocate_image_tail(unsigned long image_addr,
+					unsigned long kernel_memsize,
+					unsigned long extra_size,
+					unsigned long *reserve_addr,
+					unsigned long *reserve_size)
+{
+	efi_physical_addr_t tail_addr;
+	efi_status_t status;
+
+	/*
+	 * Since we intend to use efi_free() for reserve_addr it should be
+	 * aligned to the higher alignment since efi_free() includes rounding.
+	 * For ARM64 the kernel image is already aligned up to EFI_ALLOC_ALIGN
+	 * by the linker.
+	 */
+	tail_addr = image_addr + kernel_memsize;
+	if (!IS_ALIGNED(tail_addr, EFI_ALLOC_ALIGN))
+		return EFI_OUT_OF_RESOURCES;
+
+	extra_size = round_up(extra_size, EFI_ALLOC_ALIGN);
+	status = efi_bs_call(allocate_pages, EFI_ALLOCATE_ADDRESS,
+			     EFI_LOADER_CODE, extra_size / EFI_PAGE_SIZE,
+			     &tail_addr);
+	if (status != EFI_SUCCESS)
+		return status;
+
+	*reserve_addr = tail_addr;
+	*reserve_size = extra_size;
+	return EFI_SUCCESS;
+}
+
 /**
  * efi_kaslr_relocate_kernel() - Relocate the kernel (random if KASLR enabled)
  * @image_addr: Pointer to the current kernel location
- * @reserve_addr:	Pointer to the relocated kernel location
- * @reserve_size:	Size of the relocated kernel
+ * @reserve_addr:	Pointer to any allocated memory
+ * @reserve_size:	Size that was allocated
  * @kernel_size:	Size of the text + data
  * @kernel_codesize:	Size of the text
  * @kernel_memsize:	Size of the text + data + bss
@@ -109,6 +145,9 @@ efi_status_t efi_kaslr_relocate_kernel(unsigned long *image_addr,
 {
 	efi_status_t status;
 	u64 min_kimg_align = efi_get_kimg_min_align();
+	unsigned long extra_size = efi_drtm_get_extra_size();
+
+	*reserve_size = kernel_memsize + extra_size;
 
 	if (IS_ENABLED(CONFIG_RANDOMIZE_BASE) && phys_seed != 0) {
 		/*
@@ -125,16 +164,29 @@ efi_status_t efi_kaslr_relocate_kernel(unsigned long *image_addr,
 	}
 
 	if (status != EFI_SUCCESS) {
-		if (!check_image_region(*image_addr, kernel_memsize)) {
+		bool image_region_ok =
+			check_image_region(*image_addr, kernel_memsize);
+
+		if (!image_region_ok) {
 			efi_err("FIRMWARE BUG: Image BSS overlaps adjacent EFI memory region\n");
 		} else if (IS_ALIGNED(*image_addr, min_kimg_align) &&
-			   (unsigned long)_end < EFI_ALLOC_LIMIT) {
-			/*
-			 * Just execute from wherever we were loaded by the
-			 * UEFI PE/COFF loader if the placement is suitable.
-			 */
-			*reserve_size = 0;
-			return EFI_SUCCESS;
+			   (unsigned long)_end + extra_size < EFI_ALLOC_LIMIT) {
+			if (!extra_size) {
+				/*
+				 * Just execute from wherever we were loaded by
+				 * the UEFI PE/COFF loader if the placement is
+				 * suitable.
+				 */
+				*reserve_size = 0;
+				return EFI_SUCCESS;
+			}
+
+			status = allocate_image_tail(*image_addr,
+						     kernel_memsize, extra_size,
+						     reserve_addr,
+						     reserve_size);
+			if (status == EFI_SUCCESS)
+				return EFI_SUCCESS;
 		}
 
 		status = efi_allocate_pages_aligned(*reserve_size, reserve_addr,
diff --git a/drivers/firmware/efi/libstub/riscv-stub.c b/drivers/firmware/efi/libstub/riscv-stub.c
index 725b634c517919..a88b2d275e5bb0 100644
--- a/drivers/firmware/efi/libstub/riscv-stub.c
+++ b/drivers/firmware/efi/libstub/riscv-stub.c
@@ -37,7 +37,6 @@ efi_status_t handle_kernel_image(unsigned long *image_addr,
 	kernel_codesize = __init_text_end - _start;
 	kernel_memsize = kernel_size + (_end - _edata);
 	*image_addr = (unsigned long)_start;
-	*reserve_size = kernel_memsize;
 
 	status = efi_kaslr_relocate_kernel(image_addr,
 					   reserve_addr, reserve_size,
-- 
2.43.0