w corruption could happen (I thought overwriting the page
    table), but claude notice the possibility that staging might change
    after we computed the page table, and this fix is the result. Claude
    didn't suggest a diff, but did provide many helpful clues that lead me
    to this fix.
    
    PR: 294630
    Reviewed by: kib (prior version)
    Sponsored by: Netflix
    MFC After: insta per re@ request
    Differential Revision: https://reviews.freebsd.org/D57462
    
    (cherry picked from commit 3915ffb1c3e04b26d1506bf35d3f665b2e25a915)
---
 stand/efi/loader/arch/amd64/elf64_freebsd.c | 47 ++++++++++++++++++-----------
 stand/efi/loader/bootinfo.c                 | 19 +++++++++++-
 2 files changed, 48 insertions(+), 18 deletions(-)

diff --git a/stand/efi/loader/arch/amd64/elf64_freebsd.c b/stand/efi/loader/arch/amd64/elf64_freebsd.c
index 91dd979a677e..72c8d558d8a6 100644
--- a/stand/efi/loader/arch/amd64/elf64_freebsd.c
+++ b/stand/efi/loader/arch/amd64/elf64_freebsd.c
@@ -94,7 +94,7 @@ elf64_exec(struct preloaded_file *fp)
 	Elf_Ehdr 		*ehdr;
 	vm_offset_t		modulep, kernend, trampcode, trampstack;
 	int			err, i;
-	bool			copy_auto;
+	bool			copy_auto, needs_pt4;
 
 	copy_auto = copy_staging == COPY_STAGING_AUTO;
 	if (copy_auto)
@@ -162,6 +162,7 @@ elf64_exec(struct preloaded_file *fp)
 			PT2[i] = (pd_entry_t)i * (2 * 1024 * 1024);
 			PT2[i] |= PG_V | PG_RW | PG_PS;
 		}
+		needs_pt4 = false;
 	} else {
 		PT4 = (pml4_entry_t *)0x0000000100000000; /* 4G */
 		err = BS->AllocatePages(AllocateMaxAddress, EfiLoaderData, 9,
@@ -173,7 +174,35 @@ elf64_exec(struct preloaded_file *fp)
 				copy_staging = COPY_STAGING_AUTO;
 			return (ENOMEM);
 		}
+		needs_pt4 = true;
+	}
+
+	printf("%scopying staging tramp %p PT4 %p\n",
+	    copy_staging == COPY_STAGING_ENABLE ? "" : "not ",
+	    trampoline, PT4);
+	printf("Start @ 0x%lx ...\n", ehdr->e_entry);
+
+	/*
+	 * we have to cleanup here because net_cleanup() doesn't work after
+	 * we call ExitBootServices
+	 */
+	dev_cleanup();
+
+	efi_time_fini();
+	err = bi_load(fp->f_args, &modulep, &kernend, true);
+	if (err != 0) {
+		efi_time_init();
+		if (copy_auto)
+			copy_staging = COPY_STAGING_AUTO;
+		return (err);
+	}
 
+	/*
+	 * staging might move in bi_load because we automatiaclly move when we
+	 * copy data in. At this point, staging can't move anymore, so create
+	 * PT4 with the correct value.
+	 */
+	if (needs_pt4) {
 		bzero(PT4, 9 * EFI_PAGE_SIZE);
 
 		PT3_l = &PT4[NPML4EPG * 1];
@@ -210,22 +239,6 @@ elf64_exec(struct preloaded_file *fp)
 		}
 	}
 
-	printf("staging %#lx (%scopying) tramp %p PT4 %p\n",
-	    staging, copy_staging == COPY_STAGING_ENABLE ? "" : "not ",
-	    trampoline, PT4);
-	printf("Start @ 0x%lx ...\n", ehdr->e_entry);
-
-	efi_time_fini();
-	err = bi_load(fp->f_args, &modulep, &kernend, true);
-	if (err != 0) {
-		efi_time_init();
-		if (copy_auto)
-			copy_staging = COPY_STAGING_AUTO;
-		return (err);
-	}
-
-	dev_cleanup();
-
 	trampoline(trampstack, copy_staging == COPY_STAGING_ENABLE ?
 	    efi_copy_finish : efi_copy_finish_nop, kernend, modulep,
 	    PT4, ehdr->e_entry);
diff --git a/stand/efi/loader/bootinfo.c b/stand/efi/loader/bootinfo.c
index 2961b8b97fb7..e56cd90ed7b8 100644
--- a/stand/efi/loader/bootinfo.c
+++ b/stand/efi/loader/bootinfo.c
@@ -213,6 +213,17 @@ bi_load_efi_data(struct preloaded_file *kfp, bool exit_bs)
 	}
 #endif
 
+#if defined(__amd64__) || defined(__i386__)
+	extern uint64_t staging;
+	/*
+	 * Staging can't move after this point, so report the final value before
+	 * we try to exit boot services below. The metadata added is added to
+	 * the malloced arena that we setup when we started and doesn't interact
+	 * with boot services.
+	 */
+	printf("staging %#jx\n", (uintmax_t)staging);
+#endif
+
 	do_vmap = true;
 	efi_novmap = getenv("efi_disable_vmap");
 	if (efi_novmap != NULL)
@@ -302,14 +313,20 @@ bi_load_efi_data(struct preloaded_file *kfp, bool exit_bs)
 	 * loader.conf(5). By default we will setup the virtual
 	 * map entries.
 	 */
-
 	if (do_vmap)
 		efi_do_vmap(mm, sz, dsz, mmver);
+
+	/*
+	 * Add the memory map to the metadata. addmetadata copies the data into
+	 * the malloc arena, so we can safely free the memory map pages after.
+	 * Or could if boot services was still running.
+	 */
 	efihdr->memory_size = sz;
 	efihdr->descriptor_size = dsz;
 	efihdr->descriptor_version = mmver;
 	file_addmetadata(kfp, MODINFOMD_EFI_MAP, efisz + sz,
 	    efihdr);
+	/* BS->FreePages(addr, pages); */
 
 	return (0);
 }