Update PTE mapping to allocate from PFN database

xtoskrnl/mm/amd64/pte.cc

@@ -176,6 +176,7 @@ MM::Pte::MapP5E(IN PVOID StartAddress,
                 IN PMMP5E TemplateP5e)
 {
     PMMP5E EndSpace, PointerP5e;
+    PFN_NUMBER PageFrameNumber;
 
     /* Get P5E addresses */
     PointerP5e = MM::Paging::GetP5eAddress(StartAddress);
@@ -187,8 +188,18 @@ MM::Pte::MapP5E(IN PVOID StartAddress,
         /* Check if P5E is already mapped */
         if(!MM::Paging::PteValid(PointerP5e))
         {
-            /* Map P5E */
+            /* Attempt to allocate a physical page from the PFN database */
-            MM::Paging::SetPte(TemplateP5e, MM::Pfn::AllocateBootstrapPages(1), 0);
+            PageFrameNumber = MM::Pfn::AllocatePhysicalPage(0);
+
+            /* Check if the primary allocation failed */
+            if(!PageFrameNumber)
+            {
+                /* Allocate a physical page from the fallback allocator */
+                PageFrameNumber = MM::Pfn::AllocateBootstrapPages(1);
+            }
+
+            /* Update the template with new page frame and write the PTE */
+            MM::Paging::SetPte(TemplateP5e, PageFrameNumber, 0);
             *PointerP5e = *TemplateP5e;
 
             /* Clear the page table */
@@ -223,6 +234,7 @@ MM::Pte::MapPPE(IN PVOID StartAddress,
                 IN PMMPPE TemplatePpe)
 {
     PMMPPE EndSpace, PointerPpe;
+    PFN_NUMBER PageFrameNumber;
 
     /* Get PPE addresses */
     PointerPpe = MM::Paging::GetPpeAddress(StartAddress);
@@ -234,8 +246,18 @@ MM::Pte::MapPPE(IN PVOID StartAddress,
         /* Check if PPE is already mapped */
         if(!MM::Paging::PteValid(PointerPpe))
         {
-            /* Map PPE */
+            /* Attempt to allocate a physical page from the PFN database */
-            MM::Paging::SetPte(TemplatePpe, MM::Pfn::AllocateBootstrapPages(1), 0);
+            PageFrameNumber = MM::Pfn::AllocatePhysicalPage(0);
+
+            /* Check if the primary allocation failed */
+            if(!PageFrameNumber)
+            {
+                /* Allocate a physical page from the fallback allocator */
+                PageFrameNumber = MM::Pfn::AllocateBootstrapPages(1);
+            }
+
+            /* Update the template with new page frame and write the PTE */
+            MM::Paging::SetPte(TemplatePpe, PageFrameNumber, 0);
             *PointerPpe = *TemplatePpe;
 
             /* Clear the page table */
@@ -270,6 +292,7 @@ MM::Pte::MapPXE(IN PVOID StartAddress,
                 IN PMMPXE TemplatePxe)
 {
     PMMPXE EndSpace, PointerPxe;
+    PFN_NUMBER PageFrameNumber;
 
     /* Get PXE addresses */
     PointerPxe = MM::Paging::GetPxeAddress(StartAddress);
@@ -281,8 +304,18 @@ MM::Pte::MapPXE(IN PVOID StartAddress,
         /* Check if PTE is already mapped */
         if(!MM::Paging::PteValid(PointerPxe))
         {
-            /* Map PTE */
+            /* Attempt to allocate a physical page from the PFN database */
-            MM::Paging::SetPte(TemplatePxe, MM::Pfn::AllocateBootstrapPages(1), 0);
+            PageFrameNumber = MM::Pfn::AllocatePhysicalPage(0);
+
+            /* Check if the primary allocation failed */
+            if(!PageFrameNumber)
+            {
+                /* Allocate a physical page from the fallback allocator */
+                PageFrameNumber = MM::Pfn::AllocateBootstrapPages(1);
+            }
+
+            /* Update the template with new page frame and write the PTE */
+            MM::Paging::SetPte(TemplatePxe, PageFrameNumber, 0);
             *PointerPxe = *TemplatePxe;
 
             /* Clear the page table */

xtoskrnl/mm/mmgr.cc

@@ -256,6 +256,9 @@ MM::Manager::InitializeMemoryManager(VOID)
     /* Initialize page table */
     MM::Pte::InitializePageTable();
 
+    /* Initialize PFN database */
+    MM::Pfn::InitializePfnDatabase();
+
     /* Initialize system PTE space */
     MM::Pte::InitializeSystemPteSpace();
 
@@ -265,9 +268,6 @@ MM::Manager::InitializeMemoryManager(VOID)
     /* Initialize non-paged pool */
     MM::Pool::InitializeNonPagedPool();
 
-    /* Initialize PFN database */
-    MM::Pfn::InitializePfnDatabase();
-
     /* Initialize allocations tracking tables */
     MM::Allocator::InitializeAllocationsTracking();
     MM::Allocator::InitializeBigAllocationsTracking();

xtoskrnl/mm/pte.cc

@@ -242,6 +242,7 @@ MM::Pte::MapPDE(IN PVOID StartAddress,
                 IN PMMPDE TemplatePde)
 {
     PMMPDE EndSpace, PointerPde;
+    PFN_NUMBER PageFrameNumber;
 
     /* Get PDE addresses */
     PointerPde = MM::Paging::GetPdeAddress(StartAddress);
@@ -253,8 +254,18 @@ MM::Pte::MapPDE(IN PVOID StartAddress,
         /* Check if PDE is already mapped */
         if(!MM::Paging::PteValid(PointerPde))
         {
-            /* Map PDE */
+            /* Attempt to allocate a physical page from the PFN database */
-            MM::Paging::SetPte(TemplatePde, MM::Pfn::AllocateBootstrapPages(1), 0);
+            PageFrameNumber = MM::Pfn::AllocatePhysicalPage(0);
+
+            /* Check if the primary allocation failed */
+            if(!PageFrameNumber)
+            {
+                /* Allocate a physical page from the fallback allocator */
+                PageFrameNumber = MM::Pfn::AllocateBootstrapPages(1);
+            }
+
+            /* Update the template with new page frame and write the PTE */
+            MM::Paging::SetPte(TemplatePde, PageFrameNumber, 0);
             MM::Paging::WritePte(PointerPde, *TemplatePde);
 
             /* Clear the page table */
@@ -289,6 +300,7 @@ MM::Pte::MapPTE(IN PVOID StartAddress,
                 IN PMMPTE TemplatePte)
 {
     PMMPTE EndSpace, PointerPte;
+    PFN_NUMBER PageFrameNumber;
 
     /* Get PTE addresses */
     PointerPte = MM::Paging::GetPteAddress(StartAddress);
@@ -300,8 +312,18 @@ MM::Pte::MapPTE(IN PVOID StartAddress,
         /* Check if PTE is already mapped */
         if(!MM::Paging::PteValid(PointerPte))
         {
-            /* Map PTE */
+            /* Attempt to allocate a physical page from the PFN database */
-            MM::Paging::SetPte(TemplatePte, MM::Pfn::AllocateBootstrapPages(1), 0);
+            PageFrameNumber = MM::Pfn::AllocatePhysicalPage(0);
+
+            /* Check if the primary allocation failed */
+            if(!PageFrameNumber)
+            {
+                /* Allocate a physical page from the fallback allocator */
+                PageFrameNumber = MM::Pfn::AllocateBootstrapPages(1);
+            }
+
+            /* Update the template with new page frame and write the PTE */
+            MM::Paging::SetPte(TemplatePte, PageFrameNumber, 0);
             MM::Paging::WritePte(PointerPte, *TemplatePte);
 
             /* Clear the page table */
@@ -313,7 +335,6 @@ MM::Pte::MapPTE(IN PVOID StartAddress,
     }
 }
 
-
 /**
  * Releases a block of system PTEs into a specified pool.
  *