[reactos] 01/01: [NTOS:MM] Fix VADs being inserted even though the quota would exceed

[Tuur Martens]

https://git.reactos.org/?p=reactos.git;a=commitdiff;h=10126e7710b3d3ae1005e…

commit 10126e7710b3d3ae1005ecfefd71f0cfb1573514
Author:     Tuur Martens <tuurmartens4(a)gmail.com>
AuthorDate: Wed Jun 29 10:49:05 2022 +0200
Commit:     George Bișoc <george.bisoc(a)reactos.org>
CommitDate: Wed Jul 6 18:48:32 2022 +0200

    [NTOS:MM] Fix VADs being inserted even though the quota would exceed
    
    Since we were charging the pool quota after the VAD insertion,
    if the quota charge failed, the VAD would still have been inserted.
    This commit attempts to resolve this issue by charging quota
    before inserting the VAD thus allowing the quota charge to fail early.
    
    Addendum to 884356a0. CORE-18028
---
 ntoskrnl/mm/ARM3/mdlsup.c  | 14 ++++++++----
 ntoskrnl/mm/ARM3/procsup.c | 57 ++++++++++++++++++++++++++++------------------
 ntoskrnl/mm/ARM3/virtual.c | 25 ++++++++++++--------
 3 files changed, 59 insertions(+), 37 deletions(-)

diff --git a/ntoskrnl/mm/ARM3/mdlsup.c b/ntoskrnl/mm/ARM3/mdlsup.c
index 0d74949cd16..cc6d2ffb4ec 100644
--- a/ntoskrnl/mm/ARM3/mdlsup.c
+++ b/ntoskrnl/mm/ARM3/mdlsup.c
@@ -74,10 +74,18 @@ MiMapLockedPagesInUserSpace(
         DPRINT1("FIXME: Need to check for large pages\n");
     }
 
+    Status = PsChargeProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
+    if (!NT_SUCCESS(Status))
+    {
+        Vad = NULL;
+        goto Error;
+    }
+
     /* Allocate a VAD for our mapped region */
     Vad = ExAllocatePoolWithTag(NonPagedPool, sizeof(MMVAD_LONG), 'ldaV');
     if (Vad == NULL)
     {
+        PsReturnProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
         Status = STATUS_INSUFFICIENT_RESOURCES;
         goto Error;
     }
@@ -155,13 +163,8 @@ MiMapLockedPagesInUserSpace(
     MiLockProcessWorkingSetUnsafe(Process, Thread);
 
     ASSERT(Vad->EndingVpn >= Vad->StartingVpn);
-
     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)
     {
@@ -279,6 +282,7 @@ Error:
     if (Vad != NULL)
     {
         ExFreePoolWithTag(Vad, 'ldaV');
+        PsReturnProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
     }
     ExRaiseStatus(Status);
 }
diff --git a/ntoskrnl/mm/ARM3/procsup.c b/ntoskrnl/mm/ARM3/procsup.c
index 4111d711e3c..f2a93ce6fa5 100644
--- a/ntoskrnl/mm/ARM3/procsup.c
+++ b/ntoskrnl/mm/ARM3/procsup.c
@@ -35,9 +35,17 @@ MiCreatePebOrTeb(IN PEPROCESS Process,
     ULONG AlignedSize;
     LARGE_INTEGER CurrentTime;
 
+    Status = PsChargeProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
+    if (!NT_SUCCESS(Status))
+        return Status;
+
     /* Allocate a VAD */
     Vad = ExAllocatePoolWithTag(NonPagedPool, sizeof(MMVAD_LONG), 'ldaV');
-    if (!Vad) return STATUS_NO_MEMORY;
+    if (!Vad)
+    {
+        Status = STATUS_NO_MEMORY;
+        goto FailPath;
+    }
 
     /* Setup the primary flags with the size, and make it commited, private, RW */
     Vad->u.LongFlags = 0;
@@ -94,18 +102,18 @@ MiCreatePebOrTeb(IN PEPROCESS Process,
     if (!NT_SUCCESS(Status))
     {
         ExFreePoolWithTag(Vad, 'ldaV');
-        return STATUS_NO_MEMORY;
+        Status = STATUS_NO_MEMORY;
+        goto FailPath;
     }
 
-    Status = PsChargeProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
-    if (!NT_SUCCESS(Status))
-    {
-        ExFreePoolWithTag(Vad, 'ldaV');
-        return Status;
-    }
 
     /* Success */
     return STATUS_SUCCESS;
+
+FailPath:
+    PsReturnProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
+
+    return Status;
 }
 
 VOID
@@ -859,12 +867,23 @@ MiInsertSharedUserPageVad(
     ULONG_PTR BaseAddress;
     NTSTATUS Status;
 
+    if (Process->QuotaBlock != NULL)
+    {
+        Status = PsChargeProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("Ran out of quota.\n");
+            return Status;
+        }
+    }
+
     /* Allocate a VAD */
     Vad = ExAllocatePoolWithTag(NonPagedPool, sizeof(MMVAD_LONG), 'ldaV');
     if (Vad == NULL)
     {
         DPRINT1("Failed to allocate VAD for shared user page\n");
-        return STATUS_INSUFFICIENT_RESOURCES;
+        Status = STATUS_INSUFFICIENT_RESOURCES;
+        goto FailPath;
     }
 
     /* Setup the primary flags with the size, and make it private, RO */
@@ -898,23 +917,17 @@ MiInsertSharedUserPageVad(
     {
         DPRINT1("Failed to insert shared user VAD\n");
         ExFreePoolWithTag(Vad, 'ldaV');
-        return Status;
+        goto FailPath;
     }
 
-    if (Process->QuotaBlock != NULL)
-    {
-        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;
+
+FailPath:
+    if (Process->QuotaBlock != NULL)
+        PsReturnProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
+
+    return Status;
 }
 #endif
 
diff --git a/ntoskrnl/mm/ARM3/virtual.c b/ntoskrnl/mm/ARM3/virtual.c
index 5b54d8b4a28..e0804520435 100644
--- a/ntoskrnl/mm/ARM3/virtual.c
+++ b/ntoskrnl/mm/ARM3/virtual.c
@@ -4503,7 +4503,7 @@ NtAllocateVirtualMemory(IN HANDLE ProcessHandle,
     PETHREAD CurrentThread = PsGetCurrentThread();
     KAPC_STATE ApcState;
     ULONG ProtectionMask, QuotaCharge = 0, QuotaFree = 0;
-    BOOLEAN Attached = FALSE, ChangeProtection = FALSE;
+    BOOLEAN Attached = FALSE, ChangeProtection = FALSE, QuotaCharged = FALSE;
     MMPTE TempPte;
     PMMPTE PointerPte, LastPte;
     PMMPDE PointerPde;
@@ -4763,6 +4763,16 @@ NtAllocateVirtualMemory(IN HANDLE ProcessHandle,
             StartingAddress = (ULONG_PTR)PBaseAddress;
         }
 
+        // Charge quotas for the VAD
+        Status = PsChargeProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("Quota exceeded.\n");
+            goto FailPathNoLock;
+        }
+
+        QuotaCharged = TRUE;
+
         //
         // Allocate and initialize the VAD
         //
@@ -4796,15 +4806,6 @@ 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
@@ -5207,6 +5208,10 @@ FailPathNoLock:
         }
         _SEH2_END;
     }
+    else if (QuotaCharged)
+    {
+        PsReturnProcessNonPagedPoolQuota(Process, sizeof(MMVAD_LONG));
+    }
 
     return Status;
 }