Invalidate TLB if necessary when switching vCPUs.

Bug: 132422393
Change-Id: If5983b9c39c92604a75ea42839b73376577a4f65
diff --git a/inc/hf/vm.h b/inc/hf/vm.h
index e35edc6..6a8f8b1 100644
--- a/inc/hf/vm.h
+++ b/inc/hf/vm.h
@@ -18,6 +18,8 @@
 
 #include <stdatomic.h>
 
+#include "hf/arch/types.h"
+
 #include "hf/cpu.h"
 #include "hf/list.h"
 #include "hf/mm.h"
@@ -85,6 +87,9 @@
 	struct wait_entry wait_entries[MAX_VMS];
 
 	atomic_bool aborting;
+
+	/** Arch-specific VM information. */
+	struct arch_vm arch;
 };
 
 /** Encapsulates a VM whose lock is held. */
diff --git a/src/arch/aarch64/hypervisor/exceptions.S b/src/arch/aarch64/hypervisor/exceptions.S
index 9acd240..6524540 100644
--- a/src/arch/aarch64/hypervisor/exceptions.S
+++ b/src/arch/aarch64/hypervisor/exceptions.S
@@ -476,6 +476,16 @@
 	msr ich_hcr_el2, x3
 #endif
 
+	/*
+	 * If a different vCPU is being run on this physical CPU to the last one
+	 * which was run for this VM, invalidate the TLB. This must be called
+	 * after vttbr_el2 has been updated, so that we have the page table and
+	 * VMID of the vCPU to which we are switching.
+	 */
+	mov x19, x0
+	bl maybe_invalidate_tlb
+	mov x0, x19
+
 	/* Restore non-volatile registers. */
 	ldp x19, x20, [x0, #VCPU_REGS + 8 * 19]
 	ldp x21, x22, [x0, #VCPU_REGS + 8 * 21]
diff --git a/src/arch/aarch64/hypervisor/handler.c b/src/arch/aarch64/hypervisor/handler.c
index c580c1a..7de6ad3 100644
--- a/src/arch/aarch64/hypervisor/handler.c
+++ b/src/arch/aarch64/hypervisor/handler.c
@@ -16,6 +16,7 @@
 
 #include <stdnoreturn.h>
 
+#include "hf/arch/barriers.h"
 #include "hf/arch/init.h"
 
 #include "hf/api.h"
@@ -125,6 +126,55 @@
 	}
 }
 
+/**
+ * Ensures all explicit memory access and management instructions for
+ * non-shareable normal memory have completed before continuing.
+ */
+static void dsb_nsh(void)
+{
+	__asm__ volatile("dsb nsh");
+}
+
+/**
+ * Invalidate all stage 1 TLB entries on the current (physical) CPU for the
+ * current VMID.
+ */
+static void invalidate_vm_tlb(void)
+{
+	isb();
+	__asm__ volatile("tlbi vmalle1");
+	isb();
+	dsb_nsh();
+}
+
+/**
+ * Invalidates the TLB if a different vCPU is being run than the last vCPU of
+ * the same VM which was run on the current pCPU.
+ *
+ * This is necessary because VMs may (contrary to the architecture
+ * specification) use inconsistent ASIDs across vCPUs. c.f. KVM's similar
+ * workaround:
+ * https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/commit/?id=94d0e5980d6791b9
+ */
+void maybe_invalidate_tlb(struct vcpu *vcpu)
+{
+	size_t current_cpu_index = cpu_index(vcpu->cpu);
+	size_t new_vcpu_index = vcpu_index(vcpu);
+
+	if (vcpu->vm->arch.last_vcpu_on_cpu[current_cpu_index] !=
+	    new_vcpu_index) {
+		/*
+		 * The vCPU has changed since the last time this VM was run on
+		 * this pCPU, so we need to invalidate the TLB.
+		 */
+		invalidate_vm_tlb();
+
+		/* Record the fact that this vCPU is now running on this CPU. */
+		vcpu->vm->arch.last_vcpu_on_cpu[current_cpu_index] =
+			new_vcpu_index;
+	}
+}
+
 noreturn void irq_current_exception(uintreg_t elr, uintreg_t spsr)
 {
 	(void)elr;
diff --git a/src/arch/aarch64/inc/hf/arch/types.h b/src/arch/aarch64/inc/hf/arch/types.h
index f768309..cc1cbef 100644
--- a/src/arch/aarch64/inc/hf/arch/types.h
+++ b/src/arch/aarch64/inc/hf/arch/types.h
@@ -50,7 +50,18 @@
 static_assert(sizeof(struct float_reg) == FLOAT_REG_BYTES,
 	      "Ensure float register type is 128 bits.");
 
-/** Type to represent the register state of a VM.  */
+/** Arch-specifc information about a VM. */
+struct arch_vm {
+	/**
+	 * The index of the last vCPU of this VM which ran on each pCPU. Each
+	 * element of this array should only be read or written by code running
+	 * on that CPU, which avoids contention and so no lock is needed to
+	 * access this field.
+	 */
+	uint32_t last_vcpu_on_cpu[MAX_CPUS];
+};
+
+/** Type to represent the register state of a vCPU.  */
 struct arch_regs {
 	/* General purpose registers. */
 	uintreg_t r[31];
diff --git a/src/arch/fake/inc/hf/arch/types.h b/src/arch/fake/inc/hf/arch/types.h
index b0b6803..e4fda8d 100644
--- a/src/arch/fake/inc/hf/arch/types.h
+++ b/src/arch/fake/inc/hf/arch/types.h
@@ -34,6 +34,12 @@
 /** The integer corresponding to the native register size. */
 typedef uint64_t uintreg_t;
 
+/** Arch-specifc information about a VM. */
+struct arch_vm {
+	/* This field is only here because empty structs aren't allowed. */
+	void *dummy;
+};
+
 /** Type to represent the register state of a VM.  */
 struct arch_regs {
 	uintreg_t r[5];