[NTOS:PS] Add query support for QUOTA_LIMITS_EX

TAN-Gaming · Dec 10, 2024 · e58fb6f · e58fb6f
1 parent 16025b4
commit e58fb6f
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 31 deletions.
diff --git a/ntoskrnl/include/internal/ps_i.h b/ntoskrnl/include/internal/ps_i.h
@@ -29,7 +29,9 @@ static const INFORMATION_CLASS_INFO PsProcessInfoClass[] =
     (
         QUOTA_LIMITS,
         ULONG,
-        ICIF_QUERY | ICIF_SET | ICIF_SET_SIZE_VARIABLE
+
+        /* NOTE: ICIF_SIZE_VARIABLE is for QUOTA_LIMITS_EX support */
+        ICIF_QUERY | ICIF_SET | ICIF_SIZE_VARIABLE
     ),
 
     /* ProcessIoCounters */

diff --git a/ntoskrnl/ps/query.c b/ntoskrnl/ps/query.c
@@ -178,15 +178,17 @@ NtQueryInformationProcess(
         /* Process quota limits */
         case ProcessQuotaLimits:
         {
-            PQUOTA_LIMITS QuotaLimits = (PQUOTA_LIMITS)ProcessInformation;
+            QUOTA_LIMITS_EX QuotaLimits;
 
-            if (ProcessInformationLength != sizeof(QUOTA_LIMITS))
+            if (ProcessInformationLength != sizeof(QUOTA_LIMITS) &&
+                ProcessInformationLength != sizeof(QUOTA_LIMITS_EX))
             {
                 Status = STATUS_INFO_LENGTH_MISMATCH;
                 break;
             }
 
-            Length = sizeof(QUOTA_LIMITS);
+            /* Set return length */
+            Length = ProcessInformationLength;
 
             /* Reference the process */
             Status = ObReferenceObjectByHandle(ProcessHandle,
@@ -200,36 +202,43 @@ NtQueryInformationProcess(
             /* Indicate success */
             Status = STATUS_SUCCESS;
 
-            _SEH2_TRY
+            RtlZeroMemory(&QuotaLimits, sizeof(QuotaLimits));
+
+            /* Set max/min working set sizes */
+            QuotaLimits.MaximumWorkingSetSize =
+                Process->Vm.MaximumWorkingSetSize << PAGE_SHIFT;
+            QuotaLimits.MinimumWorkingSetSize =
+                Process->Vm.MinimumWorkingSetSize << PAGE_SHIFT;
+
+            /* Set default time limits */
+            QuotaLimits.TimeLimit.LowPart = MAXULONG;
+            QuotaLimits.TimeLimit.HighPart = MAXULONG;
+
+            /* Is quota block a default one? */
+            if (Process->QuotaBlock == &PspDefaultQuotaBlock)
+            {
+                /* Set default pools and pagefile limits */
+                QuotaLimits.PagedPoolLimit = (SIZE_T)-1;
+                QuotaLimits.NonPagedPoolLimit = (SIZE_T)-1;
+                QuotaLimits.PagefileLimit = (SIZE_T)-1;
+            }
+            else
             {
-                /* Set max/min working set sizes */
-                QuotaLimits->MaximumWorkingSetSize =
-                        Process->Vm.MaximumWorkingSetSize << PAGE_SHIFT;
-                QuotaLimits->MinimumWorkingSetSize =
-                        Process->Vm.MinimumWorkingSetSize << PAGE_SHIFT;
+                /* Get limits from non-default quota block */
+                QuotaLimits.PagedPoolLimit =
+                    Process->QuotaBlock->QuotaEntry[PsPagedPool].Limit;
+                QuotaLimits.NonPagedPoolLimit =
+                    Process->QuotaBlock->QuotaEntry[PsNonPagedPool].Limit;
+                QuotaLimits.PagefileLimit =
+                    Process->QuotaBlock->QuotaEntry[PsPageFile].Limit;
+            }
 
-                /* Set default time limits */
-                QuotaLimits->TimeLimit.LowPart = MAXULONG;
-                QuotaLimits->TimeLimit.HighPart = MAXULONG;
+            /* TODO: Set more fields of the extended version */
 
-                /* Is quota block a default one? */
-                if (Process->QuotaBlock == &PspDefaultQuotaBlock)
-                {
-                    /* Set default pools and pagefile limits */
-                    QuotaLimits->PagedPoolLimit = (SIZE_T)-1;
-                    QuotaLimits->NonPagedPoolLimit = (SIZE_T)-1;
-                    QuotaLimits->PagefileLimit = (SIZE_T)-1;
-                }
-                else
-                {
-                    /* Get limits from non-default quota block */
-                    QuotaLimits->PagedPoolLimit =
-                        Process->QuotaBlock->QuotaEntry[PsPagedPool].Limit;
-                    QuotaLimits->NonPagedPoolLimit =
-                        Process->QuotaBlock->QuotaEntry[PsNonPagedPool].Limit;
-                    QuotaLimits->PagefileLimit =
-                        Process->QuotaBlock->QuotaEntry[PsPageFile].Limit;
-                }
+            /* Protect writes with SEH */
+            _SEH2_TRY
+            {
+                RtlCopyMemory(ProcessInformation, &QuotaLimits, Length);
             }
             _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
             {