ANDROID: KVM: arm64: Handle hyp s1 allocation failures gracefully

Now that we have removed several indirection layers in the mem_protect
machinery, it becomes a little simpler to handle allocation failures
gracefully. Use this opportunity to handle hyp stage-1 mapping failures
without BUG-ing, hence paving the way for reducing its associated
carveout.

Bug: 357781595
Suggested-by: Vincent Donnefort <vdonnefort@google.com>
Change-Id: Ib8b27e59903f3223eaf82a915fe45ba4bcf976ba
Signed-off-by: Quentin Perret <qperret@google.com>

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

@@ -33,6 +33,7 @@ int __pkvm_create_private_mapping(phys_addr_t phys, size_t size,
 int pkvm_create_stack(phys_addr_t phys, unsigned long *haddr);
 int pkvm_alloc_private_va_range(size_t size, unsigned long *haddr);
 void pkvm_remove_mappings(void *from, void *to);
+unsigned long pkvm_remove_mappings_locked(void *from, void *to);
 phys_addr_t __pkvm_private_range_pa(void *va);
 int __hyp_allocator_map(unsigned long start, phys_addr_t phys);
 

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

@@ -1285,7 +1285,14 @@ int __pkvm_host_share_hyp(u64 pfn)
 	}
 
 	prot = pkvm_mkstate(PAGE_HYP, PKVM_PAGE_SHARED_BORROWED);
-	WARN_ON(pkvm_create_mappings_locked(virt, virt + size, prot));
+	ret = pkvm_create_mappings_locked(virt, virt + size, prot);
+	if (ret) {
+		WARN_ON(ret != -ENOMEM);
+		/* We might have failed halfway through, so remove anything we've installed */
+		pkvm_remove_mappings_locked(virt, virt + size);
+		goto unlock;
+	}
+
 	WARN_ON(__host_set_page_state_range(phys, size, PKVM_PAGE_SHARED_OWNED));
 
 unlock:
@@ -1477,7 +1484,13 @@ int __pkvm_host_donate_hyp_locked(u64 pfn, u64 nr_pages, enum kvm_pgtable_prot p
 	}
 
 	prot = pkvm_mkstate(prot, PKVM_PAGE_OWNED);
-	WARN_ON(pkvm_create_mappings_locked(virt, virt + size, prot));
+	ret = pkvm_create_mappings_locked(virt, virt + size, prot);
+	if (ret) {
+		WARN_ON(ret != -ENOMEM);
+		/* We might have failed halfway through, so remove anything we've installed */
+		pkvm_remove_mappings_locked(virt, virt + size);
+		goto unlock;
+	}
 	WARN_ON(host_stage2_set_owner_locked(phys, size, PKVM_ID_HYP));
 
 unlock:

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

@@ -178,7 +178,13 @@ void __pkvm_unmap_module_page(u64 pfn, void *va)
 
 int __hyp_allocator_map(unsigned long va, phys_addr_t phys)
 {
-	return __pkvm_create_mappings(va, PAGE_SIZE, phys, PAGE_HYP);
+	int ret = __pkvm_create_mappings(va, PAGE_SIZE, phys, PAGE_HYP);
+
+	/* Let's not confuse the hyp_alloc callers who will try to top-up pointlessly on -ENOMEM */
+	if (ret == -ENOMEM)
+		ret = -EBUSY;
+
+	return ret;
 }
 
 int pkvm_create_mappings_locked(void *from, void *to, enum kvm_pgtable_prot prot)
@@ -217,12 +223,19 @@ int pkvm_create_mappings(void *from, void *to, enum kvm_pgtable_prot prot)
 	return ret;
 }
 
+unsigned long pkvm_remove_mappings_locked(void *from, void *to)
+{
+	unsigned long size = (unsigned long)to - (unsigned long)from;
+
+	return kvm_pgtable_hyp_unmap(&pkvm_pgtable, (u64)from, size);
+}
+
 void pkvm_remove_mappings(void *from, void *to)
 {
 	unsigned long size = (unsigned long)to - (unsigned long)from;
 
 	hyp_spin_lock(&pkvm_pgd_lock);
-	WARN_ON(kvm_pgtable_hyp_unmap(&pkvm_pgtable, (u64)from, size) != size);
+	WARN_ON(pkvm_remove_mappings_locked(from, to) != size);
 	hyp_spin_unlock(&pkvm_pgd_lock);
 }
 

arch/arm64/kvm/pkvm.c

@@ -1639,8 +1639,10 @@ int pkvm_pgtable_stage2_map(struct kvm_pgtable *pgt, u64 addr, u64 size,
 	}
 
 	ret = kvm_call_hyp_nvhe(__pkvm_host_share_guest, pfn, gfn, prot, size / PAGE_SIZE);
-	if (WARN_ON(ret))
+	if (ret) {
+		WARN_ON(ret != -ENOMEM);
 		return ret;
+	}
 
 	swap(mapping, cache->mapping);
 	mapping->gfn = gfn;