From e558631c73cd0e9e92f9e2d78e0d13624a6d54df Mon Sep 17 00:00:00 2001
From: Corey Lewis <corey.lewis@proofcraft.systems>
Date: Tue, 10 Sep 2024 18:17:28 +1000
Subject: [PATCH] arm: flush vcpu when switching domain

This avoids information about vcpu state being leaked across domains.

Signed-off-by: Corey Lewis <corey.lewis@proofcraft.systems>
---
 include/arch/arm/arch/object/vcpu.h | 2 ++
 include/kernel/thread.h             | 1 +
 src/arch/arm/kernel/thread.c        | 7 +++++++
 src/arch/arm/object/vcpu.c          | 8 ++++++++
 src/arch/riscv/kernel/thread.c      | 5 +++++
 src/arch/x86/kernel/thread.c        | 5 +++++
 src/kernel/thread.c                 | 1 +
 7 files changed, 29 insertions(+)

diff --git a/include/arch/arm/arch/object/vcpu.h b/include/arch/arm/arch/object/vcpu.h
index cc9ac5e94..fb45fe976 100644
--- a/include/arch/arm/arch/object/vcpu.h
+++ b/include/arch/arm/arch/object/vcpu.h
@@ -123,6 +123,7 @@ exception_t decodeARMVCPUInvocation(
 
 void vcpu_restore(vcpu_t *cpu);
 void vcpu_switch(vcpu_t *cpu);
+void vcpu_flush(void);
 #ifdef ENABLE_SMP_SUPPORT
 void handleVCPUInjectInterruptIPI(vcpu_t *vcpu, unsigned long index, virq_t virq);
 #endif /* ENABLE_SMP_SUPPORT */
@@ -206,6 +207,7 @@ static inline VPPIEventIRQ_t irqVPPIEventIndex(irq_t irq)
 
 #define vcpu_boot_init() do {} while(0)
 #define vcpu_switch(x) do {} while(0)
+#define vcpu_flush() do {} while(0)
 static inline void VGICMaintenance(void) {}
 
 #endif /* end of !CONFIG_ARM_HYPERVISOR_SUPPORT */
diff --git a/include/kernel/thread.h b/include/kernel/thread.h
index e5728e231..475195482 100644
--- a/include/kernel/thread.h
+++ b/include/kernel/thread.h
@@ -164,6 +164,7 @@ void Arch_switchToThread(tcb_t *tcb);
 void Arch_switchToIdleThread(void);
 void Arch_configureIdleThread(tcb_t *tcb);
 void Arch_activateIdleThread(tcb_t *tcb);
+void Arch_prepareNextDomain(void);
 
 void NORETURN idle_thread(void);
 
diff --git a/src/arch/arm/kernel/thread.c b/src/arch/arm/kernel/thread.c
index f84d59792..5dd4c1f20 100644
--- a/src/arch/arm/kernel/thread.c
+++ b/src/arch/arm/kernel/thread.c
@@ -10,3 +10,10 @@ void Arch_postModifyRegisters(tcb_t *tptr)
 {
     /* Nothing to do */
 }
+
+void Arch_prepareNextDomain(void)
+{
+    if (config_set(CONFIG_ARM_HYPERVISOR_SUPPORT)) {
+        vcpu_flush();
+    }
+}
diff --git a/src/arch/arm/object/vcpu.c b/src/arch/arm/object/vcpu.c
index 89d690c0f..7a86fbc66 100644
--- a/src/arch/arm/object/vcpu.c
+++ b/src/arch/arm/object/vcpu.c
@@ -264,6 +264,14 @@ static void vcpu_invalidate_active(void)
     ARCH_NODE_STATE(armHSCurVCPU) = NULL;
 }
 
+void vcpu_flush(void)
+{
+    if (ARCH_NODE_STATE(armHSCurVCPU)) {
+        vcpu_save(ARCH_NODE_STATE(armHSCurVCPU), ARCH_NODE_STATE(armHSVCPUActive));
+        vcpu_invalidate_active();
+    }
+}
+
 void vcpu_finalise(vcpu_t *vcpu)
 {
     if (vcpu->vcpuTCB) {
diff --git a/src/arch/riscv/kernel/thread.c b/src/arch/riscv/kernel/thread.c
index 05fe6e8d6..a2239a593 100644
--- a/src/arch/riscv/kernel/thread.c
+++ b/src/arch/riscv/kernel/thread.c
@@ -55,3 +55,8 @@ void Arch_postModifyRegisters(tcb_t *tptr)
 {
     /* Nothing to do */
 }
+
+void Arch_prepareNextDomain(void)
+{
+    /* Don't need to do anything */
+}
diff --git a/src/arch/x86/kernel/thread.c b/src/arch/x86/kernel/thread.c
index 4029c431c..28e50aefe 100644
--- a/src/arch/x86/kernel/thread.c
+++ b/src/arch/x86/kernel/thread.c
@@ -11,3 +11,8 @@ void Arch_postModifyRegisters(tcb_t *tptr)
 {
     Mode_postModifyRegisters(tptr);
 }
+
+void Arch_prepareNextDomain(void)
+{
+    /* Don't need to do anything */
+}
diff --git a/src/kernel/thread.c b/src/kernel/thread.c
index 11594e4f5..3fc32f15a 100644
--- a/src/kernel/thread.c
+++ b/src/kernel/thread.c
@@ -314,6 +314,7 @@ void prepareSetDomain(tcb_t *tptr, dom_t dom)
 
 static void prepareNextDomain(void)
 {
+    Arch_prepareNextDomain();
 #ifdef CONFIG_HAVE_FPU
     /* Save FPU state now to avoid touching cross-domain state later */
     switchLocalFpuOwner(NULL);