[PATCH v2 2/3] x86/boot: Use separate API for memory acceptance in the EFI stub

[Ard Biesheuvel <ardb+git@google.com>]

From: Ard Biesheuvel <ardb@kernel.org>

The EFI stub runs in the execution context of the firmware, and even if
the memory acceptance that occurs when constructing the E820 memory map
after ExitBootServices(), it is still too early for SEV-SNP based memory
acceptance to use the shared GHCB page, as this requires manipulating
the firmware's live page tables.

So implement the unaccepted memory processing locally in the EFI stub.
The SEV-SNP handling will be updated in a subsequent patch.

Signed-off-by: Ard Biesheuvel <ardb@kernel.org>
---
 arch/x86/boot/compressed/mem.c                   |  2 +-
 drivers/firmware/efi/libstub/efistub.h           |  3 +-
 drivers/firmware/efi/libstub/unaccepted_memory.c | 12 +++---
 drivers/firmware/efi/libstub/x86-stub.c          | 39 ++++++++++++++++++++
 4 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/arch/x86/boot/compressed/mem.c b/arch/x86/boot/compressed/mem.c
index 6a888b80669e..caa3640fc2a0 100644
--- a/arch/x86/boot/compressed/mem.c
+++ b/arch/x86/boot/compressed/mem.c
@@ -32,7 +32,7 @@ static bool early_is_tdx_guest(void)
 	return is_tdx;
 }
 
-void arch_accept_memory(phys_addr_t start, phys_addr_t end)
+static void arch_accept_memory(phys_addr_t start, phys_addr_t end)
 {
 	/* Platform-specific memory-acceptance call goes here */
 	if (early_is_tdx_guest()) {
diff --git a/drivers/firmware/efi/libstub/efistub.h b/drivers/firmware/efi/libstub/efistub.h
index f5ba032863a9..88751d83f584 100644
--- a/drivers/firmware/efi/libstub/efistub.h
+++ b/drivers/firmware/efi/libstub/efistub.h
@@ -1231,8 +1231,7 @@ efi_zboot_entry(efi_handle_t handle, efi_system_table_t *systab);
 efi_status_t allocate_unaccepted_bitmap(__u32 nr_desc,
 					struct efi_boot_memmap *map);
 void process_unaccepted_memory(u64 start, u64 end);
-void accept_memory(phys_addr_t start, unsigned long size);
-void arch_accept_memory(phys_addr_t start, phys_addr_t end);
+void efistub_accept_memory(phys_addr_t start, phys_addr_t end);
 
 efi_status_t efi_zboot_decompress_init(unsigned long *alloc_size);
 efi_status_t efi_zboot_decompress(u8 *out, unsigned long outlen);
diff --git a/drivers/firmware/efi/libstub/unaccepted_memory.c b/drivers/firmware/efi/libstub/unaccepted_memory.c
index 02040bd6a330..4e90077ca797 100644
--- a/drivers/firmware/efi/libstub/unaccepted_memory.c
+++ b/drivers/firmware/efi/libstub/unaccepted_memory.c
@@ -118,7 +118,7 @@ void process_unaccepted_memory(u64 start, u64 end)
 	 * immediately accepted in its entirety.
 	 */
 	if (end - start < 2 * unit_size) {
-		arch_accept_memory(start, end);
+		efistub_accept_memory(start, end);
 		return;
 	}
 
@@ -129,13 +129,13 @@ void process_unaccepted_memory(u64 start, u64 end)
 
 	/* Immediately accept a <unit_size piece at the start: */
 	if (start & unit_mask) {
-		arch_accept_memory(start, round_up(start, unit_size));
+		efistub_accept_memory(start, round_up(start, unit_size));
 		start = round_up(start, unit_size);
 	}
 
 	/* Immediately accept a <unit_size piece at the end: */
 	if (end & unit_mask) {
-		arch_accept_memory(round_down(end, unit_size), end);
+		efistub_accept_memory(round_down(end, unit_size), end);
 		end = round_down(end, unit_size);
 	}
 
@@ -144,8 +144,8 @@ void process_unaccepted_memory(u64 start, u64 end)
 	 * into the bitmap.
 	 */
 	if (start < unaccepted_table->phys_base) {
-		arch_accept_memory(start,
-				   min(unaccepted_table->phys_base, end));
+		efistub_accept_memory(start,
+				      min(unaccepted_table->phys_base, end));
 		start = unaccepted_table->phys_base;
 	}
 
@@ -165,7 +165,7 @@ void process_unaccepted_memory(u64 start, u64 end)
 			     unaccepted_table->phys_base;
 		phys_end = end + unaccepted_table->phys_base;
 
-		arch_accept_memory(phys_start, phys_end);
+		efistub_accept_memory(phys_start, phys_end);
 		end = bitmap_size * unit_size * BITS_PER_BYTE;
 	}
 
diff --git a/drivers/firmware/efi/libstub/x86-stub.c b/drivers/firmware/efi/libstub/x86-stub.c
index cafc90d4caaf..7d9cf473f4d0 100644
--- a/drivers/firmware/efi/libstub/x86-stub.c
+++ b/drivers/firmware/efi/libstub/x86-stub.c
@@ -17,6 +17,7 @@
 #include <asm/boot.h>
 #include <asm/kaslr.h>
 #include <asm/sev.h>
+#include <asm/shared/tdx.h>
 
 #include "efistub.h"
 #include "x86-stub.h"
@@ -364,6 +365,44 @@ static void setup_unaccepted_memory(void)
 		efi_err("Memory acceptance protocol failed\n");
 }
 
+#ifdef CONFIG_UNACCEPTED_MEMORY
+
+static bool efistub_is_tdx_guest(void)
+{
+	u32 eax = TDX_CPUID_LEAF_ID, sig[3] = {};
+
+	if (!IS_ENABLED(CONFIG_INTEL_TDX_GUEST))
+		return false;
+
+	native_cpuid(&eax, &sig[0], &sig[2], &sig[1]);
+	return !memcmp(TDX_IDENT, sig, sizeof(sig));
+}
+
+static bool efistub_is_sevsnp_guest(void)
+{
+	return sev_get_status() & MSR_AMD64_SEV_SNP_ENABLED;
+}
+
+void efistub_accept_memory(phys_addr_t start, phys_addr_t end)
+{
+	static bool once, is_tdx, is_sevsnp;
+
+	if (!once) {
+		if (efistub_is_tdx_guest())
+			is_tdx = true;
+		else if (efistub_is_sevsnp_guest())
+			is_sevsnp = true;
+		once = true;
+	}
+
+	if (is_tdx)
+		tdx_accept_memory(start, end);
+	else if (is_sevsnp)
+		snp_accept_memory(start, end);
+}
+
+#endif
+
 static efi_char16_t *efistub_fw_vendor(void)
 {
 	unsigned long vendor = efi_table_attr(efi_system_table, fw_vendor);
-- 
2.49.0.504.g3bcea36a83-goog