ANDROID: KVM: arm64: Dedicated hyp VA space for pKVM modules

Direct branching between pKVM modules (without PLTs) requires the
modules to be within a specific range of each other. If kernel VA
allocation for modules is scattered, the hypervisor's VA space mapping
can quickly exhauts the limited space available for direct addressing:

 kern VA:
 +------------+------------+------------+---------+------------+
 | mod1 .text | mod2 .text | mod1 .data | XXXXXXX | mod2 .data |
 +------------+------------+------------+---------+------------+

 hyp VA (current):
 +--------------------------------------+--------------------- ...
 |              mod1                    |              mod2
 +--------------------------------------+--------------------- ...

To maximize inter-module symbol sharing availability, reserve a 128MiB
hyp VA region aligned with the kernel's module allocation pool. Modules
loaded within this 128MiB window will mirror their kernel VA relative
offsets in the hyp VA space.

 hyp VA (Dedicated space):
 +------------+------------+------------+---------+------------+
 |    mod1    |    mod2    |    mod1    | XXXXXXX |    mod2    |
 +------------+------------+------------+---------+------------+

Bug: 472552541
Bug: 493762423
Change-Id: I817d828aa58874464c771a91587acf5d80e975b7
Signed-off-by: Vincent Donnefort <vdonnefort@google.com>
(cherry picked from commit c42c886c1098479b7f8665bad536aab702fe00ae)

arch/arm64/include/asm/memory.h

diff --git a/arch/arm64/include/asm/memory.h b/arch/arm64/include/asm/memory.h
index f9aa425..7eef249 100644
--- a/arch/arm64/include/asm/memory.h
+++ b/arch/arm64/include/asm/memory.h
@@ -247,6 +247,11 @@ extern s64			memstart_addr;
 /* PHYS_OFFSET - the physical address of the start of memory. */
 #define PHYS_OFFSET		({ VM_BUG_ON(memstart_addr & 1); memstart_addr; })
 
+#ifdef CONFIG_EXECMEM
+extern u64 module_direct_base;
+extern u64 module_plt_base;
+#endif
+
 /* the offset between the kernel virtual and physical mappings */
 extern u64			kimage_voffset;
 

arch/arm64/kvm/pkvm.c

diff --git a/arch/arm64/kvm/pkvm.c b/arch/arm64/kvm/pkvm.c
index f11b5e2..7b1e59e 100644
--- a/arch/arm64/kvm/pkvm.c
+++ b/arch/arm64/kvm/pkvm.c
@@ -1124,27 +1124,6 @@ static struct module *pkvm_el2_mod_to_module(struct pkvm_el2_module *hyp_mod)
 	return container_of(arch, struct module, arch);
 }
 
