[PATCH v2] KVM: arm64: Ensure TLBs and I-cache are private to each vCPU in pKVM

[Fuad Tabba <fuad.tabba@linux.dev>]

All vCPUs of a VM share its VMID, so KVM flushes a physical CPU's TLBs
and I-cache for the VM whenever a different vCPU of that VM is loaded
there. pKVM guests, protected or not, run on the hypervisor's stage-2
and VMID, and the host doesn't flush for them. Nothing at EL2 flushes
instead, so a vCPU can hit TLB and I-cache entries that another vCPU of
the same VM left on that CPU.

Track at EL2 the last vCPU of each VM loaded on each CPU, and flush
when it changes. EL2 can't allocate per-CPU memory for each VM, as the
host does, so use one per-CPU array indexed by the VM's table entry.
Reset the entry's slot on every CPU when a VM reserves it, so that a new
VM starts with no last vCPU, whether the entry is unused or was left by
a VM that was torn down.

Fixes: fce886a602073 ("KVM: arm64: Plumb the pKVM MMU in KVM")
Signed-off-by: Fuad Tabba <fuad.tabba@linux.dev>
---
Changes since v1 [1]:
  - Dropped patch 1, which added a CPU number at EL2. The last vCPU is
    now tracked in a per-CPU array indexed by VM table entry. (Marc)
  - Moved the check from __pkvm_vcpu_load() to pkvm_load_hyp_vcpu().
  - The array is unsigned int, initialised to UINT_MAX. (Vincent)
  - Dropped Joey's Reviewed-by, since the patch changed.

Based on kvmarm/next (dd0bc00e3e920).

[1] https://lore.kernel.org/all/20261014125510.2928306-1-fuad.tabba@linux.dev/

 arch/arm64/kvm/hyp/nvhe/pkvm.c | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/arch/arm64/kvm/hyp/nvhe/pkvm.c b/arch/arm64/kvm/hyp/nvhe/pkvm.c
index ac020d0501d2e..1fc59229c8e86 100644
--- a/arch/arm64/kvm/hyp/nvhe/pkvm.c
+++ b/arch/arm64/kvm/hyp/nvhe/pkvm.c
@@ -38,6 +38,9 @@ unsigned int kvm_host_sve_max_vl;
  */
 static DEFINE_PER_CPU(struct pkvm_hyp_vcpu *, loaded_hyp_vcpu);
 
+/* Per VM table entry, the vcpu_idx of the last vCPU of that VM loaded on this CPU. */
+static DEFINE_PER_CPU(unsigned int, last_vcpu_ran[KVM_MAX_PVMS]);
+
 static void pkvm_vcpu_reset_hcr(struct kvm_vcpu *vcpu)
 {
 	vcpu->arch.hcr_el2 = HCR_GUEST_FLAGS;
@@ -261,6 +264,7 @@ struct pkvm_hyp_vcpu *pkvm_load_hyp_vcpu(pkvm_handle_t handle,
 {
 	struct pkvm_hyp_vcpu *hyp_vcpu = NULL;
 	struct pkvm_hyp_vm *hyp_vm;
+	unsigned int *last_ran;
 
 	/* Cannot load a new vcpu without putting the old one first. */
 	if (__this_cpu_read(loaded_hyp_vcpu))
@@ -290,8 +294,18 @@ struct pkvm_hyp_vcpu *pkvm_load_hyp_vcpu(pkvm_handle_t handle,
 unlock:
 	hyp_spin_unlock(&vm_table_lock);
 
-	if (hyp_vcpu)
-		__this_cpu_write(loaded_hyp_vcpu, hyp_vcpu);
+	if (!hyp_vcpu)
+		return NULL;
+
+	__this_cpu_write(loaded_hyp_vcpu, hyp_vcpu);
+
+	/* TLBs and I-cache are private to each vCPU, see vcpu_prepare_mmu(). */
+	last_ran = this_cpu_ptr(&last_vcpu_ran[vm_handle_to_idx(handle)]);
+	if (*last_ran != vcpu_idx) {
+		__kvm_flush_cpu_context(hyp_vcpu->vcpu.arch.hw_mmu);
+		*last_ran = vcpu_idx;
+	}
+
 	return hyp_vcpu;
 }
 
@@ -652,6 +666,7 @@ static int find_free_vm_table_entry(void)
  */
 static int allocate_vm_table_entry(void)
 {
+	unsigned int cpu;
 	int idx;
 
 	hyp_assert_lock_held(&vm_table_lock);
@@ -670,6 +685,9 @@ static int allocate_vm_table_entry(void)
 
 	vm_table[idx] = RESERVED_ENTRY;
 
+	for (cpu = 0; cpu < hyp_nr_cpus; cpu++)
+		per_cpu(last_vcpu_ran[idx], cpu) = UINT_MAX;
+
 	return idx;
 }
 

base-commit: dd0bc00e3e9201fd49b62ff9cc07890c89003afe
-- 
2.47.3