ANDROID: KVM: arm64: Reclaim guest page-table pages during teardown

The memory pages donated by the host to the hypervisor for the creation
of guest page-tables is currently never reclaimed. Fix this returning
those pages back to the host during guest teardown using a hyp_memcache.

Signed-off-by: Quentin Perret <qperret@google.com>
Bug: 209580772
Change-Id: I90615c34e61a11ca402c19d016bd40e3dd880637
Signed-off-by: Will Deacon <willdeacon@google.com>

arch/arm64/include/asm/kvm_host.h

@@ -159,6 +159,7 @@ struct kvm_arch_memory_slot {
 struct kvm_protected_vm {
 	bool enabled;
 	int shadow_handle;
+	struct kvm_hyp_memcache teardown_mc;
 };
 
 struct kvm_arch {

arch/arm64/kvm/arm.c

@@ -161,6 +161,7 @@ vm_fault_t kvm_arch_vcpu_fault(struct kvm_vcpu *vcpu, struct vm_fault *vmf)
 	return VM_FAULT_SIGBUS;
 }
 
+void free_hyp_memcache(struct kvm_hyp_memcache *mc);
 static void kvm_shadow_destroy(struct kvm *kvm)
 {
 	if (!kvm_vm_is_protected(kvm))
@@ -168,6 +169,8 @@ static void kvm_shadow_destroy(struct kvm *kvm)
 
 	if (kvm->arch.pkvm.shadow_handle)
 		WARN_ON(kvm_call_hyp_nvhe(__pkvm_teardown_shadow, kvm));
+
+	free_hyp_memcache(&kvm->arch.pkvm.teardown_mc);
 }
 
 /**

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

@@ -72,6 +72,7 @@ int hyp_pin_shared_mem(void *from, void *to);
 void hyp_unpin_shared_mem(void *from, void *to);
 int refill_memcache(struct kvm_hyp_memcache *mc, unsigned long min_pages,
 		    struct kvm_hyp_memcache *host_mc);
+void reclaim_guest_pages(struct kvm_shadow_vm *vm, struct kvm_hyp_memcache *mc);
 
 static __always_inline void __load_host_stage2(void)
 {

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

@@ -275,6 +275,60 @@ err:
 	return ret;
 }
 
+static int reclaim_walker(u64 addr, u64 end, u32 level, kvm_pte_t *ptep,
+		enum kvm_pgtable_walk_flags flag,
+		void * const arg)
+{
+	kvm_pte_t pte = *ptep;
+	phys_addr_t phys;
+
+	if (!kvm_pte_valid(pte))
+		return 0;
+
+	/*
+	 * Only update the host stage-2 -- we're about to tear-down the guest
+	 * stage-2 so no need to waste effort trying to keep it in sync.
+	 */
+	phys = kvm_pte_to_phys(pte);
+	BUG_ON(host_stage2_set_owner_locked(phys, PAGE_SIZE, pkvm_host_id));
+
+	/*
+	 * XXX: if protected guest mark the page 'dirty' instead, and zero it
+	 * lazily on host s2 aborts.
+	 */
+
+	return 0;
+}
+
+void reclaim_guest_pages(struct kvm_shadow_vm *vm, struct kvm_hyp_memcache *mc)
+{
+	struct kvm_pgtable_walker walker = {
+		.cb     = reclaim_walker,
+		.flags  = KVM_PGTABLE_WALK_LEAF
+	};
+	void *addr;
+
+	host_lock_component();
+	__guest_lock(vm);
+
+	/* Reclaim all guest pages, and dump all pgtable pages in the hyp_pool */
+	BUG_ON(kvm_pgtable_walk(&vm->pgt, 0, BIT(vm->pgt.ia_bits), &walker));
+	kvm_pgtable_stage2_destroy(&vm->pgt);
+	vm->arch.mmu.pgd_phys = 0ULL;
+
+	__guest_unlock(vm);
+	host_unlock_component();
+
+	/* Drain the hyp_pool into the memcache */
+	addr = hyp_alloc_pages(&vm->pool, 0);
+	while (addr) {
+		memset(hyp_virt_to_page(addr), 0, sizeof(struct hyp_page));
+		push_hyp_memcache(mc, addr, hyp_virt_to_phys);
+		WARN_ON(__pkvm_hyp_donate_host(hyp_virt_to_pfn(addr), 1));
+		addr = hyp_alloc_pages(&vm->pool, 0);
+	}
+}
+
 int __pkvm_prot_finalize(void)
 {
 	struct kvm_s2_mmu *mmu = &host_kvm.arch.mmu;

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

@@ -605,6 +605,7 @@ int __pkvm_teardown_shadow(struct kvm *kvm)
 
 	shadow_size = vm->shadow_area_size;
 
+	reclaim_guest_pages(vm, &vm->host_kvm->arch.pkvm.teardown_mc);
 	unpin_host_vcpus(vm);
 	hyp_unpin_shared_mem(vm->host_kvm, vm->host_kvm + 1);
 	remove_shadow_table(shadow_handle);