Separate Hypervisor Core from OS Layer. Don't touch IRQL in Hypervisor. Separate Load vs Unload callback. Misc. portability fixes.

Create a layer of OS-specific functions to handle the various
requirements around memory allocation, context save/restore,
entrypoint/unloadpoint, and multi-CPU execution and topology
information. SimpleVisor no longer uses NT-specific functions (some
structures and types still remain).
Additionally, the hypervisor should not know that "NT" is running
underneath, so it has no business touching the IRQL. As we won't call
Windows functions, and as interrupts are disabled, this doesn't 'change'
anything and is correct.
Don't use the same callback for load and unload. We can make unload its
own callback now, as we've separated out the DPC-specific logic. This
makes the load callback cleaner as well.
Remove NT_ASSERTS which don't work anyway, and use portable
definitions/types when possible (more to do here).
Return the failed CPU and status in all cases during load. Sometimes
this wasn't done before.

9 files changed, 340 insertions, 230 deletions

diff --git a/ntint.h b/ntint.h
index a78068c..22fff93 100644
--- a/ntint.h
+++ b/ntint.h
@@ -23,44 +23,10 @@ Environment:
 #pragma once
 
 //
-// Define pseudo descriptor structures for both 64- and 32-bit mode.
+// Definitions from ntosp.h
 //
-
-typedef struct _KDESCRIPTOR {
-    USHORT Pad[3];
-    USHORT Limit;
-    PVOID Base;
-} KDESCRIPTOR, *PKDESCRIPTOR;
-
-//
-// Define descriptor privilege levels for user and system.
-//
-
-#define DPL_USER 3
-#define DPL_SYSTEM 0
-
-//
-// Define limit granularity.
-//
-
-#define GRANULARITY_BYTE 0
-#define GRANULARITY_PAGE 1
-
-//
-// Define processor number packing constants.
-//
-// The compatibility processor number is encoded in the FS segment descriptor.
-//
-// Bits 19:14 of the segment limit encode the compatible processor number.
-// Bits 13:10 are set to ones to ensure that segment limit is at least 15360.
-// Bits 9:0 of the segment limit encode the extended processor number.
-//
-
-#define KGDT_LEGACY_LIMIT_SHIFT 14
-#define KGDT_LIMIT_ENCODE_MASK (0xf << 10)
-
-#define SELECTOR_TABLE_INDEX 0x04
-
+#define DPL_USER            3
+#define DPL_SYSTEM          0
 #define KGDT64_NULL         0x00
 #define KGDT64_R0_CODE      0x10
 #define KGDT64_R0_DATA      0x18
