[ros-diffs] [reactos] 02/11: [NTOS:EX] Manage quotas when allocating or freeing pool tables

[George Bișoc]

https://git.reactos.org/?p=reactos.git;a=commitdiff;h=71a4921f8ab706f81264e…

commit 71a4921f8ab706f81264e53df5dacd241c8c22f1
Author:     George Bișoc <george.bisoc(a)reactos.org>
AuthorDate: Thu Dec 30 21:05:27 2021 +0100
Commit:     George Bișoc <george.bisoc(a)reactos.org>
CommitDate: Tue Jan 11 10:11:08 2022 +0100

    [NTOS:EX] Manage quotas when allocating or freeing pool tables
    
    This fixes an assertion where QuotaUsage == 0 is actually not 0 when a process is
about to be destroyed.
---
 ntoskrnl/ex/handle.c | 40 +++++++++++++++++++++++++++++++++++-----
 1 file changed, 35 insertions(+), 5 deletions(-)

diff --git a/ntoskrnl/ex/handle.c b/ntoskrnl/ex/handle.c
index be4406bc56a..94b345caeb3 100644
--- a/ntoskrnl/ex/handle.c
+++ b/ntoskrnl/ex/handle.c
@@ -102,6 +102,7 @@ ExpAllocateTablePagedPool(IN PEPROCESS Process OPTIONAL,
                           IN SIZE_T Size)
 {
     PVOID Buffer;
+    NTSTATUS Status;
 
     /* Do the allocation */
     Buffer = ExAllocatePoolWithTag(PagedPool, Size, TAG_OBJECT_TABLE);
@@ -113,7 +114,13 @@ ExpAllocateTablePagedPool(IN PEPROCESS Process OPTIONAL,
         /* Check if we have a process to charge quota */
         if (Process)
         {
-            /* FIXME: Charge quota */
+            /* Charge quota */
+            Status = PsChargeProcessPagedPoolQuota(Process, Size);
+            if (!NT_SUCCESS(Status))
+            {
+                ExFreePoolWithTag(Buffer, TAG_OBJECT_TABLE);
+                return NULL;
+            }
         }
     }
 
@@ -127,6 +134,7 @@ ExpAllocateTablePagedPoolNoZero(IN PEPROCESS Process OPTIONAL,
                                 IN SIZE_T Size)
 {
     PVOID Buffer;
+    NTSTATUS Status;
 
     /* Do the allocation */
     Buffer = ExAllocatePoolWithTag(PagedPool, Size, TAG_OBJECT_TABLE);
@@ -135,7 +143,13 @@ ExpAllocateTablePagedPoolNoZero(IN PEPROCESS Process OPTIONAL,
         /* Check if we have a process to charge quota */
         if (Process)
         {
-            /* FIXME: Charge quota */
+            /* Charge quota */
+            Status = PsChargeProcessPagedPoolQuota(Process, Size);
+            if (!NT_SUCCESS(Status))
+            {
+                ExFreePoolWithTag(Buffer, TAG_OBJECT_TABLE);
+                return NULL;
+            }
         }
     }
 
@@ -153,7 +167,8 @@ ExpFreeTablePagedPool(IN PEPROCESS Process OPTIONAL,
     ExFreePoolWithTag(Buffer, TAG_OBJECT_TABLE);
     if (Process)
     {
-        /* FIXME: Release quota */
+        /* Release quota */
+        PsReturnProcessPagedPoolQuota(Process, Size);
     }
 }
 
@@ -248,7 +263,8 @@ ExpFreeHandleTable(IN PHANDLE_TABLE HandleTable)
     ExFreePoolWithTag(HandleTable, TAG_OBJECT_TABLE);
     if (Process)
     {
-        /* FIXME: TODO */
+        /* Release the quota it was taking up */
+        PsReturnProcessPagedPoolQuota(Process, sizeof(HANDLE_TABLE));
     }
 }
 
@@ -312,6 +328,7 @@ ExpAllocateHandleTable(IN PEPROCESS Process OPTIONAL,
     PHANDLE_TABLE HandleTable;
     PHANDLE_TABLE_ENTRY HandleTableTable, HandleEntry;
     ULONG i;
+    NTSTATUS Status;
     PAGED_CODE();
 
     /* Allocate the table */
@@ -323,7 +340,13 @@ ExpAllocateHandleTable(IN PEPROCESS Process OPTIONAL,
     /* Check if we have a process */
     if (Process)
     {
-        /* FIXME: Charge quota */
+        /* Charge quota */
+        Status = PsChargeProcessPagedPoolQuota(Process, sizeof(HANDLE_TABLE));
+        if (!NT_SUCCESS(Status))
+        {
+            ExFreePoolWithTag(HandleTable, TAG_OBJECT_TABLE);
+            return NULL;
+        }
     }
 
     /* Clear the table */
@@ -335,6 +358,13 @@ ExpAllocateHandleTable(IN PEPROCESS Process OPTIONAL,
     {
         /* Failed, free the table */
         ExFreePoolWithTag(HandleTable, TAG_OBJECT_TABLE);
+
+        /* Return the quota it was taking up */
+        if (Process)
+        {
+            PsReturnProcessPagedPoolQuota(Process, sizeof(HANDLE_TABLE));
+        }
+
         return NULL;
     }