Handle non-paged pool overflow

xtoskrnl/mm/i686/mmgr.cc

@@ -396,25 +396,34 @@ MM::Manager::InitializeMemoryLayout(VOID)
     /* Compute system PTE size */
     ComputeSystemPteSize(&NumberOfSystemPtes);
 
-    /* Compute non-paged pool size */
+    /* Compute the initial and maximum non-paged pool sizes */
     ComputeNonPagedPoolSize(&MemoryLayout.NonPagedPoolSize);
     ComputeMaximumNonPagedPoolSize(&MaximumNonPagedPoolSize);
 
     /* Compute paged pool size */
     ComputePagedPoolSize(&MemoryLayout.PagedPoolSize);
 
-    /* Insert the PFN database right after the loader mappings */
+    /* Position the PFN database right after the loader mappings */
     MemoryLayout.PfnDatabase = (PMMPFN)MemoryLayout.LoaderMappingsEnd;
 
     /* Compute the PFN database end address */
     PfnDatabaseEnd = (ULONG_PTR)MemoryLayout.PfnDatabase + (MemoryLayout.PfnDatabaseSize * MM_PAGE_SIZE);
 
-    /* Check in non-paged pool fits before session space */
+    /* Position the initial non-paged pool immediately after the PFN database */
-    if(MemoryLayout.NonPagedPoolSize * MM_PAGE_SIZE <= ((ULONG_PTR)MemoryLayout.SessionSpaceStart - PfnDatabaseEnd))
-    {
-        /* Set non-paged pool start and end addresses */
     MemoryLayout.NonPagedPoolStart = (PVOID)PfnDatabaseEnd;
-        MemoryLayout.NonPagedPoolEnd = (PVOID)(PfnDatabaseEnd + MemoryLayout.NonPagedPoolSize * MM_PAGE_SIZE);
+
+    /* Check if the calculated non-paged pool size fits in the KVA */
+    if((MemoryLayout.NonPagedPoolSize * MM_PAGE_SIZE) >
+       ((ULONG_PTR)MemoryLayout.SessionSpaceStart - (ULONG_PTR)MemoryLayout.NonPagedPoolStart))
+    {
+        /* Set the final size for the non-paged pool */
+        MemoryLayout.NonPagedPoolSize = ((ULONG_PTR)MemoryLayout.NonPagedPoolEnd -
+                                        (ULONG_PTR)MemoryLayout.NonPagedPoolStart) / MM_PAGE_SIZE;
+    }
+
+    /* Set the final non-paged pool end address */
+    MemoryLayout.NonPagedPoolEnd = (PVOID)((ULONG_PTR)MemoryLayout.NonPagedPoolStart +
+                                           MemoryLayout.NonPagedPoolSize * MM_PAGE_SIZE);
 
     /* Check if non-paged expansion pool overflows */
     if((ULONG_PTR)MemoryLayout.NonPagedExpansionPoolStart + MaximumNonPagedPoolSize *
@@ -424,7 +433,7 @@ MM::Manager::InitializeMemoryLayout(VOID)
         if((ULONG_PTR)MemoryLayout.NonPagedExpansionPoolStart + MaximumNonPagedPoolSize *
             MM_PAGE_SIZE <= (ULONG_PTR)MemoryLayout.NonPagedExpansionPoolEnd)
         {
-                /* Set non-paged expansion pool end address */
+            /* Set new non-paged expansion pool end address */
             MemoryLayout.NonPagedExpansionPoolEnd = (PVOID)((ULONG_PTR)MemoryLayout.NonPagedExpansionPoolStart +
                                                             MaximumNonPagedPoolSize * MM_PAGE_SIZE);
         }
@@ -433,15 +442,6 @@ MM::Manager::InitializeMemoryLayout(VOID)
     /* Compute non-paged expansion pool size */
     MemoryLayout.NonPagedExpansionPoolSize = ((ULONG_PTR)MemoryLayout.NonPagedExpansionPoolEnd -
                                                 (ULONG_PTR)MemoryLayout.NonPagedExpansionPoolStart) / MM_PAGE_SIZE;
-    }
-    else
-    {
-        /* Unfortunally non-paged pool does not fit before session space. What can we do? */
-    }
-
-    /* Update paged pool end address */
-    MemoryLayout.PagedPoolEnd = (PVOID)(((ULONG_PTR)MemoryLayout.PagedPoolStart +
-                                         MemoryLayout.PagedPoolSize * MM_PAGE_SIZE) - 1);
 
     /* Dump memory layout */
     DumpMemoryLayout();