From 5941251f671a1cce282af7647e3c3fa87b4106c0 Mon Sep 17 00:00:00 2001
From: xerox <xerox@back.engineering>
Date: Wed, 28 Oct 2020 18:36:05 -0700
Subject: [PATCH] dynamically find an empty pml4e index to insert into....

---
 nasa-tables/mem_ctx/mem_ctx.cpp | 57 +++++++--------------------------
 nasa-tables/mem_ctx/mem_ctx.hpp |  4 +--
 2 files changed, 12 insertions(+), 49 deletions(-)

diff --git a/nasa-tables/mem_ctx/mem_ctx.cpp b/nasa-tables/mem_ctx/mem_ctx.cpp
index 4bf3a68..05b1f6e 100644
--- a/nasa-tables/mem_ctx/mem_ctx.cpp
+++ b/nasa-tables/mem_ctx/mem_ctx.cpp
@@ -8,9 +8,12 @@ namespace nasa
 		dirbase(get_dirbase(krnl_ctx, pid)),
 		pid(pid)
 	{
-		//
+		// find an empty pml4e...
+		for (auto idx = 100u; idx > 0u; --idx)
+			if (!k_ctx->rkm<pml4e>(k_ctx->get_virtual((reinterpret_cast<::ppml4e>(get_dirbase()) + idx))).present)
+				this->pml4e_index = idx;
+
 		// allocate a pdpt
-		//
 		this->new_pdpt.second = 
 			reinterpret_cast<ppdpte>(
 				VirtualAlloc(
@@ -21,30 +24,18 @@ namespace nasa
 				));
 		PAGE_IN(this->new_pdpt.second, PAGE_SIZE);
 
-		//
 		// get page table entries for new pdpt
-		//
 		pt_entries new_pdpt_entries;
-		hyperspace_entries(
-			new_pdpt_entries,
-			new_pdpt.second
-		);
+		hyperspace_entries(new_pdpt_entries, new_pdpt.second);
 		this->new_pdpt.first = reinterpret_cast<ppdpte>(new_pdpt_entries.pt.second.pfn << 12);
 
-		//
 		// make a new pml4e that points to our new pdpt.
-		//
 		new_pdpt_entries.pml4.second.pfn = new_pdpt_entries.pt.second.pfn;
 
-
-		//
 		// set the pml4e to point to the new pdpt
-		//
-		set_pml4e(reinterpret_cast<::ppml4e>(get_dirbase()) + PML4E_INDEX, new_pdpt_entries.pml4.second, true);
+		set_pml4e(reinterpret_cast<::ppml4e>(get_dirbase()) + this->pml4e_index, new_pdpt_entries.pml4.second, true);
 
-		//
 		// make a new pd
-		//
 		this->new_pd.second = 
 			reinterpret_cast<ppde>(
 				VirtualAlloc(
@@ -78,31 +69,22 @@ namespace nasa
 				));
 		PAGE_IN(this->new_pt.second, PAGE_SIZE);
 
-		//
 		// get paging table entries for pt
-		//
 		pt_entries new_pt_entries;
-		hyperspace_entries(
-			new_pt_entries,
-			this->new_pt.second
-		);
+		hyperspace_entries(new_pt_entries, this->new_pt.second);
 		this->new_pt.first = reinterpret_cast<ppte>(new_pt_entries.pt.second.pfn << 12);
 	}
 
 	mem_ctx::~mem_ctx()
 	{
-		//
 		// remove pml4e
-		//
 		pml4e null_value{ NULL };
-		set_pml4e(reinterpret_cast<::ppml4e>(get_dirbase()) + PML4E_INDEX, null_value, true);
+		set_pml4e(reinterpret_cast<::ppml4e>(get_dirbase()) + this->pml4e_index, null_value, true);
 	}
 
 	void* mem_ctx::set_page(void* addr)
 	{
-		//
 		// table entry change.
-		//
 		{
 			++pte_index;
 			if (pte_index >= 511)
@@ -128,9 +110,7 @@ namespace nasa
 		new_pdpte.user_supervisor = true;
 		new_pdpte.accessed = true;
 
-		//
 		// set pdpte entry
-		//
 		*reinterpret_cast<pdpte*>(new_pdpt.second + pdpte_index) = new_pdpte;
 
 		pde new_pde = { NULL };
@@ -140,9 +120,7 @@ namespace nasa
 		new_pde.user_supervisor = true;
 		new_pde.accessed = true;
 
-		//
 		// set pde entry
-		//
 		*reinterpret_cast<pde*>(new_pd.second + pde_index) = new_pde;
 
 		pte new_pte = { NULL };
@@ -152,25 +130,19 @@ namespace nasa
 		new_pte.user_supervisor = true;
 		new_pte.accessed = true;
 
-		//
 		// set pte entry
-		//
 		*reinterpret_cast<pte*>(new_pt.second + pte_index) = new_pte;
 
-		//
 		// set page offset
-		//
 		this->page_offset = virt_addr_t{ addr }.offset;
 		return get_page();
 	}
 
 	void* mem_ctx::get_page() const
 	{
-		//
 		// builds a new address given the state of all table indexes
-		//
 		virt_addr_t new_addr;
-		new_addr.pml4_index = PML4E_INDEX;
+		new_addr.pml4_index = this->pml4e_index;
 		new_addr.pdpt_index = this->pdpte_index;
 		new_addr.pd_index = this->pde_index;
 		new_addr.pt_index = this->pte_index;
@@ -180,18 +152,11 @@ namespace nasa
 
 	void* mem_ctx::get_dirbase(kernel_ctx& k_ctx, DWORD pid)
 	{
-		if (!pid)
-			return NULL;
-
 		const auto peproc =
 			reinterpret_cast<std::uint64_t>(k_ctx.get_peprocess(pid));
 
-		if (!peproc)
-			return NULL;
-
 		pte dirbase = k_ctx.rkm<pte>(
-			reinterpret_cast<void*>(peproc + 0x28)
-		);
+			reinterpret_cast<void*>(peproc + 0x28));
 
 		return reinterpret_cast<void*>(dirbase.pfn << 12);
 	}
diff --git a/nasa-tables/mem_ctx/mem_ctx.hpp b/nasa-tables/mem_ctx/mem_ctx.hpp
index b99228c..ff0e792 100644
--- a/nasa-tables/mem_ctx/mem_ctx.hpp
+++ b/nasa-tables/mem_ctx/mem_ctx.hpp
@@ -3,8 +3,6 @@
 #include "../kernel_ctx/kernel_ctx.h"
 
 #define PAGE_IN(addr, size) memset(addr, NULL, size)
-#define PML4E_INDEX 57
-
 struct pt_entries
 {
 	std::pair<ppml4e, pml4e>	pml4;
@@ -112,7 +110,7 @@ namespace nasa
 		bool hyperspace_entries(pt_entries& entries, void* addr);
 		void* dirbase;
 		kernel_ctx* k_ctx;
-		std::uint16_t pde_index, pte_index, pdpte_index, page_offset;
+		std::uint16_t pml4e_index, pdpte_index, pde_index, pte_index, page_offset;
 
 		/// first == physical
 		/// second == virtual