-#ifdef CONFIG_PKVM_STACKTRACE
-unsigned long pkvm_el2_mod_kern_va(unsigned long addr)
-{
-	struct pkvm_el2_module *mod;
-
-	list_for_each_entry(mod, &pkvm_modules, node) {
-		unsigned long hyp_va = (unsigned long)mod->hyp_va;
-		size_t len = (unsigned long)mod->sections.end -
-			     (unsigned long)mod->sections.start;
-
-		if (addr >= hyp_va && addr < (hyp_va + len))
-			return (unsigned long)mod->sections.start +
-				(addr - hyp_va);
-	}
-
-	return 0;
-}
-#else
-unsigned long pkvm_el2_mod_kern_va(unsigned long addr) { return 0; }
-#endif
-
 static struct pkvm_el2_module *pkvm_el2_mod_lookup_symbol(const char *name,
 							  unsigned long *addr)
 {
@@ -1380,6 +1359,78 @@ static void pkvm_module_kmemleak(struct module *this,
 	kmemleak_scan_area(start, end - start, GFP_KERNEL);
 }
 
+#define PKVM_EL2_MOD_DIRECT_RANGE	SZ_128M
+
+static unsigned long mod_direct_kern_base;
+static unsigned long mod_direct_hyp_base;
+
+unsigned long pkvm_el2_mod_kern_va(unsigned long addr)
+{
+#ifdef CONFIG_PKVM_STACKTRACE
+	struct pkvm_el2_module *mod;
+
+	/* Fast lookup into the direct range */
+	if (mod_direct_hyp_base &&
+	    addr >= mod_direct_hyp_base &&
+	    addr < mod_direct_hyp_base + PKVM_EL2_MOD_DIRECT_RANGE)
+		return mod_direct_kern_base + (addr - mod_direct_hyp_base);
+
+	list_for_each_entry(mod, &pkvm_modules, node) {
+		unsigned long hyp_va = (unsigned long)mod->hyp_va;
+		size_t len = (unsigned long)mod->sections.end -
+			     (unsigned long)mod->sections.start;
+
+		if (addr >= hyp_va && addr < (hyp_va + len))
+			return (unsigned long)mod->sections.start +
+				(addr - hyp_va);
+	}
+#endif
+	return 0;
+}
+
+static void *pkvm_el2_mod_alloc_hyp_va(size_t size)
+{
+	struct arm_smccc_res res;
+
+	arm_smccc_1_1_hvc(KVM_HOST_SMCCC_FUNC(__pkvm_alloc_module_va), size >> PAGE_SHIFT, &res);
+	if (res.a0 != SMCCC_RET_SUCCESS || !res.a1)
+		return NULL;
+
+	return (void *)res.a1;
+}
+
+static __ref void *pkvm_el2_mod_alloc_mod_direct(unsigned long start, unsigned long end)
+{
+	if (mod_direct_hyp_base)
+		goto alloc;
+
+	mod_direct_kern_base = module_direct_base ?: module_plt_base;
+	if (WARN_ON(!mod_direct_kern_base))
+		return NULL;
+
+	mod_direct_hyp_base = (unsigned long)pkvm_el2_mod_alloc_hyp_va(PKVM_EL2_MOD_DIRECT_RANGE);
+	if (!mod_direct_hyp_base) {
+		mod_direct_kern_base = 0;
+		return NULL;
+	}
+
+alloc:
+	if (start < mod_direct_kern_base ||
+	    end > mod_direct_kern_base + PKVM_EL2_MOD_DIRECT_RANGE)
+		return NULL;
+
+	return (void *)(mod_direct_hyp_base + (start - mod_direct_kern_base));
+}
+
+static void *pkvm_el2_mod_alloc_va(unsigned long start, unsigned long end)
+{
+	void *hyp_va;
+
+	hyp_va = pkvm_el2_mod_alloc_mod_direct(start, end);
+
+	return hyp_va ?: pkvm_el2_mod_alloc_hyp_va(end - start);
+}
+
 int __pkvm_load_el2_module(struct module *this, unsigned long *token)
 {
 	struct pkvm_el2_module *mod = &this->arch.hyp;
@@ -1392,7 +1443,6 @@ int __pkvm_load_el2_module(struct module *this, unsigned long *token)
 		{ &mod->data, KVM_PGTABLE_PROT_R | KVM_PGTABLE_PROT_W },
 	};
 	void *start, *end, *hyp_va, *mod_remap;
-	struct arm_smccc_res res;
 	kvm_nvhe_reloc_t *endrel;
 	int ret, i, secs_first;
 	size_t size;
@@ -1426,14 +1476,12 @@ int __pkvm_load_el2_module(struct module *this, unsigned long *token)
 	mod->sections.start = start;
 	mod->sections.end = end;
 
-	arm_smccc_1_1_hvc(KVM_HOST_SMCCC_FUNC(__pkvm_alloc_module_va),
-			  size >> PAGE_SHIFT, &res);
-	if (res.a0 != SMCCC_RET_SUCCESS || !res.a1) {
+	hyp_va = pkvm_el2_mod_alloc_va((unsigned long)start, (unsigned long)end);
+	if (!hyp_va) {
 		kvm_err("Failed to allocate hypervisor VA space for EL2 module\n");
 		module_put(this);
-		return res.a0 == SMCCC_RET_SUCCESS ? -ENOMEM : -EPERM;
+		return -ENOMEM;
 	}
-	hyp_va = (void *)res.a1;
 	mod->hyp_va = hyp_va;
 
 	/*

arch/arm64/mm/init.c

diff --git a/arch/arm64/mm/init.c b/arch/arm64/mm/init.c
index 446060a..2c46c1d 100644
--- a/arch/arm64/mm/init.c
+++ b/arch/arm64/mm/init.c
@@ -459,8 +459,8 @@ void dump_mem_limit(void)
 }
 
 #ifdef CONFIG_EXECMEM
-static u64 module_direct_base __ro_after_init = 0;
-static u64 module_plt_base __ro_after_init = 0;
+u64 module_direct_base __ro_after_init;
+u64 module_plt_base __ro_after_init;
 
 /*
  * Choose a random page-aligned base address for a window of 'size' bytes which