[PATCH 10/16] KVM: planes: Keep current_vcpu in sync in vcpu_load()

[David Woodhouse <dwmw2@infradead.org>]

From: David Woodhouse <dwmw@amazon.co.uk>

The preempt notifiers save and restore common->current_vcpu
(kvm_sched_out()/kvm_sched_in()), but vcpu_load() does not update it.

An ioctl on a plane-N fd loads that plane's vCPU while current_vcpu
may still name the plane that ran last. A reschedule inside the ioctl
then arch-puts the stale twin, whose state is not resident, and
arch-loads it back over the state the ioctl made resident: on VMX the
ioctl's VMCS writes land in the other plane's VMCS, and the trailing
vcpu_put() unloads a non-resident vCPU, corrupting the per-pCPU
loaded_vmcs tracking. The kick paths read current_vcpu->cpu too, and
the next patch, which swaps each twin in around its update after
KVM_SET_CPUID2, relies on the loaded vCPU being the one the preempt
notifiers see.

Set common->current_vcpu in vcpu_load(), so every load path (the run
loop and vCPU ioctls on any plane's fd) agrees with what the preempt
notifiers will operate on. kvm_vcpu_select_plane() keeps its own
assignment; both stores use WRITE_ONCE() to pair with the lock-free
readers in the kick paths.

That makes twin creation a writer as well: kvm_arch_vcpu_create()
calls vcpu_load() on the new twin while its plane-0 twin may be in
KVM_RUN on another pCPU, and the shared preempt notifier would be
registered twice. Hold the shared vCPU mutex across creation of a
non-zero plane's vCPU, from before kvm_arch_vcpu_create() to the end,
including the error-path kvm_arch_vcpu_destroy(), which can
vcpu_load() too (nested_vmx_free_vcpu()). Since vcpu_put() leaves
current_vcpu naming the new twin, put back the twin that owned it once
the arch code returns, on success and on every failure path; otherwise
a failed creation (create_vcpu_fd() can be made to fail with
RLIMIT_NOFILE) frees a vCPU that the lock-free kick paths still
dereference.

Found by a local Sashiko review.

Fixes the current_vcpu tracking in "kvm: Move VCPU scheduling state to
struct kvm_vcpu_common".

Fixes: https://lore.kernel.org/all/20260608144252.351443-16-joro@8bytes.org/
Signed-off-by: David Woodhouse <dwmw@amazon.co.uk>
---
 virt/kvm/kvm_main.c | 54 ++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 48 insertions(+), 6 deletions(-)

diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index 24d2291b60a2..b6bdbe5c5133 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -174,6 +174,16 @@ void vcpu_load(struct kvm_vcpu *vcpu)
 {
 	int cpu = get_cpu();
 
+	/*
+	 * The preempt notifiers operate on common->current_vcpu, so it
+	 * must name the plane vCPU being loaded here: an ioctl taken on
+	 * another plane's fd than the one that ran last would otherwise
+	 * have a preemption arch-put the stale twin (whose state is not
+	 * resident) and arch-load it back over the state this load makes
+	 * resident.  The kick paths also read it to find the running
+	 * twin's pCPU.
+	 */
+	WRITE_ONCE(vcpu->common->current_vcpu, vcpu);
 	__this_cpu_write(kvm_running_vcpu, vcpu->common);
 	preempt_notifier_register(&vcpu->common->preempt_notifier);
 	kvm_arch_vcpu_load(vcpu, cpu);
@@ -4470,7 +4480,7 @@ static void kvm_create_vcpu_debugfs(struct kvm_vcpu *vcpu)
 static int kvm_plane_ioctl_create_vcpu(struct kvm_plane *plane, unsigned long id)
 {
 	struct kvm *kvm = plane->kvm;
-	struct kvm_vcpu *vcpu;
+	struct kvm_vcpu *vcpu, *prev = NULL;
 	int r;
 
 	guard(mutex)(&kvm->lock);
@@ -4513,9 +4523,34 @@ static int kvm_plane_ioctl_create_vcpu(struct kvm_plane *plane, unsigned long id
 
 	kvm_vcpu_init(vcpu, kvm, id);
 
+	/*
+	 * A non-zero plane shares 'common' with a plane-0 twin which may
+	 * already be running.  kvm_arch_vcpu_create() calls vcpu_load(),
+	 * which publishes the half-created vCPU through
+	 * common->current_vcpu and registers the shared preempt notifier,
+	 * so it must not run concurrently with a twin's KVM_RUN: hold the
+	 * shared vCPU mutex from here to the end of creation, including
+	 * the error-path kvm_arch_vcpu_destroy(), which can vcpu_load()
+	 * too (nested VMX teardown).  It nests inside the kvm->lock held
+	 * for all of vCPU creation, the usual kvm->lock -> vcpu mutex order.
+	 *
+	 * vcpu_put() leaves current_vcpu naming the new twin; put back
+	 * the twin that owned it, so that the lock-free kick paths never
+	 * see a vCPU that is not reachable yet, or that a failed creation
+	 * frees.
+	 */
+	if (plane->level > 0) {
+		r = mutex_lock_killable(kvm_vcpu_mutex(vcpu));
+		if (r)
+			goto vcpu_free_common;
+		prev = vcpu->common->current_vcpu;
+	}
+
 	r = kvm_arch_vcpu_create(vcpu);
+	if (plane->level > 0)
+		WRITE_ONCE(vcpu->common->current_vcpu, prev);
 	if (r)
-		goto vcpu_free_common;
+		goto unlock_vcpu_free_common;
 
 	/*
 	 * Set the vCPU's index *before* the vCPU is reachable by other tasks.
@@ -4534,9 +4569,10 @@ static int kvm_plane_ioctl_create_vcpu(struct kvm_plane *plane, unsigned long id
 	 * visible (per online_vcpus), e.g. so that KVM doesn't get tricked
 	 * into a NULL-pointer dereference because KVM thinks the _current_
 	 * vCPU doesn't exist.  As a bonus, taking vcpu->mutex ensures lockdep
-	 * knows it's taken *inside* kvm->lock.
+	 * knows it's taken *inside* kvm->lock.  A twin holds it already.
 	 */
-	kvm_vcpu_lock(vcpu);
+	if (!plane->level)
+		kvm_vcpu_lock(vcpu);
 	kvm_get_kvm(kvm);
 	r = create_vcpu_fd(vcpu);
 	if (r < 0)
@@ -4550,12 +4586,18 @@ static int kvm_plane_ioctl_create_vcpu(struct kvm_plane *plane, unsigned long id
 	return r;
 
 kvm_put_xa_erase:
-	kvm_vcpu_unlock(vcpu);
+	if (!plane->level)
+		kvm_vcpu_unlock(vcpu);
 	kvm_put_kvm_no_destroy(kvm);
 	xa_erase(&plane->vcpu_array, vcpu->vcpu_idx);
 unlock_vcpu_destroy:
 	vcpu->vcpu_idx = -1;
 	kvm_arch_vcpu_destroy(vcpu);
+	if (plane->level > 0)
+		WRITE_ONCE(vcpu->common->current_vcpu, prev);
+unlock_vcpu_free_common:
+	if (plane->level > 0)
+		kvm_vcpu_unlock(vcpu);
 vcpu_free_common:
 	kvm_vcpu_common_destroy(vcpu);
 vcpu_free:
@@ -5136,7 +5178,7 @@ struct kvm_vcpu *kvm_vcpu_select_plane(struct kvm_vcpu *vcpu)
 		ret->plane_state = RUNNABLE;
 	}
 
-	common->current_vcpu = ret;
+	WRITE_ONCE(common->current_vcpu, ret);
 
 	common->plane_switch = false;
 
-- 
2.50.1