From 0ea1ae9643dbc73d35a22eea29de7e59737086a8 Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann@gmail.com>
Date: Sat, 2 Apr 2022 17:44:58 +0200
Subject: [PATCH] More refactoring

---
 src/driver/hypervisor.cpp  | 33 +++++++++++++++-------------
 src/driver/std_include.hpp |  2 ++
 src/driver/vmx.hpp         | 45 ++++++++------------------------------
 3 files changed, 29 insertions(+), 51 deletions(-)

diff --git a/src/driver/hypervisor.cpp b/src/driver/hypervisor.cpp
index 851e664..ba30eb7 100644
--- a/src/driver/hypervisor.cpp
+++ b/src/driver/hypervisor.cpp
@@ -99,16 +99,19 @@ void hypervisor::enable()
 {
 	const auto cr3 = __readcr3();
 
-	bool success = true;
+	volatile long failures = 0;
 	thread::dispatch_on_all_cores([&]()
 	{
-		success &= this->try_enable_core(cr3);
-	}, true);
+		if(!this->try_enable_core(cr3))
+		{
+			InterlockedIncrement(&failures);
+		}
+	});
 
-	if (!success)
+	if (failures)
 	{
 		this->disable();
-		//throw std::runtime_error("Hypervisor initialization failed");
+		throw std::runtime_error("Hypervisor initialization failed");
 	}
 }
 
