4 files changed, 64 insertions, 16 deletions

diff --git a/drivers/firmware/efi/arm-init.c b/drivers/firmware/efi/arm-init.c
index c49d50e68aee..5a2df3fefccc 100644
--- a/drivers/firmware/efi/arm-init.c
+++ b/drivers/firmware/efi/arm-init.c
@@ -211,12 +211,11 @@ static __init void reserve_regions(void)
 			memblock_mark_nomap(paddr, size);
 
 	}
-
-	set_bit(EFI_MEMMAP, &efi.flags);
 }
 
 void __init efi_init(void)
 {
+	struct efi_memory_map_data data;
 	struct efi_fdt_params params;
 
 	/* Grab UEFI information placed in FDT by stub */
@@ -225,9 +224,12 @@ void __init efi_init(void)
 
 	efi_system_table = params.system_table;
 
-	efi.memmap.phys_map = params.mmap;
-	efi.memmap.map = early_memremap_ro(params.mmap, params.mmap_size);
-	if (efi.memmap.map == NULL) {
+	data.desc_version = params.desc_ver;
+	data.desc_size = params.desc_size;
+	data.size = params.mmap_size;
+	data.phys_map = params.mmap;
+
+	if (efi_memmap_init_early(&data) < 0) {
 		/*
 		* If we are booting via UEFI, the UEFI memory map is the only
 		* description of memory we have, so there is little point in
@@ -235,9 +237,6 @@ void __init efi_init(void)
 		*/
 		panic("Unable to map EFI memory map.\n");
 	}
-	efi.memmap.map_end = efi.memmap.map + params.mmap_size;
-	efi.memmap.desc_size = params.desc_size;
-	efi.memmap.desc_version = params.desc_ver;
 
 	WARN(efi.memmap.desc_version != 1,
 	     "Unexpected EFI_MEMORY_DESCRIPTOR version %ld",
@@ -248,7 +247,7 @@ void __init efi_init(void)
 
 	reserve_regions();
 	efi_memattr_init();
-	early_memunmap(efi.memmap.map, params.mmap_size);
+	efi_memmap_unmap();
 
 	memblock_reserve(params.mmap & PAGE_MASK,
 			 PAGE_ALIGN(params.mmap_size +
diff --git a/drivers/firmware/efi/arm-runtime.c b/drivers/firmware/efi/arm-runtime.c
index c394b81fe452..eedb30351a68 100644
--- a/drivers/firmware/efi/arm-runtime.c
+++ b/drivers/firmware/efi/arm-runtime.c
@@ -114,7 +114,7 @@ static int __init arm_enable_runtime_services(void)
 
 	pr_info("Remapping and enabling EFI services.\n");
 
-	mapsize = efi.memmap.map_end - efi.memmap.map;
+	mapsize = efi.memmap.desc_size * efi.memmap.nr_map;
 
 	efi.memmap.map = memremap(efi.memmap.phys_map, mapsize, MEMREMAP_WB);
 	if (!efi.memmap.map) {
diff --git a/drivers/firmware/efi/efi.c b/drivers/firmware/efi/efi.c
index 5a2631af7410..c1879999abe7 100644
--- a/drivers/firmware/efi/efi.c
+++ b/drivers/firmware/efi/efi.c
@@ -544,6 +544,52 @@ int __init efi_config_init(efi_config_table_type_t *arch_tables)
 	return ret;
 }
 
+/**
+ * efi_memmap_init_early - Map the EFI memory map data structure
+ * @data: EFI memory map data
+ *
+ * Use early_memremap() to map the passed in EFI memory map and assign
+ * it to efi.memmap.
+ */
+int __init efi_memmap_init_early(struct efi_memory_map_data *data)
+{
+	struct efi_memory_map map;
+
+	if (efi_enabled(EFI_PARAVIRT))
+		return 0;
+
+	map.phys_map = data->phys_map;
+
+	map.map = early_memremap(data->phys_map, data->size);
+	if (!map.map) {
+		pr_err("Could not map the memory map!\n");
+		return -ENOMEM;
+	}
+
+	map.nr_map = data->size / data->desc_size;
+	map.map_end = map.map + data->size;
+
+	map.desc_version = data->desc_version;
+	map.desc_size = data->desc_size;
+
+	set_bit(EFI_MEMMAP, &efi.flags);
+
+	efi.memmap = map;
+
+	return 0;
+}
+
+void __init efi_memmap_unmap(void)
+{
+	unsigned long size;
+
+	size = efi.memmap.desc_size * efi.memmap.nr_map;
+
+	early_memunmap(efi.memmap.map, size);
+	efi.memmap.map = NULL;
+	clear_bit(EFI_MEMMAP, &efi.flags);
+}
+
 #ifdef CONFIG_EFI_VARS_MODULE
 static int __init efi_load_efivars(void)
 {
diff --git a/drivers/firmware/efi/fake_mem.c b/drivers/firmware/efi/fake_mem.c
index c437388a7b85..939eec47139f 100644
--- a/drivers/firmware/efi/fake_mem.c
+++ b/drivers/firmware/efi/fake_mem.c
@@ -57,6 +57,7 @@ static int __init cmp_fake_mem(const void *x1, const void *x2)
 void __init efi_fake_memmap(void)
 {
 	u64 start, end, m_start, m_end, m_attr;
+	struct efi_memory_map_data data;
 	int new_nr_map = efi.memmap.nr_map;
 	efi_memory_desc_t *md;
 	phys_addr_t new_memmap_phy;
@@ -180,12 +181,14 @@ void __init efi_fake_memmap(void)
 	}
 
 	/* swap into new EFI memmap */
-	efi_unmap_memmap();
-	efi.memmap.map = new_memmap;
-	efi.memmap.phys_map = new_memmap_phy;
-	efi.memmap.nr_map = new_nr_map;
-	efi.memmap.map_end = efi.memmap.map + efi.memmap.nr_map * efi.memmap.desc_size;
-	set_bit(EFI_MEMMAP, &efi.flags);
+	early_memunmap(new_memmap, efi.memmap.desc_size * new_nr_map);
+	efi_memmap_unmap();
+
+	data.phys_map = new_memmap_phy;
+	data.size = efi.memmap.desc_size * new_nr_map;
+	data.desc_version = efi.memmap.desc_version;
+	data.desc_size = efi.memmap.desc_size;
+	efi_memmap_init_early(&data);
 
 	/* print new EFI memmap */
 	efi_print_memmap();