[PATCH 2/4] KVM: x86/mmu: Harden "map private PFN" against unexpected root invalidation

Move kvm_tdp_mmu_map_private_pfn()'s reload of the MMU into its tight loop
so that an unexpected root invalidation has a better chance of being
handled gracefully, even though it should be impossible for the vCPU's root
to be invalidated after the initial reload.  As is, encountering an invalid
root is *guaranteed* to put the task into an infinite loop (albeit a
breakable loop that honors NEED_RESCHED).

Add a WARN to try and detect bugs that break KVM's expectations, along with
a comment to explain why it should be impossible for the root to be
invalidated.

Cc: Kai Huang <[email protected]>
Cc: Yan Zhao <[email protected]>
Cc: Rick Edgecombe <[email protected]>
Signed-off-by: Sean Christopherson <[email protected]>
---
 arch/x86/kvm/mmu/mmu.c | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/arch/x86/kvm/mmu/mmu.c b/arch/x86/kvm/mmu/mmu.c
index 621b0a42f2a1..c6cac893cbad 100644
--- a/arch/x86/kvm/mmu/mmu.c
+++ b/arch/x86/kvm/mmu/mmu.c
@@ -5184,10 +5184,6 @@ int kvm_tdp_mmu_map_private_pfn(struct kvm_vcpu *vcpu, gfn_t gfn, kvm_pfn_t pfn)
 	if (kvm_gfn_is_write_tracked(kvm, fault.slot, fault.gfn))
 		return -EPERM;
 
-	r = kvm_mmu_reload(vcpu);
-	if (r)
-		return r;
-
 	r = mmu_topup_memory_caches(vcpu, false);
 	if (r)
 		return r;
@@ -5199,10 +5195,21 @@ int kvm_tdp_mmu_map_private_pfn(struct kvm_vcpu *vcpu, gfn_t gfn, kvm_pfn_t pfn)
 		if (kvm_test_request(KVM_REQ_VM_DEAD, vcpu))
 			return -EIO;
 
+		r = kvm_mmu_reload(vcpu);
+		if (r)
+			return r;
+
 		cond_resched();
 
 		guard(read_lock)(&kvm->mmu_lock);
 
+		/*
+		 * Because slots_lock is held, it should be impossible for roots
+		 * to be invalidated after the initial MMU reload.  WARN, but
+		 * continue and re-reload the MMU to try and keep the VM alive.
+		 */
+		WARN_ON_ONCE(kvm_test_request(KVM_REQ_MMU_FREE_OBSOLETE_ROOTS, vcpu));
+
 		r = kvm_tdp_mmu_map(vcpu, &fault);
 	} while (r == RET_PF_RETRY);
 
-- 
2.55.0.679.g6767b8d81c-goog