@@ -277,9 +280,6 @@ ShvVmxMtrrAdjustEffectiveMemoryType(
 
 void ShvVmxEptInitialize(vmx::vm_state* VpData)
 {
-	UINT32 i, j;
-	vmx::pdpte tempEpdpte;
-
 	//
 	// Fill out the EPML4E which covers the first 512GB of RAM
 	//
@@ -292,14 +292,17 @@ void ShvVmxEptInitialize(vmx::vm_state* VpData)
 	//
 	// Fill out a RWX PDPTE
 	//
-	tempEpdpte.full = 0;
-	tempEpdpte.read = tempEpdpte.write = tempEpdpte.execute = 1;
+	epdpte temp_epdpte;
+	temp_epdpte.flags = 0;
+	temp_epdpte.read_access = 1;
+	temp_epdpte.write_access = 1;
+	temp_epdpte.execute_access = 1;
 
 	//
 	// Construct EPT identity map for every 1GB of RAM
 	//
-	__stosq((UINT64*)VpData->epdpt, tempEpdpte.full, PDPTE_ENTRY_COUNT);
-	for (i = 0; i < PDPTE_ENTRY_COUNT; i++)
+	__stosq((UINT64*)VpData->epdpt, temp_epdpte.flags, EPT_PDPTE_ENTRY_COUNT);
+	for (auto i = 0; i < EPT_PDPTE_ENTRY_COUNT; i++)
 	{
 		//
 		// Set the page frame number of the PDE table
@@ -320,13 +323,13 @@ void ShvVmxEptInitialize(vmx::vm_state* VpData)
 	//
 	// Loop every 1GB of RAM (described by the PDPTE)
 	//
-	__stosq((UINT64*)VpData->epde, temp_epde.flags, PDPTE_ENTRY_COUNT * PDE_ENTRY_COUNT);
-	for (i = 0; i < PDPTE_ENTRY_COUNT; i++)
+	__stosq((UINT64*)VpData->epde, temp_epde.flags, EPT_PDPTE_ENTRY_COUNT * EPT_PDE_ENTRY_COUNT);
+	for (auto i = 0; i < EPT_PDPTE_ENTRY_COUNT; i++)
 	{
 		//
 		// Construct EPT identity map for every 2MB of RAM
 		//
-		for (j = 0; j < PDE_ENTRY_COUNT; j++)
+		for (auto j = 0; j < EPT_PDE_ENTRY_COUNT; j++)
 		{
 			VpData->epde[i][j].page_frame_number = (i * 512) + j;
 			VpData->epde[i][j].memory_type = ShvVmxMtrrAdjustEffectiveMemoryType(VpData,
diff --git a/src/driver/std_include.hpp b/src/driver/std_include.hpp
index 8691c80..cff740a 100644
--- a/src/driver/std_include.hpp
+++ b/src/driver/std_include.hpp
@@ -3,6 +3,8 @@
 #include <ntddk.h>
 #include <intrin.h>
 
+#include <ia32.hpp>
+
 #include "stdint.hpp"
 #include "nt_ext.hpp"
 #include "new.hpp"
diff --git a/src/driver/vmx.hpp b/src/driver/vmx.hpp
index 76646fa..ab665fe 100644
--- a/src/driver/vmx.hpp
+++ b/src/driver/vmx.hpp
@@ -1,11 +1,5 @@
 #pragma once
 
-#include <ia32.hpp>
-
-#define PML4E_ENTRY_COUNT   512 // EPT_PML4E_ENTRY_COUNT
-#define PDPTE_ENTRY_COUNT   512 // EPT_PDPTE_ENTRY_COUNT
-#define PDE_ENTRY_COUNT     512 // EPT_PDE_ENTRY_COUNT
-
 namespace vmx
 {
 	struct vmcs
@@ -38,29 +32,6 @@ namespace vmx
 		};
 	};
 
-	struct pdpte
-	{
-		union
-		{
-			struct
-			{
-				uint64_t read : 1;
-				uint64_t write : 1;
-				uint64_t execute : 1;
-				uint64_t reserved : 5;
-				uint64_t accessed : 1;
-				uint64_t software_use : 1;
-				uint64_t user_mode_execute : 1;
-				uint64_t software_use2 : 1;
-				uint64_t page_frame_number : 36;
-				uint64_t reserved_high : 4;
-				uint64_t software_use_high : 12;
-			};
-
-			uint64_t full;
-		};
-	};
-
 	struct kdescriptor
 	{
 		uint16_t pad[3];
@@ -90,11 +61,13 @@ namespace vmx
 		uint64_t physical_address_max;
 	};
 
+#define DECLSPEC_PAGE_ALIGN DECLSPEC_ALIGN(PAGE_SIZE)
+
 	struct vm_state
 	{
 		union
 		{
-			DECLSPEC_ALIGN(PAGE_SIZE) uint8_t stack_buffer[KERNEL_STACK_SIZE]{};
+			DECLSPEC_PAGE_ALIGN uint8_t stack_buffer[KERNEL_STACK_SIZE]{};
 
 			struct
 			{
@@ -111,12 +84,12 @@ namespace vmx
 			};
 		};
 
-		DECLSPEC_ALIGN(PAGE_SIZE) uint8_t msr_bitmap[PAGE_SIZE]{};
-		DECLSPEC_ALIGN(PAGE_SIZE) epml4e epml4[PML4E_ENTRY_COUNT]{};
-		DECLSPEC_ALIGN(PAGE_SIZE) pdpte epdpt[PDPTE_ENTRY_COUNT]{};
-		DECLSPEC_ALIGN(PAGE_SIZE) epde_2mb epde[PDPTE_ENTRY_COUNT][PDE_ENTRY_COUNT]{};
+		DECLSPEC_PAGE_ALIGN uint8_t msr_bitmap[PAGE_SIZE]{};
+		DECLSPEC_PAGE_ALIGN epml4e epml4[EPT_PML4E_ENTRY_COUNT]{};
+		DECLSPEC_PAGE_ALIGN epdpte epdpt[EPT_PDPTE_ENTRY_COUNT]{};
+		DECLSPEC_PAGE_ALIGN epde_2mb epde[EPT_PDPTE_ENTRY_COUNT][EPT_PDE_ENTRY_COUNT]{};
 
-		DECLSPEC_ALIGN(PAGE_SIZE) vmcs vmx_on{};
-		DECLSPEC_ALIGN(PAGE_SIZE) vmcs vmcs{};
+		DECLSPEC_PAGE_ALIGN vmcs vmx_on{};
+		DECLSPEC_PAGE_ALIGN vmcs vmcs{};
 	};
 }