arm64/efi: move SetVirtualAddressMap() to UEFI stub

In order to support kexec, the kernel needs to be able to deal with the
state of the UEFI firmware after SetVirtualAddressMap() has been called.
To avoid having separate code paths for non-kexec and kexec, let's move
the call to SetVirtualAddressMap() to the stub: this will guarantee us
that it will only be called once (since the stub is not executed during
kexec), and ensures that the UEFI state is identical between kexec and
normal boot.

This implies that the layout of the virtual mapping needs to be created
by the stub as well. All regions are rounded up to a naturally aligned
multiple of 64 KB (for compatibility with 64k pages kernels) and recorded
in the UEFI memory map. The kernel proper reads those values and installs
the mappings in a dedicated set of page tables that are swapped in during
UEFI Runtime Services calls.

Acked-by: Leif Lindholm <leif.lindholm@linaro.org>
Acked-by: Matt Fleming <matt.fleming@intel.com>
Tested-by: Leif Lindholm <leif.lindholm@linaro.org>
Signed-off-by: Ard Biesheuvel <ard.biesheuvel@linaro.org>
diff --git a/arch/arm64/include/asm/efi.h b/arch/arm64/include/asm/efi.h
index 7129125..effef37 100644
--- a/arch/arm64/include/asm/efi.h
+++ b/arch/arm64/include/asm/efi.h
@@ -7,28 +7,36 @@
 #ifdef CONFIG_EFI
 extern void efi_init(void);
 extern void efi_idmap_init(void);
+extern void efi_virtmap_init(void);
 #else
 #define efi_init()
 #define efi_idmap_init()
+#define efi_virtmap_init()
 #endif
 
 #define efi_call_virt(f, ...)						\
 ({									\
-	efi_##f##_t *__f = efi.systab->runtime->f;			\
+	efi_##f##_t *__f;						\
 	efi_status_t __s;						\
 									\
 	kernel_neon_begin();						\
+	efi_virtmap_load();						\
+	__f = efi.systab->runtime->f;					\
 	__s = __f(__VA_ARGS__);						\
+	efi_virtmap_unload();						\
 	kernel_neon_end();						\
 	__s;								\
 })
 
 #define __efi_call_virt(f, ...)						\
 ({									\
-	efi_##f##_t *__f = efi.systab->runtime->f;			\
+	efi_##f##_t *__f;						\
 									\
 	kernel_neon_begin();						\
+	efi_virtmap_load();						\
+	__f = efi.systab->runtime->f;					\
 	__f(__VA_ARGS__);						\
+	efi_virtmap_unload();						\
 	kernel_neon_end();						\
 })
 
@@ -46,4 +54,26 @@
 
 #define EFI_ALLOC_ALIGN		SZ_64K
 
+/*
+ * On ARM systems, virtually remapped UEFI runtime services are set up in three
+ * distinct stages:
+ * - The stub retrieves the final version of the memory map from UEFI, populates
+ *   the virt_addr fields and calls the SetVirtualAddressMap() [SVAM] runtime
+ *   service to communicate the new mapping to the firmware (Note that the new
+ *   mapping is not live at this time)
+ * - During early boot, the page tables are allocated and populated based on the
+ *   virt_addr fields in the memory map, but only if all descriptors with the
+ *   EFI_MEMORY_RUNTIME attribute have a non-zero value for virt_addr. If this
+ *   succeeds, the EFI_VIRTMAP flag is set to indicate that the virtual mappings
+ *   have been installed successfully.
+ * - During an early initcall(), the UEFI Runtime Services are enabled and the
+ *   EFI_RUNTIME_SERVICES bit set if some conditions are met, i.e., we need a
+ *   non-early mapping of the UEFI system table, and we need to have the virtmap
+ *   installed.
+ */
+#define EFI_VIRTMAP		EFI_ARCH_1
+
+void efi_virtmap_load(void);
+void efi_virtmap_unload(void);
+
 #endif /* _ASM_EFI_H */