[ros-diffs] [reactos] 01/01: [NTOS:MM] Charge and free quotas for VAD allocations

[Tuur Martens]

https://git.reactos.org/?p=reactos.git;a=commitdiff;h=884356a06e1589271595e…

commit 884356a06e1589271595e198e0a32c35d29d6059
Author:     Tuur Martens <tuurmartens4(a)gmail.com>
AuthorDate: Thu May 12 08:11:52 2022 +0200
Commit:     George Bișoc <george.bisoc(a)reactos.org>
CommitDate: Sat May 14 15:09:50 2022 +0200

    [NTOS:MM] Charge and free quotas for VAD allocations
    
    Charge quotas for VAD allocations and free the quotas again when the VADs are freed.
    
    CORE-18028
---
 ntoskrnl/mm/ARM3/mdlsup.c  |  5 +++++
 ntoskrnl/mm/ARM3/procsup.c | 27 +++++++++++++++++++++++++--
 ntoskrnl/mm/ARM3/section.c |  9 +++++++++
 ntoskrnl/mm/ARM3/virtual.c | 26 ++++++++++++++++++++++----
 4 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/ntoskrnl/mm/ARM3/mdlsup.c b/ntoskrnl/mm/ARM3/mdlsup.c
index 9332ec3153e..0d74949cd16 100644
--- a/ntoskrnl/mm/ARM3/mdlsup.c
+++ b/ntoskrnl/mm/ARM3/mdlsup.c
@@ -158,6 +158,10 @@ MiMapLockedPagesInUserSpace(
 
     MiInsertVad((PMMVAD)Vad, &Process->VadRoot);
 
+    Status = PsChargeProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
+    if (!NT_SUCCESS(Status))
+        goto Error;
+
     /* Check if this is uncached */
     if (CacheAttribute != MiCached)
     {
@@ -319,6 +323,7 @@ MiUnmapLockedPagesInUserSpace(
     /* Remove it from the process VAD tree */
     ASSERT(Process->VadRoot.NumberGenericTableElements >= 1);
     MiRemoveNode((PMMADDRESS_NODE)Vad, &Process->VadRoot);
+    PsReturnProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
 
     /* MiRemoveNode should have removed us if we were the hint */
     ASSERT(Process->VadRoot.NodeHint != Vad);
diff --git a/ntoskrnl/mm/ARM3/procsup.c b/ntoskrnl/mm/ARM3/procsup.c
index dd04b68fb83..2ddd91c5ca2 100644
--- a/ntoskrnl/mm/ARM3/procsup.c
+++ b/ntoskrnl/mm/ARM3/procsup.c
@@ -97,6 +97,13 @@ MiCreatePebOrTeb(IN PEPROCESS Process,
         return STATUS_NO_MEMORY;
     }
 
+    Status = PsChargeProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
+    if (!NT_SUCCESS(Status))
+    {
+        ExFreePoolWithTag(Vad, 'ldaV');
+        return Status;
+    }
+
     /* Success */
     return STATUS_SUCCESS;
 }
@@ -154,6 +161,9 @@ MmDeleteTeb(IN PEPROCESS Process,
 
         /* Remove the VAD */
         ExFreePool(Vad);
+
+        /* Return the quota the VAD used */
+        PsReturnProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
     }
 
     /* Release the address space lock */
@@ -842,7 +852,8 @@ MmCreateTeb(IN PEPROCESS Process,
 #ifdef _M_AMD64
 static
 NTSTATUS
-MiInsertSharedUserPageVad(VOID)
+MiInsertSharedUserPageVad(
+    _In_ PEPROCESS Process)
 {
     PMMVAD_LONG Vad;
     ULONG_PTR BaseAddress;
@@ -890,6 +901,15 @@ MiInsertSharedUserPageVad(VOID)
         return Status;
     }
 
+    Status = PsChargeProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("Ran out of quota.\n");
+        ExFreePoolWithTag(Vad, 'ldaV');
+        return Status;
+    }
+
+
     /* Success */
     return STATUS_SUCCESS;
 }
@@ -1012,7 +1032,7 @@ MmInitializeProcessAddressSpace(IN PEPROCESS Process,
 
 #ifdef _M_AMD64
     /* On x64 we need a VAD for the shared user page */
-    Status = MiInsertSharedUserPageVad();
+    Status = MiInsertSharedUserPageVad(Process);
     if (!NT_SUCCESS(Status))
     {
         DPRINT1("MiCreateSharedUserPageVad() failed: 0x%lx\n", Status);
@@ -1300,6 +1320,9 @@ MmCleanProcessAddressSpace(IN PEPROCESS Process)
 
         /* Free the VAD memory */
         ExFreePool(Vad);
+
+        /* Return the quota the VAD used */
+        PsReturnProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
     }
 
     /* Lock the working set */
diff --git a/ntoskrnl/mm/ARM3/section.c b/ntoskrnl/mm/ARM3/section.c
index 4a65209d38b..40d942f7703 100644
--- a/ntoskrnl/mm/ARM3/section.c
+++ b/ntoskrnl/mm/ARM3/section.c
@@ -915,6 +915,7 @@ MiUnmapViewOfSection(IN PEPROCESS Process,
     /* Remove the VAD */
     ASSERT(Process->VadRoot.NumberGenericTableElements >= 1);
     MiRemoveNode((PMMADDRESS_NODE)Vad, &Process->VadRoot);
+    PsReturnProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
 
     /* Remove the PTEs for this view, which also releases the working set lock */
     MiRemoveMappedView(Process, Vad);
@@ -1489,6 +1490,13 @@ MiMapViewOfDataSection(IN PCONTROL_AREA ControlArea,
         StartAddress = 0;
     }
 
+    Status = PsChargeProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
+    if (!NT_SUCCESS(Status))
+    {
+        ExFreePoolWithTag(Vad, 'ldaV');
+        return Status;
+    }
+
     /* Insert the VAD */
     Status = MiInsertVadEx((PMMVAD)Vad,
                            &StartAddress,
@@ -1498,6 +1506,7 @@ MiMapViewOfDataSection(IN PCONTROL_AREA ControlArea,
                            AllocationType);
     if (!NT_SUCCESS(Status))
     {
+        PsReturnProcessNonPagedPoolQuota(PsGetCurrentProcess(), sizeof(MMVAD_LONG));
         return Status;
     }
 
diff --git a/ntoskrnl/mm/ARM3/virtual.c b/ntoskrnl/mm/ARM3/virtual.c
index f178edbb909..c5c5a5283b0 100644
--- a/ntoskrnl/mm/ARM3/virtual.c
+++ b/ntoskrnl/mm/ARM3/virtual.c
@@ -4790,6 +4790,15 @@ NtAllocateVirtualMemory(IN HANDLE ProcessHandle,
             goto FailPathNoLock;
         }
 
+        // Charge quotas for the VAD
+        Status = PsChargeProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("Quota exceeded.\n");
+            ExFreePoolWithTag(Vad, 'SdaV');
+            goto FailPathNoLock;
+        }
+
         //
         // Detach and dereference the target process if
         // it was different from the current process
@@ -5416,6 +5425,7 @@ NtFreeVirtualMemory(IN HANDLE ProcessHandle,
             MiLockProcessWorkingSetUnsafe(Process, CurrentThread);
             ASSERT(Process->VadRoot.NumberGenericTableElements >= 1);
             MiRemoveNode((PMMADDRESS_NODE)Vad, &Process->VadRoot);
+            PsReturnProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
         }
         else
         {
@@ -5448,6 +5458,7 @@ NtFreeVirtualMemory(IN HANDLE ProcessHandle,
                     MiLockProcessWorkingSetUnsafe(Process, CurrentThread);
                     ASSERT(Process->VadRoot.NumberGenericTableElements >= 1);
                     MiRemoveNode((PMMADDRESS_NODE)Vad, &Process->VadRoot);
+                    PsReturnProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
                 }
                 else
                 {
@@ -5522,6 +5533,17 @@ NtFreeVirtualMemory(IN HANDLE ProcessHandle,
                         goto FailPath;
                     }
 
+                    // Charge quota for the new VAD
+                    Status = PsChargeProcessNonPagedPoolQuota(Process,
sizeof(MMVAD_LONG));
+
+                    if (!NT_SUCCESS(Status))
+                    {
+                        DPRINT1("Ran out of process quota whilst creating new
VAD!\n");
+                        ExFreePoolWithTag(NewVad, 'SdaV');
+                        Status = STATUS_QUOTA_EXCEEDED;
+                        goto FailPath;
+                    }
+
                     //
                     // This new VAD describes the second chunk, so we keep the end
                     // address of the original and adjust the start to point past
@@ -5534,10 +5556,6 @@ NtFreeVirtualMemory(IN HANDLE ProcessHandle,
                     NewVad->u.VadFlags.CommitCharge = 0;
                     ASSERT(NewVad->EndingVpn >= NewVad->StartingVpn);
 
-                    //
-                    // TODO: charge quota for the new VAD
-                    //
-
                     //
                     // Get the commit charge for the released region
                     //