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

[Sean Christopherson <seanjc@google.com>]

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.

Note, the loop in question doesn't actually check for a stale page fault,
i.e. likely won't detect an invalid loop in the first place.  That bug will
be addressed shortly.

Cc: Kai Huang <kai.huang@intel.com>
Cc: Yan Zhao <yan.y.zhao@intel.com>
Cc: Rick Edgecombe <rick.p.edgecombe@intel.com>
Signed-off-by: Sean Christopherson <seanjc@google.com>
---
 arch/x86/kvm/mmu/mmu.c | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/arch/x86/kvm/mmu/mmu.c b/arch/x86/kvm/mmu/mmu.c
index 3220f05387b5..1969c26861e5 100644
--- a/arch/x86/kvm/mmu/mmu.c
+++ b/arch/x86/kvm/mmu/mmu.c
@@ -5209,10 +5209,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;
@@ -5224,10 +5220,22 @@ 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 *any*
+		 * roots to be invalidated after the initial MMU reload.  WARN,
+		 * but continue on; the above MMU reload will do the right thing
+		 * if the current root is actually invalid.
+		 */
+		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.860.g4b6b3295ed-goog