[reactos] 01/01: [NTOS:PS] Implement PS_QUOTA_TYPE and let the quota code use it (#3389)

[George Bișoc]

https://git.reactos.org/?p=reactos.git;a=commitdiff;h=6170b574f0fb00460da7a…

commit 6170b574f0fb00460da7a3e547cec7e9fa41b621
Author:     George Bișoc <george.bisoc(a)reactos.org>
AuthorDate: Tue Mar 2 20:09:58 2021 +0100
Commit:     GitHub <noreply(a)github.com>
CommitDate: Tue Mar 2 20:09:58 2021 +0100

    [NTOS:PS] Implement PS_QUOTA_TYPE and let the quota code use it (#3389)
    
    This will replace the PoolIndex variable and as such we'll only be using the
PS_QUOTA_TYPE enumeration, as Windows does. Both QuotaEntry, QuotaUsage and QuotaPeak
depend explicitly or implicitly on this enumeration. Further details about this enum can
be found in the following articles.
   
https://www.geoffchappell.com/studies/windows/km/ntoskrnl/api/ps/psquota/ty…
   
https://www.geoffchappell.com/studies/windows/km/ntoskrnl/api/ps/psquota/bl…
    https://www.geoffchappell.com/studies/windows/km/ntoskrnl/inc/ps/eprocess/i…
(see QuotaPeak and QuotaUsage)
---
 ntoskrnl/ex/sysinfo.c     | 12 +++++------
 ntoskrnl/ps/query.c       | 12 +++++------
 ntoskrnl/ps/quota.c       | 52 ++++++++++++++++++++++-------------------------
 sdk/include/ndk/pstypes.h | 23 ++++++++++++++++++---
 4 files changed, 56 insertions(+), 43 deletions(-)

diff --git a/ntoskrnl/ex/sysinfo.c b/ntoskrnl/ex/sysinfo.c
index ffb8ddc906d..f02f0a1cce1 100644
--- a/ntoskrnl/ex/sysinfo.c
+++ b/ntoskrnl/ex/sysinfo.c
@@ -1025,12 +1025,12 @@ QSI_DEF(SystemProcessInformation)
                 SpiCurrent->PageFaultCount = Process->Vm.PageFaultCount;
                 SpiCurrent->PeakWorkingSetSize = Process->Vm.PeakWorkingSetSize;
                 SpiCurrent->WorkingSetSize = Process->Vm.WorkingSetSize;
-                SpiCurrent->QuotaPeakPagedPoolUsage = Process->QuotaPeak[0];
-                SpiCurrent->QuotaPagedPoolUsage = Process->QuotaUsage[0];
-                SpiCurrent->QuotaPeakNonPagedPoolUsage = Process->QuotaPeak[1];
-                SpiCurrent->QuotaNonPagedPoolUsage = Process->QuotaUsage[1];
-                SpiCurrent->PagefileUsage = Process->QuotaUsage[2];
-                SpiCurrent->PeakPagefileUsage = Process->QuotaPeak[2];
+                SpiCurrent->QuotaPeakPagedPoolUsage =
Process->QuotaPeak[PsPagedPool];
+                SpiCurrent->QuotaPagedPoolUsage =
Process->QuotaUsage[PsPagedPool];
+                SpiCurrent->QuotaPeakNonPagedPoolUsage =
Process->QuotaPeak[PsNonPagedPool];
+                SpiCurrent->QuotaNonPagedPoolUsage =
Process->QuotaUsage[PsNonPagedPool];
+                SpiCurrent->PagefileUsage = Process->QuotaUsage[PsPageFile];
+                SpiCurrent->PeakPagefileUsage = Process->QuotaPeak[PsPageFile];
                 SpiCurrent->PrivatePageCount = Process->CommitCharge;
                 ThreadInfo = (PSYSTEM_THREAD_INFORMATION)(SpiCurrent + 1);
 
diff --git a/ntoskrnl/ps/query.c b/ntoskrnl/ps/query.c
index 60f440ab9bc..cad72c4d60a 100644
--- a/ntoskrnl/ps/query.c
+++ b/ntoskrnl/ps/query.c
@@ -460,12 +460,12 @@ NtQueryInformationProcess(IN HANDLE ProcessHandle,
                 VmCounters->PageFaultCount = Process->Vm.PageFaultCount;
                 VmCounters->PeakWorkingSetSize = Process->Vm.PeakWorkingSetSize;
                 VmCounters->WorkingSetSize = Process->Vm.WorkingSetSize;
-                VmCounters->QuotaPeakPagedPoolUsage = Process->QuotaPeak[0];
-                VmCounters->QuotaPagedPoolUsage = Process->QuotaUsage[0];
-                VmCounters->QuotaPeakNonPagedPoolUsage = Process->QuotaPeak[1];
-                VmCounters->QuotaNonPagedPoolUsage = Process->QuotaUsage[1];
-                VmCounters->PagefileUsage = Process->QuotaUsage[2] <<
PAGE_SHIFT;
-                VmCounters->PeakPagefileUsage = Process->QuotaPeak[2] <<
PAGE_SHIFT;
+                VmCounters->QuotaPeakPagedPoolUsage =
Process->QuotaPeak[PsPagedPool];
+                VmCounters->QuotaPagedPoolUsage =
Process->QuotaUsage[PsPagedPool];
+                VmCounters->QuotaPeakNonPagedPoolUsage =
Process->QuotaPeak[PsNonPagedPool];
+                VmCounters->QuotaNonPagedPoolUsage =
Process->QuotaUsage[PsNonPagedPool];
+                VmCounters->PagefileUsage = Process->QuotaUsage[PsPageFile]
<< PAGE_SHIFT;
+                VmCounters->PeakPagefileUsage = Process->QuotaPeak[PsPageFile]
<< PAGE_SHIFT;
                 //VmCounters->PrivateUsage = Process->CommitCharge <<
PAGE_SHIFT;
                 //
 
diff --git a/ntoskrnl/ps/quota.c b/ntoskrnl/ps/quota.c
index 6d3751090ca..0645654a320 100644
--- a/ntoskrnl/ps/quota.c
+++ b/ntoskrnl/ps/quota.c
@@ -28,38 +28,36 @@ static KSPIN_LOCK PspQuotaLock;
 
 /*
  * Private helper to charge the specified process quota.
- * ReturnsSTATUS_QUOTA_EXCEEDED on quota limit check failure.
- * Updates QuotaPeak as needed for specified PoolIndex.
- * TODO: Research and possibly add (the undocumented) enum type PS_QUOTA_TYPE
- *       to replace UCHAR for 'PoolIndex'.
+ * Returns STATUS_QUOTA_EXCEEDED on quota limit check failure.
+ * Updates QuotaPeak as needed for specified quota type in PS_QUOTA_TYPE enum.
  * Notes: Conceptually translation unit local/private.
  */
 NTSTATUS
 NTAPI
 PspChargeProcessQuotaSpecifiedPool(IN PEPROCESS Process,
-                                   IN UCHAR     PoolIndex,
+                                   IN PS_QUOTA_TYPE QuotaType,
                                    IN SIZE_T    Amount)
 {
     ASSERT(Process);
     ASSERT(Process != PsInitialSystemProcess);
-    ASSERT(PoolIndex <= 2);
+    ASSERT(QuotaType < PsQuotaTypes);
     ASSERT(Process->QuotaBlock);
 
     /* Note: Race warning. TODO: Needs to add/use lock for this */
-    if (Process->QuotaUsage[PoolIndex] + Amount >
-        Process->QuotaBlock->QuotaEntry[PoolIndex].Limit)
+    if (Process->QuotaUsage[QuotaType] + Amount >
+        Process->QuotaBlock->QuotaEntry[QuotaType].Limit)
     {
         DPRINT1("Quota exceeded, but ROS will let it slide...\n");
         return STATUS_SUCCESS;
         //return STATUS_QUOTA_EXCEEDED; /* caller raises the exception */
     }
 
-    InterlockedExchangeAdd((LONG*)&Process->QuotaUsage[PoolIndex], Amount);
+    InterlockedExchangeAdd((LONG*)&Process->QuotaUsage[QuotaType], Amount);
 
     /* Note: Race warning. TODO: Needs to add/use lock for this */
-    if (Process->QuotaPeak[PoolIndex] < Process->QuotaUsage[PoolIndex])
+    if (Process->QuotaPeak[QuotaType] < Process->QuotaUsage[QuotaType])
     {
-        Process->QuotaPeak[PoolIndex] = Process->QuotaUsage[PoolIndex];
+        Process->QuotaPeak[QuotaType] = Process->QuotaUsage[QuotaType];
     }
 
     return STATUS_SUCCESS;
@@ -67,27 +65,25 @@ PspChargeProcessQuotaSpecifiedPool(IN PEPROCESS Process,
 
 /*
  * Private helper to remove quota charge from the specified process quota.
- * TODO: Research and possibly add (the undocumented) enum type PS_QUOTA_TYPE
- *       to replace UCHAR for 'PoolIndex'.
  * Notes: Conceptually translation unit local/private.
  */
 VOID
 NTAPI
 PspReturnProcessQuotaSpecifiedPool(IN PEPROCESS Process,
-                                   IN UCHAR     PoolIndex,
+                                   IN PS_QUOTA_TYPE QuotaType,
                                    IN SIZE_T    Amount)
 {
     ASSERT(Process);
     ASSERT(Process != PsInitialSystemProcess);
-    ASSERT(PoolIndex <= 2);
+    ASSERT(QuotaType < PsQuotaTypes);
     ASSERT(!(Amount & 0x80000000)); /* we need to be able to negate it */
-    if (Process->QuotaUsage[PoolIndex] < Amount)
+    if (Process->QuotaUsage[QuotaType] < Amount)
     {
         DPRINT1("WARNING: Process->QuotaUsage sanity check failed.\n");
     }
     else
     {
-        InterlockedExchangeAdd((LONG*)&Process->QuotaUsage[PoolIndex],
+        InterlockedExchangeAdd((LONG*)&Process->QuotaUsage[QuotaType],
                                -(LONG)Amount);
     }
 }
@@ -100,9 +96,9 @@ NTAPI
 PsInitializeQuotaSystem(VOID)
 {
     RtlZeroMemory(&PspDefaultQuotaBlock, sizeof(PspDefaultQuotaBlock));
-    PspDefaultQuotaBlock.QuotaEntry[PagedPool].Limit = (SIZE_T)-1;
-    PspDefaultQuotaBlock.QuotaEntry[NonPagedPool].Limit = (SIZE_T)-1;
-    PspDefaultQuotaBlock.QuotaEntry[2].Limit = (SIZE_T)-1; /* Page file */
+    PspDefaultQuotaBlock.QuotaEntry[PsNonPagedPool].Limit = (SIZE_T)-1;
+    PspDefaultQuotaBlock.QuotaEntry[PsPagedPool].Limit = (SIZE_T)-1;
+    PspDefaultQuotaBlock.QuotaEntry[PsPageFile].Limit = (SIZE_T)-1;
     PsGetCurrentProcess()->QuotaBlock = &PspDefaultQuotaBlock;
 }
 
@@ -164,7 +160,7 @@ PsChargeProcessPageFileQuota(IN PEPROCESS Process,
     /* Don't do anything for the system process */
     if (Process == PsInitialSystemProcess) return STATUS_SUCCESS;
 
-    return PspChargeProcessQuotaSpecifiedPool(Process, 2, Amount);
+    return PspChargeProcessQuotaSpecifiedPool(Process, PsPageFile, Amount);
 }
 
 /*
@@ -284,7 +280,7 @@ PsReturnProcessPageFileQuota(IN PEPROCESS Process,
     /* Don't do anything for the system process */
     if (Process == PsInitialSystemProcess) return STATUS_SUCCESS;
 
-    PspReturnProcessQuotaSpecifiedPool(Process, 2, Amount);
+    PspReturnProcessQuotaSpecifiedPool(Process, PsPageFile, Amount);
     return STATUS_SUCCESS;
 }
 
@@ -419,12 +415,12 @@ PspSetQuotaLimits(
         /* Initialize the quota block */
         QuotaBlock->ReferenceCount = 1;
         QuotaBlock->ProcessCount = 1;
-        QuotaBlock->QuotaEntry[0].Peak = Process->QuotaPeak[0];
-        QuotaBlock->QuotaEntry[1].Peak = Process->QuotaPeak[1];
-        QuotaBlock->QuotaEntry[2].Peak = Process->QuotaPeak[2];
-        QuotaBlock->QuotaEntry[0].Limit = PspDefaultQuotaBlock.QuotaEntry[0].Limit;
-        QuotaBlock->QuotaEntry[1].Limit = PspDefaultQuotaBlock.QuotaEntry[1].Limit;
-        QuotaBlock->QuotaEntry[2].Limit = PspDefaultQuotaBlock.QuotaEntry[2].Limit;
+        QuotaBlock->QuotaEntry[PsNonPagedPool].Peak =
Process->QuotaPeak[PsNonPagedPool];
+        QuotaBlock->QuotaEntry[PsPagedPool].Peak =
Process->QuotaPeak[PsPagedPool];
+        QuotaBlock->QuotaEntry[PsPageFile].Peak = Process->QuotaPeak[PsPageFile];
+        QuotaBlock->QuotaEntry[PsNonPagedPool].Limit =
PspDefaultQuotaBlock.QuotaEntry[PsNonPagedPool].Limit;
+        QuotaBlock->QuotaEntry[PsPagedPool].Limit =
PspDefaultQuotaBlock.QuotaEntry[PsPagedPool].Limit;
+        QuotaBlock->QuotaEntry[PsPageFile].Limit =
PspDefaultQuotaBlock.QuotaEntry[PsPageFile].Limit;
 
         /* Try to exchange the quota block, if that failed, just drop it */
         OldQuotaBlock =
InterlockedCompareExchangePointer((PVOID*)&Process->QuotaBlock,
diff --git a/sdk/include/ndk/pstypes.h b/sdk/include/ndk/pstypes.h
index d10199097a3..a1868e1d59a 100644
--- a/sdk/include/ndk/pstypes.h
+++ b/sdk/include/ndk/pstypes.h
@@ -972,6 +972,23 @@ typedef struct _JOB_SET_ARRAY
     ULONG Flags;
 } JOB_SET_ARRAY, *PJOB_SET_ARRAY;
 
+//
+// Process Quota Type
+//
+typedef enum _PS_QUOTA_TYPE
+{
+    PsNonPagedPool = 0,
+    PsPagedPool,
+    PsPageFile,
+#if (NTDDI_VERSION >= NTDDI_LONGHORN)
+    PsWorkingSet,
+#endif
+#if (NTDDI_VERSION == NTDDI_LONGHORN)
+    PsCpuRate,
+#endif
+    PsQuotaTypes
+} PS_QUOTA_TYPE;
+
 //
 // EPROCESS Quota Structures
 //
@@ -985,7 +1002,7 @@ typedef struct _EPROCESS_QUOTA_ENTRY
 
 typedef struct _EPROCESS_QUOTA_BLOCK
 {
-    EPROCESS_QUOTA_ENTRY QuotaEntry[3];
+    EPROCESS_QUOTA_ENTRY QuotaEntry[PsQuotaTypes];
     LIST_ENTRY QuotaList;
     ULONG ReferenceCount;
     ULONG ProcessCount;
@@ -1208,8 +1225,8 @@ typedef struct _EPROCESS
     EX_RUNDOWN_REF RundownProtect;
     HANDLE UniqueProcessId;
     LIST_ENTRY ActiveProcessLinks;
-    SIZE_T QuotaUsage[3]; /* 0=PagedPool, 1=NonPagedPool, 2=Pagefile */
-    SIZE_T QuotaPeak[3];  /* ditto */
+    SIZE_T QuotaUsage[PsQuotaTypes];
+    SIZE_T QuotaPeak[PsQuotaTypes];
     SIZE_T CommitCharge;
     SIZE_T PeakVirtualSize;
     SIZE_T VirtualSize;