[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

@@ -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);
 }

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
 

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;
 }