WIP Cleaner way to launch.

4 files changed, 35 insertions, 32 deletions

diff --git a/shv.h b/shv.h
index e5f34ba..2095239 100644
--- a/shv.h
+++ b/shv.h
@@ -143,6 +143,11 @@ ShvVmxEptInitialize (
     _In_ PSHV_VP_DATA VpData
     );
 
+DECLSPEC_NORETURN
+VOID
+ShvVpRestoreAfterLaunch (
+    VOID
+    );
 
 typedef struct _SHV_DPC_CONTEXT
 {
diff --git a/shvvmx.c b/shvvmx.c
index 584f77f..df24586 100644
--- a/shvvmx.c
+++ b/shvvmx.c
@@ -375,7 +375,7 @@ ShvVmxSetupVmcsForVp (
     // to inside of ShvVpInitialize.
     //
     __vmx_vmwrite(GUEST_RSP, context->Rsp);
-    __vmx_vmwrite(GUEST_RIP, context->Rip);
+    __vmx_vmwrite(GUEST_RIP, (ULONG_PTR)ShvVpRestoreAfterLaunch);
     __vmx_vmwrite(GUEST_RFLAGS, context->EFlags);
 
     //
@@ -464,11 +464,6 @@ ShvVmxLaunchOnVp (
         // Initialize the VMCS, both guest and host state.
         //
         ShvVmxSetupVmcsForVp(VpData);
-        
-        //
-        // Record that VMX is now enabled
-        //
-        VpData->VmxEnabled = 1;
 
         //
         // Launch the VMCS, based on the guest data that was loaded into the
diff --git a/shvvmxhv.c b/shvvmxhv.c
index 6acb57f..68a1fbd 100644
--- a/shvvmxhv.c
+++ b/shvvmxhv.c
@@ -16,7 +16,7 @@ Author:
 
 Environment:
 
-    Hypervisor mode only, IRQL DIRQL_MAX
+    Hypervisor mode only, IRQL MAX_IRQL
 
 --*/
 
diff --git a/shvvp.c b/shvvp.c
index d2632a7..e11b046 100644
--- a/shvvp.c
+++ b/shvvp.c
@@ -79,6 +79,29 @@ ShvCaptureSpecialRegisters (
     _sldt(&SpecialRegisters->Ldtr);
 }
 
+DECLSPEC_NORETURN
+VOID
+ShvVpRestoreAfterLaunch (
+    VOID
+    )
+{
+    PSHV_VP_DATA vpData = ShvGlobalData[KeGetCurrentProcessorNumberEx(NULL)];
+
+    //
+    // Record that VMX is now enabled
+    //
+    vpData->VmxEnabled = 1;
+
+    //
+    // And finally, restore the context, so that all register and stack
+    // state is finally restored. Note that by continuing to reference the
+    // per-VP data this way, the compiler will continue to generate non-
+    // optimized accesses, guaranteeing that no previous register state
+    // will be used.
+    //
+    RtlRestoreContext(&vpData->ContextFrame, NULL);
+}
+
 VOID
 ShvVpInitialize (
     _In_ PSHV_VP_DATA Data,
@@ -114,34 +137,14 @@ ShvVpInitialize (
     // variable combined with an API call, we also make sure that the compiler
     // will not optimize this access in any way, even on LTGC/Ox builds.
     //
-    if (ShvGlobalData[KeGetCurrentProcessorNumberEx(NULL)]->VmxEnabled == 1)
+    if (ShvGlobalData[KeGetCurrentProcessorNumberEx(NULL)]->VmxEnabled == 0)
     {
         //
-        // We now indicate that the VM has launched, and that we are about to
-        // restore the GPRs back to their original values. This will have the
-        // effect of putting us yet *AGAIN* at the previous line of code, but
-        // this time the value of VmxEnabled will be two, bypassing the if and
-        // else if checks.
-        //
-        ShvGlobalData[KeGetCurrentProcessorNumberEx(NULL)]->VmxEnabled = 2;
-
-        //
-        // And finally, restore the context, so that all register and stack
-        // state is finally restored. Note that by continuing to reference the
-        // per-VP data this way, the compiler will continue to generate non-
-        // optimized accesses, guaranteeing that no previous register state
-        // will be used.
+        // If we are in this branch comparison, it means that we have not yet
+        // attempted to launch the VM, nor that we have launched it. In other
+        // words, this is the first time in ShvVpInitialize. Because of this,
+        // we are free to use all register state, as it is ours to use.
         //
-        RtlRestoreContext(&ShvGlobalData[KeGetCurrentProcessorNumberEx(NULL)]->ContextFrame, NULL);
-    }
-    //
-    // If we are in this branch comparison, it means that we have not yet
-    // attempted to launch the VM, nor that we have launched it. In other
-    // words, this is the first time in ShvVpInitialize. Because of this,
-    // we are free to use all register state, as it is ours to use.
-    //
-    else if (Data->VmxEnabled == 0)
-    {
         //
         // First, capture the value of the PML4 for the SYSTEM process, so that
         // all virtual processors, regardless of which process the current LP