[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
This commit is contained in:
Tuur Martens 2022-05-12 08:11:52 +02:00 committed by George Bișoc
parent 5ee633f7d7
commit 884356a06e
4 changed files with 61 additions and 6 deletions

View file

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

View file

@ -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 */

View file

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

View file

@ -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
//