From a08a850914a2e4930ce4f00a1e5b54d9eea0b9e4 Mon Sep 17 00:00:00 2001
From: Joshua Rogers <Joshua@Joshua.Hu>
Date: Wed, 11 Sep 2024 19:42:25 +0200
Subject: [PATCH] vmm: Correctly suspend and resume the vmm driver.

Previously, VMXON would be executed on a resume, contrary to proper
initalization. The MSR lock may be lost on suspension, therefore must be taken
again. Likewise, the VMX Enable bit may be cleared upon suspend, requiring to
be re-set.

Concretely disable VMX on suspend, and re-enable it on resume.

Note: any IOMMU context will remain lost for any enabled vmm devices.

Signed-off-by: Joshua Rogers <Joshua@Joshua.Hu>
---
 sys/amd64/acpica/acpi_wakeup.c |  3 +++
 sys/amd64/amd64/machdep.c      |  1 +
 sys/amd64/include/cpu.h        |  3 ++-
 sys/amd64/include/vmm.h        |  2 ++
 sys/amd64/vmm/amd/svm.c        |  9 ++++++++-
 sys/amd64/vmm/intel/vmx.c      | 11 ++++++++++-
 sys/amd64/vmm/vmm.c            |  1 +
 sys/x86/x86/mp_x86.c           |  5 +++++
 8 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/sys/amd64/acpica/acpi_wakeup.c b/sys/amd64/acpica/acpi_wakeup.c
index 351cd16ff9ca78..51d6d5e3684002 100644
--- a/sys/amd64/acpica/acpi_wakeup.c
+++ b/sys/amd64/acpica/acpi_wakeup.c
@@ -202,6 +202,9 @@ acpi_sleep_machdep(struct acpi_softc *sc, int state)
 
 	intr_suspend();
 
