From 7492e52c7376caafbc9323128e18744fb0e79cdb Mon Sep 17 00:00:00 2001
From: _xeroxz <xerox@back.engineer>
Date: Fri, 2 Apr 2021 00:12:26 -0700
Subject: [PATCH] more cleaning

---
 exit_handler.cpp | 12 ++++++++++++
 hv_types.hpp     |  1 -
 ia32.hpp         | 10 +++++++++-
 vmcs.cpp         |  2 ++
 4 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/exit_handler.cpp b/exit_handler.cpp
index b39d324..080bdbb 100644
--- a/exit_handler.cpp
+++ b/exit_handler.cpp
@@ -32,10 +32,22 @@ auto exit_handler(hv::pguest_registers regs) -> void
 		ia32_vmx_procbased_ctls_register procbased_ctls;
 		__vmx_vmread(VMCS_CTRL_PROCESSOR_BASED_VM_EXECUTION_CONTROLS, &procbased_ctls.flags);
 
+		// if an NMI happens right here and my NMI handler sets nmi window exiting on...
+		// its going to be overwritten with off right here... sus...
+
 		procbased_ctls.nmi_window_exiting = false;
 		__vmx_vmwrite(VMCS_CTRL_PROCESSOR_BASED_VM_EXECUTION_CONTROLS, procbased_ctls.flags);
 		goto dont_advance;
 	}
+	case VMX_EXIT_REASON_EXCEPTION_OR_NMI:
+	{
+		ia32_vmx_procbased_ctls_register procbased_ctls;
+		__vmx_vmread(VMCS_CTRL_PROCESSOR_BASED_VM_EXECUTION_CONTROLS, &procbased_ctls.flags);
+
+		procbased_ctls.nmi_window_exiting = true;
+		__vmx_vmwrite(VMCS_CTRL_PROCESSOR_BASED_VM_EXECUTION_CONTROLS, procbased_ctls.flags);
+		goto dont_advance;
+	}
 	case VMX_EXIT_REASON_EXECUTE_XSETBV:
 	{
 		if (handle::xsetbv(regs))
diff --git a/hv_types.hpp b/hv_types.hpp
index e56dc15..479536b 100644
--- a/hv_types.hpp
+++ b/hv_types.hpp
@@ -689,7 +689,6 @@ namespace hv
 		u64 vmcs_phys;
 		u64 vmxon_phys;
 		u64 error_code;
-		u64 nmi_code;
 	} vcpu_ctx, * pvcpu_ctx;
 
 	typedef struct _vmx_ctx
diff --git a/ia32.hpp b/ia32.hpp
index 2510e6c..7bb650e 100644
--- a/ia32.hpp
+++ b/ia32.hpp
@@ -8638,6 +8638,10 @@ typedef union
 #define IA32_MTRR_PHYSBASE7                                          0x0000020E
 #define IA32_MTRR_PHYSBASE8                                          0x00000210
 #define IA32_MTRR_PHYSBASE9                                          0x00000212
+
+#define IA32_MTRR_PHYSBASEN(n) \
+        IA32_MTRR_PHYSBASE0 + (n * 2)
+
 /**
  * @}
  */
@@ -8697,7 +8701,7 @@ typedef union
 
     uint64_t flags;
 } ia32_mtrr_physmask_register;
-
+    
 #define IA32_MTRR_PHYSMASK0                                          0x00000201
 #define IA32_MTRR_PHYSMASK1                                          0x00000203
 #define IA32_MTRR_PHYSMASK2                                          0x00000205
@@ -8708,6 +8712,10 @@ typedef union
 #define IA32_MTRR_PHYSMASK7                                          0x0000020F
 #define IA32_MTRR_PHYSMASK8                                          0x00000211
 #define IA32_MTRR_PHYSMASK9                                          0x00000213
+
+#define IA32_MTRR_PHYSMASKN(n) \
+        IA32_MTRR_PHYSMASK0 + (n * 2)
+
 /**
  * @}
  */
diff --git a/vmcs.cpp b/vmcs.cpp
index b9d1941..49c6861 100644
--- a/vmcs.cpp
+++ b/vmcs.cpp
@@ -153,6 +153,8 @@ namespace vmcs
 		ia32_vmx_basic_register vmx_basic{ __readmsr(IA32_VMX_BASIC) };
 
 		pinbased_ctls.virtual_nmi = true;
+		pinbased_ctls.nmi_exiting = true;
+
 		procbased_ctls.activate_secondary_controls = true;
 		exit_ctls.host_address_space_size = true;