ANDROID: KVM: arm64: Copy pvmfw into guest pages during donation from the host

When the host donates a page to a protected guest at an IPA which
coincides with the PVM firmware load address, copy-in the relevant
firmware page after unmapping it from the host but before mapping it
into the guest.

Signed-off-by: Will Deacon <will@kernel.org>
Bug: 209580772
Change-Id: I8cec813fa52938945f3122655deb785523a96ec8
Signed-off-by: Will Deacon <willdeacon@google.com>

arch/arm64/include/asm/kvm_host.h

@@ -167,6 +167,7 @@ struct kvm_protected_vm {
 	struct mutex shadow_lock;
 	struct kvm_hyp_memcache teardown_mc;
 	struct list_head pinned_pages;
+	gpa_t pvmfw_load_addr;
 };
 
 struct kvm_arch {

arch/arm64/include/asm/kvm_pkvm.h

@@ -15,6 +15,7 @@
 #define KVM_MAX_PVMS 255
 
 #define HYP_MEMBLOCK_REGIONS 128
+#define PVMFW_INVALID_LOAD_ADDR	(-1)
 
 int create_el2_shadow(struct kvm *kvm);
 

arch/arm64/kvm/hyp/include/nvhe/pkvm.h

@@ -51,6 +51,9 @@ struct kvm_shadow_vm {
 
 extern struct kvm_shadow_vm **shadow_table;
 
+extern phys_addr_t pvmfw_base;
+extern phys_addr_t pvmfw_size;
+
 int __pkvm_init_shadow(struct kvm *kvm, void *shadow_va, size_t size, void *pgd);
 int __pkvm_teardown_shadow(struct kvm *kvm);
 struct kvm_vcpu *get_shadow_vcpu(int shadow_handle, int vcpu_idx);
@@ -68,4 +71,22 @@ bool kvm_handle_pvm_hvc64(struct kvm_vcpu *vcpu, u64 *exit_code);
 
 struct kvm_vcpu *pvm_mpidr_to_vcpu(struct kvm_shadow_vm *vm, unsigned long mpidr);
 
+static inline bool pvm_has_pvmfw(struct kvm_shadow_vm *vm)
+{
+	return vm->arch.pkvm.pvmfw_load_addr != PVMFW_INVALID_LOAD_ADDR;
+}
+
+static inline bool ipa_in_pvmfw_region(struct kvm_shadow_vm *vm, u64 ipa)
+{
+	struct kvm_protected_vm *pkvm = &vm->arch.pkvm;
+
+	if (!pvm_has_pvmfw(vm))
+		return false;
+
+	return ipa - pkvm->pvmfw_load_addr < pvmfw_size;
+}
+
+int pkvm_load_pvmfw_pages(struct kvm_shadow_vm *vm, u64 ipa, phys_addr_t phys,
+			  u64 size);
+
 #endif /* __ARM64_KVM_NVHE_PKVM_H__ */

arch/arm64/kvm/hyp/nvhe/mem_protect.c

@@ -1060,10 +1060,18 @@ static int guest_complete_donation(u64 addr, const struct pkvm_mem_transition *t
 	enum kvm_pgtable_prot prot = pkvm_mkstate(KVM_PGTABLE_PROT_RWX, PKVM_PAGE_OWNED);
 	struct kvm_vcpu *vcpu = tx->completer.guest.vcpu;
 	struct kvm_shadow_vm *vm = vcpu->arch.pkvm.shadow_vm;
+	phys_addr_t phys = tx->completer.guest.phys;
 	u64 size = tx->nr_pages * PAGE_SIZE;
+	int err;
 
-	return kvm_pgtable_stage2_map(&vm->pgt, addr, size, tx->completer.guest.phys,
-				      prot, &vcpu->arch.pkvm_memcache);
+	if (tx->initiator.id == PKVM_ID_HOST && ipa_in_pvmfw_region(vm, addr)) {
+		err = pkvm_load_pvmfw_pages(vm, addr, phys, size);
+		if (err)
+			return err;
+	}
+
+	return kvm_pgtable_stage2_map(&vm->pgt, addr, size, phys, prot,
+				      &vcpu->arch.pkvm_memcache);
 }
 
 static int __guest_get_completer_addr(u64 *completer_addr, phys_addr_t phys,

arch/arm64/kvm/hyp/nvhe/pkvm.c

@@ -356,6 +356,7 @@ static int init_shadow_structs(struct kvm *kvm, struct kvm_shadow_vm *vm, int nr
 
 	vm->host_kvm = kvm;
 	vm->created_vcpus = 0;
+	vm->arch.pkvm.pvmfw_load_addr = PVMFW_INVALID_LOAD_ADDR;
 
 	for (i = 0; i < nr_vcpus; i++) {
 		struct kvm_vcpu *host_vcpu = kern_hyp_va(kvm->vcpus[i]);
@@ -653,6 +654,43 @@ err_unlock:
 	return err;
 }
 
+int pkvm_load_pvmfw_pages(struct kvm_shadow_vm *vm, u64 ipa, phys_addr_t phys,
+			  u64 size)
+{
+	struct kvm_protected_vm *pkvm = &vm->arch.pkvm;
+	u64 npages, offset = ipa - pkvm->pvmfw_load_addr;
+	void *src = hyp_phys_to_virt(pvmfw_base) + offset;
+
+	if (offset >= pvmfw_size)
+		return -EINVAL;
+
+	size = min(size, pvmfw_size - offset);
+	if (!PAGE_ALIGNED(size) || !PAGE_ALIGNED(src))
+		return -EINVAL;
+
+	npages = size >> PAGE_SHIFT;
+	while (npages--) {
+		void *dst;
+
+		dst = hyp_fixmap_map(phys);
+		if (!dst)
+			return -EINVAL;
+
+		/*
+		 * No need for cache maintenance here, as the pgtable code will
+		 * take care of this when installing the pte in the guest's
+		 * stage-2 page table.
+		 */
+		memcpy(dst, src, PAGE_SIZE);
+
+		hyp_fixmap_unmap();
+		src += PAGE_SIZE;
+		phys += PAGE_SIZE;
+	}
+
+	return 0;
+}
+
 /*
  * This function sets the registers on the vcpu to their architecturally defined
  * reset values.