+	if (vmm_suspend_p != NULL)
+		vmm_suspend_p();
+
 	pcb = &susppcbs[0]->sp_pcb;
 	if (savectx(pcb)) {
 		fpususpend(susppcbs[0]->sp_fpususpend);
diff --git a/sys/amd64/amd64/machdep.c b/sys/amd64/amd64/machdep.c
index 025c3c365de51c..6b2623144c5a97 100644
--- a/sys/amd64/amd64/machdep.c
+++ b/sys/amd64/amd64/machdep.c
@@ -213,6 +213,7 @@ struct mem_range_softc mem_range_softc;
 
 struct mtx dt_lock;	/* lock for GDT and LDT */
 
+void (*vmm_suspend_p)(void);
 void (*vmm_resume_p)(void);
 
 bool efi_boot;
diff --git a/sys/amd64/include/cpu.h b/sys/amd64/include/cpu.h
index 99b8806ba0bacc..57f2a0b59bbb4c 100644
--- a/sys/amd64/include/cpu.h
+++ b/sys/amd64/include/cpu.h
@@ -70,7 +70,8 @@ extern char	btext[];
 extern char	_end[];
 extern char	etext[];
 
-/* Resume hook for VMM. */
+/* Suspend and resume hook for VMM. */
+extern	void (*vmm_suspend_p)(void);
 extern	void (*vmm_resume_p)(void);
 
 void	cpu_halt(void);
diff --git a/sys/amd64/include/vmm.h b/sys/amd64/include/vmm.h
index 37972d54bd9975..64d91d51e6fd97 100644
--- a/sys/amd64/include/vmm.h
+++ b/sys/amd64/include/vmm.h
@@ -170,6 +170,7 @@ struct vm_eventinfo {
 
 typedef int	(*vmm_init_func_t)(int ipinum);
 typedef int	(*vmm_cleanup_func_t)(void);
+typedef void	(*vmm_suspend_func_t)(void);
 typedef void	(*vmm_resume_func_t)(void);
 typedef void *	(*vmi_init_func_t)(struct vm *vm, struct pmap *pmap);
 typedef int	(*vmi_run_func_t)(void *vcpui, register_t rip,
@@ -194,6 +195,7 @@ typedef int	(*vmi_restore_tsc_t)(void *vcpui, uint64_t now);
 struct vmm_ops {
 	vmm_init_func_t		modinit;	/* module wide initialization */
 	vmm_cleanup_func_t	modcleanup;
+	vmm_resume_func_t	modsuspend;
 	vmm_resume_func_t	modresume;
 
 	vmi_init_func_t		init;		/* vm-specific initialization */
diff --git a/sys/amd64/vmm/amd/svm.c b/sys/amd64/vmm/amd/svm.c
index b35cf53be080b1..78fa726dc45df5 100644
--- a/sys/amd64/vmm/amd/svm.c
+++ b/sys/amd64/vmm/amd/svm.c
@@ -276,12 +276,19 @@ svm_modinit(int ipinum)
 	return (0);
 }
 
+static void
+svm_modsuspend(void)
+{
+
+	return;
+}
+
 static void
 svm_modresume(void)
 {
 
 	svm_enable(NULL);
-}		
+}
 
 #ifdef BHYVE_SNAPSHOT
 void
diff --git a/sys/amd64/vmm/intel/vmx.c b/sys/amd64/vmm/intel/vmx.c
index 31ca9256b527ed..b594854ded1be7 100644
--- a/sys/amd64/vmm/intel/vmx.c
+++ b/sys/amd64/vmm/intel/vmx.c
@@ -647,12 +647,20 @@ vmx_enable(void *arg __unused)
 		vmxon_enabled[curcpu] = 1;
 }
 
+static void
+vmx_modsuspend(void)
+{
+
+	if (vmxon_enabled[curcpu])
+		vmx_disable(NULL);
+}
+
 static void
 vmx_modresume(void)
 {
 
 	if (vmxon_enabled[curcpu])
-		vmxon(&vmxon_region[curcpu * PAGE_SIZE]);
+		vmx_enable(NULL);
 }
 
 static int
@@ -4270,6 +4278,7 @@ vmx_restore_tsc(void *vcpui, uint64_t offset)
 const struct vmm_ops vmm_ops_intel = {
 	.modinit	= vmx_modinit,
 	.modcleanup	= vmx_modcleanup,
+	.modsuspend	= vmx_modsuspend,
 	.modresume	= vmx_modresume,
 	.init		= vmx_init,
 	.run		= vmx_run,
diff --git a/sys/amd64/vmm/vmm.c b/sys/amd64/vmm/vmm.c
index f399f876717dc7..030e76a5e67fba 100644
--- a/sys/amd64/vmm/vmm.c
+++ b/sys/amd64/vmm/vmm.c
@@ -231,6 +231,7 @@ vmmops_panic(void)
 
 DEFINE_VMMOPS_IFUNC(int, modinit, (int ipinum))
 DEFINE_VMMOPS_IFUNC(int, modcleanup, (void))
+DEFINE_VMMOPS_IFUNC(void, modsuspend, (void))
 DEFINE_VMMOPS_IFUNC(void, modresume, (void))
 DEFINE_VMMOPS_IFUNC(void *, init, (struct vm *vm, struct pmap *pmap))
 DEFINE_VMMOPS_IFUNC(int, run, (void *vcpui, register_t rip, struct pmap *pmap,
diff --git a/sys/x86/x86/mp_x86.c b/sys/x86/x86/mp_x86.c
index 0a683b3e841a78..493017e303e3d9 100644
--- a/sys/x86/x86/mp_x86.c
+++ b/sys/x86/x86/mp_x86.c
@@ -1591,6 +1591,11 @@ cpususpend_handler(void)
 
 	mtx_assert(&smp_ipi_mtx, MA_NOTOWNED);
 
+#ifdef __amd64__
+	if (vmm_suspend_p)
+		vmm_suspend_p();
+#endif
+
 	cpu = PCPU_GET(cpuid);
 
 #ifdef XENHVM