@@ -70,29 +36,38 @@ typedef struct _KDESCRIPTOR {
 #define KGDT64_SYS_TSS      0x40
 #define KGDT64_R3_CMTEB     0x50
 #define KGDT64_R0_LDT       0x60
-
 #define MSR_GS_BASE         0xC0000101
 #define MSR_DEBUG_CTL       0x1D9
-
-#define RPL_MASK 3
-
-#define MTRR_TYPE_WB 6
-
+#define RPL_MASK            3
+#define MTRR_TYPE_WB        6
 #define EFLAGS_ALIGN_CHECK  0x40000
 
-typedef union _KGDTENTRY64 {
-    struct {
+//
+// Structures from ntosp.h
+//
+typedef struct _KDESCRIPTOR
+{
+    USHORT Pad[3];
+    USHORT Limit;
+    PVOID Base;
+} KDESCRIPTOR, *PKDESCRIPTOR;
+typedef union _KGDTENTRY64
+{
+    struct
+    {
         USHORT LimitLow;
         USHORT BaseLow;
-        union {
-            struct {
+        union
+        {
+            struct
+            {
                 UCHAR BaseMiddle;
                 UCHAR Flags1;
                 UCHAR Flags2;
                 UCHAR BaseHigh;
             } Bytes;
-
-            struct {
+            struct 
+            {
                 ULONG BaseMiddle : 8;
                 ULONG Type : 5;
                 ULONG Dpl : 2;
@@ -105,49 +80,13 @@ typedef union _KGDTENTRY64 {
                 ULONG BaseHigh : 8;
             } Bits;
         };
-
         ULONG BaseUpper;
         ULONG MustBeZero;
     };
-
-    struct {
+    struct
+    {
         LONG64 DataLow;
         LONG64 DataHigh;
     };
 
 } KGDTENTRY64, *PKGDTENTRY64;
-
-NTKERNELAPI
-_IRQL_requires_max_(APC_LEVEL)
-_IRQL_requires_min_(PASSIVE_LEVEL)
-_IRQL_requires_same_
-VOID
-KeGenericCallDpc (
-    _In_ PKDEFERRED_ROUTINE Routine,
-    _In_opt_ PVOID Context
-    );
-
-NTKERNELAPI
-_IRQL_requires_(DISPATCH_LEVEL)
-_IRQL_requires_same_
-VOID
-KeSignalCallDpcDone (
-    _In_ PVOID SystemArgument1
-    );
-
-NTKERNELAPI
-_IRQL_requires_(DISPATCH_LEVEL)
-_IRQL_requires_same_
-LOGICAL
-KeSignalCallDpcSynchronize (
-    _In_ PVOID SystemArgument2
-    );
-
-DECLSPEC_NORETURN
-NTSYSAPI
-VOID
-__cdecl
-RtlRestoreContext(
-    _In_ PCONTEXT ContextRecord,
-    _In_opt_ struct _EXCEPTION_RECORD * ExceptionRecord
-    );
diff --git a/shv.c b/shv.c
index e10973c..8672c8c 100644
--- a/shv.c
+++ b/shv.c
@@ -24,38 +24,32 @@ Environment:
 
 VOID
 ShvUnload (
-    _In_ PDRIVER_OBJECT DriverObject
+    VOID
     )
 {
-    SHV_DPC_CONTEXT dpcContext;
-    UNREFERENCED_PARAMETER(DriverObject);
-
     //
     // Attempt to exit VMX root mode on all logical processors. This will
-    // broadcast a DPC interrupt which will execute the callback routine in
-    // parallel on the LPs. Send the callback routine a NULL context in order
-    // to indicate that this is the unload, not load, path.
+    // broadcast an interrupt which will execute the callback routine in
+    // parallel on the LPs.
     //
     // Note that if SHV is not loaded on any of the LPs, this routine will not
     // perform any work, but will not fail in any way.
     //
-    dpcContext.Cr3 = 0;
-    KeGenericCallDpc(ShvVpCallbackDpc, &dpcContext);
+    ShvOsRunCallbackOnProcessors(ShvVpUnloadCallback, NULL);
 
     //
     // Indicate unload
     //
-    DbgPrintEx(77, 0, "The SHV has been uninstalled.\n");
+    ShvOsDebugPrint("The SHV has been uninstalled.\n");
 }
 
 NTSTATUS
-ShvInitialize (
-    _In_ PDRIVER_OBJECT DriverObject,
-    _In_ PUNICODE_STRING RegistryPath
+ShvLoad (
+    VOID
     )
 {
-    SHV_DPC_CONTEXT dpcContext;
-    UNREFERENCED_PARAMETER(RegistryPath);
+    SHV_CALLBACK_CONTEXT callbackContext;
+    ULONG cpuCount;
 
     //
     // Attempt to enter VMX root mode on all logical processors. This will
@@ -64,12 +58,11 @@ ShvInitialize (
     // the PML4 of the system process, which is what this driver entrypoint
     // should be executing in.
     //
-    NT_ASSERT(PsGetCurrentProcess() == PsInitialSystemProcess);
-    dpcContext.Cr3 = __readcr3();
-    dpcContext.FailureStatus = STATUS_SUCCESS;
-    dpcContext.FailedCpu = -1;
-    dpcContext.InitCount = 0;
-    KeGenericCallDpc(ShvVpCallbackDpc, &dpcContext);
+    callbackContext.Cr3 = __readcr3();
+    callbackContext.FailureStatus = STATUS_SUCCESS;
+    callbackContext.FailedCpu = -1;
+    callbackContext.InitCount = 0;
+    ShvOsRunCallbackOnProcessors(ShvVpLoadCallback, &callbackContext);
 
     //
     // Check if all LPs are now hypervised. Return the failure code of at least
@@ -77,18 +70,17 @@ ShvInitialize (
     //
     // Note that each VP is responsible for freeing its VP data on failure.
     //
-    if (dpcContext.InitCount != KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS))
+    cpuCount = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS);
+    if (callbackContext.InitCount != cpuCount)
     {
-        DbgPrintEx(77, 0, "The SHV failed to initialize (0x%lX) Failed CPU: %d\n",
-                   dpcContext.FailureStatus, dpcContext.FailedCpu);
-        NT_ASSERT(dpcContext.FailureStatus != STATUS_SUCCESS);
-        return dpcContext.FailureStatus;
+        ShvOsDebugPrint("The SHV failed to initialize (0x%lX) Failed CPU: %d\n",
+                        callbackContext.FailureStatus, callbackContext.FailedCpu);
+        return callbackContext.FailureStatus;
     }
 
     //
-    // Make the driver (and SHV itself) unloadable, and indicate success.
+    // Indicate success.
     //
-    DriverObject->DriverUnload = ShvUnload;
-    DbgPrintEx(77, 0, "The SHV has been installed.\n");
+    ShvOsDebugPrint("The SHV has been installed.\n");
     return STATUS_SUCCESS;
 }
diff --git a/shv.h b/shv.h
index 7331a0c..e78e557 100644
--- a/shv.h
+++ b/shv.h
@@ -79,10 +79,24 @@ typedef struct _SHV_VP_STATE
     ULONG_PTR GuestRsp;
     ULONG_PTR GuestEFlags;
     USHORT ExitReason;
-    KIRQL GuestIrql;
     BOOLEAN ExitVm;
 } SHV_VP_STATE, *PSHV_VP_STATE;
 
+typedef struct _SHV_CALLBACK_CONTEXT
+{
+    ULONG64 Cr3;
+    volatile ULONG InitCount;
+    LONG FailedCpu;
+    NTSTATUS FailureStatus;
+} SHV_CALLBACK_CONTEXT, *PSHV_CALLBACK_CONTEXT;
+
+typedef
+VOID
+SHV_CPU_CALLBACK (
+    _In_ PSHV_CALLBACK_CONTEXT Context
+    );
+typedef SHV_CPU_CALLBACK *PSHV_CPU_CALLBACK;
+
 VOID
 ShvVmxEntry (
     VOID
@@ -142,21 +156,48 @@ ShvVmxEptInitialize (
     _In_ PSHV_VP_DATA VpData
     );
 
+NTSTATUS
+ShvLoad (
+    VOID
+    );
+
+VOID
+ShvUnload (
+    VOID
+    );
+
+DECLSPEC_NORETURN
+VOID
+__cdecl
+ShvOsRestoreContext (
+    _In_ PCONTEXT ContextRecord
+    );
+
+VOID
+ShvOsFreeContiguousAlignedMemory (
+    _In_ PVOID BaseAddress
+    );
+
+PVOID
+ShvOsAllocateContigousAlignedMemory (
+    _In_ SIZE_T Size
+    );
+
 DECLSPEC_NORETURN
 VOID
 ShvVpRestoreAfterLaunch (
     VOID
     );
 
-typedef struct _SHV_DPC_CONTEXT
-{
-    ULONG64 Cr3;
-    volatile ULONG InitCount;
-    LONG FailedCpu;
-    NTSTATUS FailureStatus;
-} SHV_DPC_CONTEXT, *PSHV_DPC_CONTEXT;
+VOID
+ShvOsRunCallbackOnProcessors (
+    _In_ PSHV_CPU_CALLBACK Routine,
+    _In_opt_ PVOID Context
+    );
 
-KDEFERRED_ROUTINE ShvVpCallbackDpc;
+SHV_CPU_CALLBACK ShvVpLoadCallback;
+SHV_CPU_CALLBACK ShvVpUnloadCallback;
 
 extern PSHV_VP_DATA* ShvGlobalData;
 
+#define ShvOsDebugPrint(format, ...) DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, format, __VA_ARGS__)
diff --git a/shv.vcxproj b/shv.vcxproj
index df17b2b..96471bd 100644
--- a/shv.vcxproj
+++ b/shv.vcxproj
@@ -27,6 +27,9 @@
     <DriverTargetPlatform>Desktop</DriverTargetPlatform>
   </PropertyGroup>
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.props" />
+  <PropertyGroup>
+    <EnableInf2cat>false</EnableInf2cat>
+  </PropertyGroup>
   <ItemDefinitionGroup>
     <ClCompile>
       <Optimization>Full</Optimization>
@@ -41,7 +44,7 @@
     <Link />
     <Link>
       <AdditionalDependencies>$(KernelBufferOverflowLib);$(DDK_LIB_PATH)ntoskrnl.lib;$(DDK_LIB_PATH)hal.lib;$(DDK_LIB_PATH)wmilib.lib</AdditionalDependencies>
-      <EntryPointSymbol>ShvInitialize</EntryPointSymbol>
+      <EntryPointSymbol>DriverEntry</EntryPointSymbol>
       <LinkTimeCodeGeneration>UseLinkTimeCodeGeneration</LinkTimeCodeGeneration>
       <Profile>false</Profile>
       <MergeSections>_TEXT=.text;_PAGE=PAGE</MergeSections>
@@ -52,6 +55,7 @@
     </Link>
   </ItemDefinitionGroup>
   <ItemGroup>
+    <ClCompile Include="shvos.c" />
     <ClCompile Include="shv.c" />
     <ClCompile Include="shvutil.c" />
     <ClCompile Include="shvvmx.c" />
diff --git a/shvos.c b/shvos.c
new file mode 100644
index 0000000..ec914da
--- /dev/null
+++ b/shvos.c
@@ -0,0 +1,189 @@
+/*++
+
+Copyright (c) Alex Ionescu.  All rights reserved.
+
+Module Name:
+
+    shvos.c
+
+Abstract:
+
+    This module implements the OS-facing Windows stubs for SimpleVisor.
+
+Author:
+
+    Alex Ionescu (@aionescu) 29-Aug-2016 - Initial version
+
+Environment:
+
+    Kernel mode only.
+
+--*/
+
+#include "shv.h"
+
+NTKERNELAPI
+_IRQL_requires_max_(APC_LEVEL)
+_IRQL_requires_min_(PASSIVE_LEVEL)
+_IRQL_requires_same_
+VOID
+KeGenericCallDpc (
+    _In_ PKDEFERRED_ROUTINE Routine,
+    _In_opt_ PVOID Context
+    );
+
+NTKERNELAPI
+_IRQL_requires_(DISPATCH_LEVEL)
+_IRQL_requires_same_
+VOID
+KeSignalCallDpcDone (
+    _In_ PVOID SystemArgument1
+    );
+
+NTKERNELAPI
+_IRQL_requires_(DISPATCH_LEVEL)
+_IRQL_requires_same_
+LOGICAL
+KeSignalCallDpcSynchronize (
+    _In_ PVOID SystemArgument2
+    );
+
+DECLSPEC_NORETURN
+NTSYSAPI
+VOID
+__cdecl
+RtlRestoreContext (
+    _In_ PCONTEXT ContextRecord,
+    _In_opt_ struct _EXCEPTION_RECORD * ExceptionRecord
+    );
+
+typedef struct _SHV_DPC_CONTEXT
+{
+    PSHV_CPU_CALLBACK Routine;
+    PSHV_CALLBACK_CONTEXT Context;
+} SHV_DPC_CONTEXT, *PSHV_DPC_CONTEXT;
+
+VOID
+ShvOsDpcRoutine (
+    _In_ struct _KDPC *Dpc,
+    _In_opt_ PVOID DeferredContext,
+    _In_opt_ PVOID SystemArgument1,
+    _In_opt_ PVOID SystemArgument2
+    )
+{
+    PSHV_DPC_CONTEXT dpcContext = DeferredContext;
+    UNREFERENCED_PARAMETER(Dpc);
+
+    //
+    // Execute the internal callback function
+    //
+    dpcContext->Routine(dpcContext->Context);
+
+    //
+    // Wait for all DPCs to synchronize at this point
+    //
+    KeSignalCallDpcSynchronize(SystemArgument2);
+
+    //
+    // Mark the DPC as being complete
+    //
+    KeSignalCallDpcDone(SystemArgument1);
+}
+
+VOID
+ShvOsFreeContiguousAlignedMemory (
+    _In_ PVOID BaseAddress
+    )
+{
+    //
+    // Free the memory
+    //
+    MmFreeContiguousMemory(BaseAddress);
+}
+
+PVOID
+ShvOsAllocateContigousAlignedMemory (
+    _In_ SIZE_T Size
+    )
+{
+    PHYSICAL_ADDRESS lowest, highest;
+
+    //
+    // The entire address range is OK for this allocation
+    //
+    lowest.QuadPart = 0;
+    highest.QuadPart = lowest.QuadPart - 1;
+
+    //
+    // Allocate a contiguous chunk of RAM to back this allocation and make sure
+    // that it is RW only, instead of RWX, by using the new Windows 8 API.
+    //
+    return MmAllocateContiguousNodeMemory(Size,
+                                          lowest,
+                                          highest,
+                                          lowest,
+                                          PAGE_READWRITE,
+                                          KeGetCurrentNodeNumber());
+}
+
+VOID
+ShvOsRunCallbackOnProcessors (
+    _In_ PSHV_CPU_CALLBACK Routine,
+    _In_opt_ PVOID Context
+    )
+{
+    SHV_DPC_CONTEXT dpcContext;
+
+    //
+    // Wrap the internal routine and context under a Windows DPC
+    //
+    dpcContext.Routine = Routine;
+    dpcContext.Context = Context;
+    KeGenericCallDpc(ShvOsDpcRoutine, &dpcContext);
+}
+
+DECLSPEC_NORETURN
+VOID
+__cdecl
+ShvOsRestoreContext (
+    _In_ PCONTEXT ContextRecord
+    )
+{
+    //
+    // Windows provides a nice OS function to do this
+    //
+    RtlRestoreContext(ContextRecord, NULL);
+}
+
+VOID
+DriverUnload (
+    _In_ PDRIVER_OBJECT DriverObject
+    )
+{
+    UNREFERENCED_PARAMETER(DriverObject);
+
+    //
+    // Unload the hypervisor
+    //
+    ShvUnload();
+}
+
+NTSTATUS
+DriverEntry (
+    _In_ PDRIVER_OBJECT DriverObject,
+    _In_ PUNICODE_STRING RegistryPath
+    )
+{
+    UNREFERENCED_PARAMETER(RegistryPath);
+
+    //
+    // Make the driver (and SHV itself) unloadable
+    //
+    DriverObject->DriverUnload = DriverUnload;
+
+    //
+    // Load the hypervisor
+    //
+    return ShvLoad();
+}
+
diff --git a/shvutil.c b/shvutil.c
index 6f36c09..5009f35 100644
--- a/shvutil.c
+++ b/shvutil.c
@@ -36,7 +36,6 @@ ShvUtilConvertGdtEntry (
     // Windows does not use an LDT for these selectors in kernel, so the TI bit
     // should never be set.
     //
-    NT_ASSERT((Selector & SELECTOR_TABLE_INDEX) == 0);
     gdtEntry = (PKGDTENTRY64)((ULONG_PTR)GdtBase + (Selector & ~RPL_MASK));
 
     //
diff --git a/shvvmx.c b/shvvmx.c
index 0fb4338..eb080ff 100644
--- a/shvvmx.c
+++ b/shvvmx.c
@@ -445,7 +445,7 @@ ShvVmxLaunchOnVp (
     //
     // Initialize all the VMX-related MSRs by reading their value
     //
-    for (i = 0; i < RTL_NUMBER_OF(VpData->MsrData); i++)
+    for (i = 0; i < sizeof(VpData->MsrData) / sizeof(VpData->MsrData[0]); i++)
     {
         VpData->MsrData[i].QuadPart = __readmsr(MSR_IA32_VMX_BASIC + i);
     }
diff --git a/shvvmxhv.c b/shvvmxhv.c
index 6a54fda..879928b 100644
--- a/shvvmxhv.c
+++ b/shvvmxhv.c
@@ -189,7 +189,6 @@ ShvVmxHandleExit (
         ShvVmxHandleVmx(VpState);
         break;
     default:
-        NT_ASSERT(FALSE);
         break;
     }
 
@@ -203,7 +202,6 @@ ShvVmxHandleExit (
 }
 
 DECLSPEC_NORETURN
-EXTERN_C
 VOID
 ShvVmxEntryHandler (
     _In_ PCONTEXT Context
@@ -213,19 +211,6 @@ ShvVmxEntryHandler (
     PSHV_VP_DATA vpData;
 
     //
-    // Because we run with interrupts disabled during the entire hypervisor's
-    // exit handling, raise the IRQL to HIGH_LEVEL which matches the reality of
-    // the situation. This will block IPIs and the clock interrupt timer, which
-    // means that it's critical to spend as little time here as possible. You
-    // can expect CLOCK_WATCHDOG_TIMEOUT bugchecks to happen otherwise. If you
-    // chose to enable interrupts note that this will result in further crashes
-    // as we are not on a correct OS stack, and you will be hitting crashes if
-    // RtlpCheckStackLimits is ever called, or if PatchGuard validates the RSP
-    // value.
-    //
-    KeRaiseIrql(HIGH_LEVEL, &guestContext.GuestIrql);
-
-    //
     // Because we had to use RCX when calling RtlCaptureContext, its true value
     // was actually pushed on the stack right before the call. Go dig into the
     // stack to find it, and overwrite the bogus value that's there now.
@@ -326,17 +311,12 @@ ShvVmxEntryHandler (
     }
 
     //
-    // Restore the IRQL back to the original level
-    //
-    KeLowerIrql(guestContext.GuestIrql);
-
-    //
     // Restore the context to either ShvVmxResume, in which case the CPU's VMX
     // facility will do the "true" return back to the VM (but without restoring
     // GPRs, which is why we must do it here), or to the original guest's RIP,
     // which we use in case an exit was requested. In this case VMX must now be
     // off, and this will look like a longjmp to the original stack and RIP.
     //
-    RtlRestoreContext(Context, NULL);
+    ShvOsRestoreContext(Context);
 }
 
diff --git a/shvvp.c b/shvvp.c
index b729b94..27b0d88 100644
--- a/shvvp.c
+++ b/shvvp.c
@@ -106,7 +106,7 @@ ShvVpRestoreAfterLaunch (
     // And finally, restore the context, so that all register and stack
     // state is finally restored.
     //
-    RtlRestoreContext(&vpData->ContextFrame, NULL);
+    ShvOsRestoreContext(&vpData->ContextFrame);
 }
 
 VOID
@@ -141,12 +141,13 @@ ShvVpInitialize (
 }
 
 VOID
-ShvVpUninitialize (
-    VOID
+ShvVpUnloadCallback (
+    _In_ PSHV_CALLBACK_CONTEXT Context
     )
 {
     INT cpuInfo[4];
     PSHV_VP_DATA vpData;
+    UNREFERENCED_PARAMETER(Context);
 
     //
     // Send the magic shutdown instruction sequence. It will return in EAX:EBX
@@ -154,7 +155,7 @@ ShvVpUninitialize (
     //
     __cpuidex(cpuInfo, 0x41414141, 0x42424242);
     vpData = (PSHV_VP_DATA)((ULONG64)cpuInfo[0] << 32 | cpuInfo[1]);
-    MmFreeContiguousMemory(vpData);
+    ShvOsFreeContiguousAlignedMemory(vpData);
 
     //
     // The processor will return here after the hypervisor issues a VMXOFF
@@ -181,31 +182,17 @@ ShvVpAllocateData (
     VOID
     )
 {
-    PHYSICAL_ADDRESS lowest, highest;
     PSHV_VP_DATA data;
 
     //
-    // The entire address range is OK for this allocation
-    //
-    lowest.QuadPart = 0;
-    highest.QuadPart = lowest.QuadPart - 1;
-
-    //
-    // Allocate a contiguous chunk of RAM to back this allocation and make sure
-    // that it is RW only, instead of RWX, by using the new Windows 8 API.
-    //
-    data = MmAllocateContiguousNodeMemory(sizeof(SHV_VP_DATA),
-                                          lowest,
-                                          highest,
-                                          lowest,
-                                          PAGE_READWRITE,
-                                          KeGetCurrentNodeNumber());
+    // Allocate a contiguous chunk of RAM to back this allocation
+    data = ShvOsAllocateContigousAlignedMemory(sizeof(*data));
     if (data != NULL)
     {
         //
         // Zero out the entire data region
         //
-        __stosq((PULONG64)data, 0, sizeof(SHV_VP_DATA) / sizeof(ULONG64));
+        __stosq((PULONG64)data, 0, sizeof(*data) / sizeof(ULONG64));
     }
 
     //
@@ -215,17 +202,13 @@ ShvVpAllocateData (
 }
 
 VOID
-ShvVpCallbackDpc (
-    _In_ PRKDPC Dpc,
-    _In_opt_ PVOID Context,
-    _In_opt_ PVOID SystemArgument1,
-    _In_opt_ PVOID SystemArgument2
+ShvVpLoadCallback (
+    _In_ PSHV_CALLBACK_CONTEXT Context
     )
 {
-    PSHV_DPC_CONTEXT dpcContext = Context;
     ULONG cpuIndex;
     PSHV_VP_DATA vpData;
-    UNREFERENCED_PARAMETER(Dpc);
+    NTSTATUS status;
 
     //
     // Detect if the hardware appears to support VMX root mode to start.
@@ -233,75 +216,58 @@ ShvVpCallbackDpc (
     //
     if (!ShvVmxProbe())
     {
-        dpcContext->FailureStatus = STATUS_HV_FEATURE_UNAVAILABLE;
-        goto Quickie;
+        status = STATUS_HV_FEATURE_UNAVAILABLE;
+        goto Failure;
     }
 
     //
-    // Check if we are loading, or unloading, and which CPU this is
+    // Allocate the per-VP data for this logical processor
     //
-    cpuIndex = KeGetCurrentProcessorNumberEx(NULL);
-    if (dpcContext->Cr3 != 0)
+    vpData = ShvVpAllocateData();
+    if (vpData == NULL)
     {
-        //
-        // Allocate the per-VP data for this logical processor
-        //
-        vpData = ShvVpAllocateData();
-        if (vpData == NULL)
-        {
-            dpcContext->FailureStatus = STATUS_HV_NO_RESOURCES;
-            goto Quickie;
-        }
-
-        //
-        // First, capture the value of the PML4 for the SYSTEM process, so that
-        // all virtual processors, regardless of which process the current LP
-        // has interrupted, can share the correct kernel address space.
-        //
-        vpData->SystemDirectoryTableBase = dpcContext->Cr3;
+        status = STATUS_HV_NO_RESOURCES;
+        goto Failure;
+    }
 
-        //
-        // Initialize the virtual processor
-        //
-        ShvVpInitialize(vpData);
+    //
+    // First, capture the value of the PML4 for the SYSTEM process, so that all
+    // virtual processors, regardless of which process the current LP has
+    // interrupted, can share the correct kernel address space.
+    //
+    vpData->SystemDirectoryTableBase = Context->Cr3;
 
-        //
-        // Our hypervisor should now be seen as present on this LP,
-        // as the SHV correctly handles CPUID ECX features register.
-        //
-        if (ShvIsOurHypervisorPresent() == FALSE)
-        {
-            //
-            // Free the per-processor data
-            //
-            MmFreeContiguousMemory(vpData);
-            dpcContext->FailureStatus = STATUS_HV_NOT_PRESENT;
-            dpcContext->FailedCpu = cpuIndex;
-            goto Quickie;
-        }
+    //
+    // Initialize the virtual processor
+    //
+    ShvVpInitialize(vpData);
 
-        //
-        // This CPU is hyperjacked!
-        //
-        InterlockedIncrement((PLONG)&dpcContext->InitCount);
-    }
-    else
+    //
+    // Our hypervisor should now be seen as present on this LP, as the SHV
+    // correctly handles CPUID ECX features register.
+    //
+    if (ShvIsOurHypervisorPresent() == FALSE)
     {
         //
-        // Tear down the virtual processor
+        // Free the per-processor data
         //
-        ShvVpUninitialize();
-        NT_ASSERT(ShvIsOurHypervisorPresent() == FALSE);
+        ShvOsFreeContiguousAlignedMemory(vpData);
+        status = STATUS_HV_NOT_PRESENT;
+        goto Failure;
     }
 
-Quickie:
     //
-    // Wait for all DPCs to synchronize at this point
+    // This CPU is hyperjacked!
     //
-    KeSignalCallDpcSynchronize(SystemArgument2);
+    _InterlockedIncrement((PLONG)&Context->InitCount);
+    return;
 
+Failure:
     //
-    // Mark the DPC as being complete
+    // Return failure
     //
-    KeSignalCallDpcDone(SystemArgument1);
+    cpuIndex = KeGetCurrentProcessorNumberEx(NULL);
+    Context->FailedCpu = cpuIndex;
+    Context->FailureStatus = status;
+    return;
 }