From bee91d0c71dba813d2b29a9801f187994f644fec Mon Sep 17 00:00:00 2001
From: Aiken Harris <harraiken91@gmail.com>
Date: Tue, 6 Jan 2026 14:49:30 +0100
Subject: [PATCH] Correctly setup PFN database for ROM and in-use pages

---
 xtoskrnl/mm/i686/pfn.cc | 55 ++++++++++++++++++++++++++---------------
 1 file changed, 35 insertions(+), 20 deletions(-)

diff --git a/xtoskrnl/mm/i686/pfn.cc b/xtoskrnl/mm/i686/pfn.cc
index 89403a7..c44949c 100644
--- a/xtoskrnl/mm/i686/pfn.cc
+++ b/xtoskrnl/mm/i686/pfn.cc
@@ -177,7 +177,8 @@ MM::Pfn::ProcessMemoryDescriptor(IN PFN_NUMBER BasePage,
                                  IN PFN_NUMBER PageCount,
                                  IN LOADER_MEMORY_TYPE MemoryType)
 {
-    PFN_NUMBER CurrentPage, PageNumber;
+    PFN_NUMBER PageNumber;
+    PMMPDE PointerPde;
     PMMPFN Pfn;
 
     /* Check if the memory descriptor describes a free memory region */
@@ -186,13 +187,8 @@ MM::Pfn::ProcessMemoryDescriptor(IN PFN_NUMBER BasePage,
         /* Iterate over each page in this free memory run */
         for(PageNumber = 0; PageNumber < PageCount; PageNumber++)
         {
-            /* Get the PFN entry for the current page and set its initial cache attribute */
-            CurrentPage = BasePage + PageNumber;
-            Pfn = GetPfnEntry(CurrentPage);
-            Pfn->u3.e1.CacheAttribute = PfnNonCached;
-
             /* Add the page to the free list to make it available for allocation */
-            LinkFreePage(CurrentPage);
+            LinkFreePage(BasePage + PageNumber);
         }
     }
     else
@@ -215,17 +211,24 @@ MM::Pfn::ProcessMemoryDescriptor(IN PFN_NUMBER BasePage,
                     /* Get the PFN entry for the current ROM page */
                     Pfn = GetPfnEntry(BasePage + PageNumber);
 
-                    /* Initialize the PFN entry to represent a ROM page */
-                    Pfn->PteAddress = 0;
-                    Pfn->u1.Flink = 0;
-                    Pfn->u2.ShareCount = 0;
-                    Pfn->u3.e1.CacheAttribute = PfnNonCached;
-                    Pfn->u3.e1.PageLocation = 0;
-                    Pfn->u3.e1.PrototypePte = 1;
-                    Pfn->u3.e1.Rom = 1;
-                    Pfn->u3.e2.ReferenceCount = 0;
-                    Pfn->u4.InPageError = 0;
-                    Pfn->u4.PteFrame = 0;
+                    /* Ensure that the page is not already in-use */
+                    if(Pfn->u3.e2.ReferenceCount == 0)
+                    {
+                        /* Get the page directory entry for the current page */
+                        PointerPde = MM::Paging::GetPdeAddress((PVOID)(KSEG0_BASE + (BasePage << MM_PAGE_SHIFT)));
+
+                        /* Initialize the PFN entry to represent a ROM page */
+                        Pfn->PteAddress = MM::Paging::GetPteAddress((PVOID)(KSEG0_BASE + (BasePage << MM_PAGE_SHIFT)));
+                        Pfn->u1.Flink = 0;
+                        Pfn->u2.ShareCount = 0;
+                        Pfn->u3.e1.CacheAttribute = PfnNonCached;
+                        Pfn->u3.e1.PageLocation = 0;
+                        Pfn->u3.e1.PrototypePte = 1;
+                        Pfn->u3.e1.Rom = 1;
+                        Pfn->u3.e2.ReferenceCount = 0;
+                        Pfn->u4.InPageError = 0;
+                        Pfn->u4.PteFrame = MM::Paging::GetPageFrameNumber(PointerPde);
+                    }
                 }
                 break;
             default:
@@ -235,8 +238,20 @@ MM::Pfn::ProcessMemoryDescriptor(IN PFN_NUMBER BasePage,
                     /* Get the PFN entry for the current in-use page */
                     Pfn = GetPfnEntry(BasePage + PageNumber);
 
-                    /* Mark the PFN as active and valid to prevent it from being allocated */
-                    Pfn->u3.e1.PageLocation = ActiveAndValid;
+                    /* Ensure that the page is not already in-use */
+                    if(Pfn->u3.e2.ReferenceCount == 0)
+                    {
+                        /* Get the page directory entry for the current page */
+                        PointerPde = MM::Paging::GetPdeAddress((PVOID)(KSEG0_BASE + (BasePage << MM_PAGE_SHIFT)));
+
+                        /* Initialize the PFN entry to represent an in-use page and prevent it from being allocated */
+                        Pfn->PteAddress = MM::Paging::GetPteAddress((PVOID)(KSEG0_BASE + (BasePage << MM_PAGE_SHIFT)));
+                        Pfn->u2.ShareCount++;
+                        Pfn->u3.e1.CacheAttribute = PfnNonCached;
+                        Pfn->u3.e1.PageLocation = ActiveAndValid;
+                        Pfn->u3.e2.ReferenceCount = 1;
+                        Pfn->u4.PteFrame = MM::Paging::GetPageFrameNumber(PointerPde);
+                    }
                 }
                 break;
         }