[KMTESTS:PS] Write some tests for process quota management

diff --git a/modules/rostests/kmtests/CMakeLists.txt b/modules/rostests/kmtests/CMakeLists.txt
index 4bfeb36..588dd3d 100644 (file)
--- a/modules/rostests/kmtests/CMakeLists.txt
+++ b/modules/rostests/kmtests/CMakeLists.txt
@@ -94,6 +94,7 @@ list(APPEND KMTEST_DRV_SOURCE
     ntos_ob/ObTypes.c
     ntos_ob/ObWait.c
     ntos_ps/PsNotify.c
+    ntos_ps/PsQuota.c
     ntos_se/SeHelpers.c
     ntos_se/SeInheritance.c
     ntos_se/SeLogonSession.c

diff --git a/modules/rostests/kmtests/kmtest_drv/testlist.c b/modules/rostests/kmtests/kmtest_drv/testlist.c
index 36bec00..c9c3781 100644 (file)
--- a/modules/rostests/kmtests/kmtest_drv/testlist.c
+++ b/modules/rostests/kmtests/kmtest_drv/testlist.c
@@ -63,6 +63,7 @@ KMT_TESTFUNC Test_ObTypeClean;
 KMT_TESTFUNC Test_ObTypeNoClean;
 KMT_TESTFUNC Test_ObTypes;
 KMT_TESTFUNC Test_PsNotify;
+KMT_TESTFUNC Test_PsQuota;
 KMT_TESTFUNC Test_SeInheritance;
 KMT_TESTFUNC Test_SeLogonSession;
 KMT_TESTFUNC Test_SeQueryInfoToken;
@@ -142,6 +143,7 @@ const KMT_TEST TestList[] =
     { "-ObTypeNoClean",                     Test_ObTypeNoClean },
     { "ObTypes",                            Test_ObTypes },
     { "PsNotify",                           Test_PsNotify },
+    { "PsQuota",                            Test_PsQuota },
     { "RtlAvlTreeKM",                       Test_RtlAvlTree },
     { "RtlExceptionKM",                     Test_RtlException },
     { "RtlIntSafeKM",                       Test_RtlIntSafe },

diff --git a/modules/rostests/kmtests/ntos_ps/PsQuota.c b/modules/rostests/kmtests/ntos_ps/PsQuota.c
new file mode 100644 (file)
index 0000000..8eb3fe2
--- /dev/null
+++ b/modules/rostests/kmtests/ntos_ps/PsQuota.c
@@ -0,0 +1,133 @@
+/*
+ * PROJECT:     ReactOS kernel-mode tests
+ * LICENSE:     GPL-2.0-or-later (https://spdx.org/licenses/GPL-2.0-or-later)
+ * PURPOSE:     Kernel mode tests for process quotas
+ * COPYRIGHT:   Copyright 2021 George Bișoc <george.bisoc@reactos.org>
+ */
+
+#include <kmt_test.h>
+
+START_TEST(PsQuota)
+{
+    NTSTATUS Status;
+    VM_COUNTERS VmCounters;
+    QUOTA_LIMITS QuotaLimits;
+    SIZE_T NonPagedUsage, PagedUsage;
+    PEPROCESS Process = PsGetCurrentProcess();
+
+    /* Guard the quota operations in a guarded region */
+    KeEnterGuardedRegion();
+
+    /* Report the current process' quota limits */
+    Status = ZwQueryInformationProcess(NtCurrentProcess(),
+                                       ProcessQuotaLimits,
+                                       &QuotaLimits,
+                                       sizeof(QuotaLimits),
+                                       NULL);
+    if (skip(NT_SUCCESS(Status), "Failed to query quota limits -- %lx\n", Status))
+    {
+        return;
+    }
+
+    trace("Process paged pool quota limit -- %lu\n", QuotaLimits.PagedPoolLimit);
+    trace("Process non paged pool quota limit -- %lu\n", QuotaLimits.NonPagedPoolLimit);
+    trace("Process page file quota limit -- %lu\n\n", QuotaLimits.PagefileLimit);
+
+    /* Query the quota usage */
+    Status = ZwQueryInformationProcess(NtCurrentProcess(),
+                                       ProcessVmCounters,
+                                       &VmCounters,
+                                       sizeof(VmCounters),
+                                       NULL);
+    if (skip(NT_SUCCESS(Status), "Failed to query quota usage -- %lx\n", Status))
+    {
+        return;
+    }
+
+    /* Test that quotas usage are within limits */
+    ok(VmCounters.QuotaNonPagedPoolUsage < QuotaLimits.NonPagedPoolLimit, "Non paged quota over limits (usage -> %lu || limit -> %lu)\n",
+       VmCounters.QuotaNonPagedPoolUsage, QuotaLimits.NonPagedPoolLimit);
+    ok(VmCounters.QuotaPagedPoolUsage < QuotaLimits.PagedPoolLimit, "Paged quota over limits (usage -> %lu || limit -> %lu)\n",
+       VmCounters.QuotaPagedPoolUsage, QuotaLimits.PagedPoolLimit);
+
+    /* Cache the quota usage pools for later checks  */
+    NonPagedUsage = VmCounters.QuotaNonPagedPoolUsage;
+    PagedUsage = VmCounters.QuotaPagedPoolUsage;
+
+    /* Charge some paged and non paged quotas */
+    Status = PsChargeProcessNonPagedPoolQuota(Process, 0x200);
+    ok_irql(PASSIVE_LEVEL);
+    ok_eq_hex(Status, STATUS_SUCCESS);
+
+    Status = PsChargeProcessPagedPoolQuota(Process, 0x500);
+    ok_irql(PASSIVE_LEVEL);
+    ok_eq_hex(Status, STATUS_SUCCESS);
+
+    /* Query the quota usage again */
+    Status = ZwQueryInformationProcess(NtCurrentProcess(),
+                                       ProcessVmCounters,
+                                       &VmCounters,
+                                       sizeof(VmCounters),
+                                       NULL);
+    if (skip(NT_SUCCESS(Status), "Failed to query quota usage -- %lx\n", Status))
+    {
+        return;
+    }
+
+
+    /* Test again the usage that's within limits */
+    ok(VmCounters.QuotaNonPagedPoolUsage < QuotaLimits.NonPagedPoolLimit, "Non paged quota over limits (usage -> %lu || limit -> %lu)\n",
+       VmCounters.QuotaNonPagedPoolUsage, QuotaLimits.NonPagedPoolLimit);
+    ok(VmCounters.QuotaPagedPoolUsage < QuotaLimits.PagedPoolLimit, "Paged quota over limits (usage -> %lu || limit -> %lu)\n",
+       VmCounters.QuotaPagedPoolUsage, QuotaLimits.PagedPoolLimit);
+
+    /*
+     * Make sure the results are consistent, that nobody else
+     * is charging quotas other than us.
+     */
+    ok_eq_size(VmCounters.QuotaNonPagedPoolUsage, NonPagedUsage + 0x200);
+    ok_eq_size(VmCounters.QuotaPagedPoolUsage, PagedUsage + 0x500);
+
+    /* Report the quota usage */
+    trace("=== QUOTA USAGE AFTER CHARGE ===\n\n");
+    trace("Process paged pool quota usage -- %lu\n", VmCounters.QuotaPagedPoolUsage);
+    trace("Process paged pool quota peak -- %lu\n", VmCounters.QuotaPeakPagedPoolUsage);
+    trace("Process non paged pool quota usage -- %lu\n", VmCounters.QuotaNonPagedPoolUsage);
+    trace("Process non paged pool quota peak -- %lu\n", VmCounters.QuotaPeakNonPagedPoolUsage);
+    trace("Process page file quota usage -- %lu\n", VmCounters.PagefileUsage);
+    trace("Process page file quota peak -- %lu\n\n", VmCounters.PeakPagefileUsage);
+
+    /* Return the quotas we've charged up */
+    PsReturnProcessNonPagedPoolQuota(Process, 0x200);
+    PsReturnProcessPagedPoolQuota(Process, 0x500);
+
+    /* Query the quota usage again */
+    Status = ZwQueryInformationProcess(NtCurrentProcess(),
+                                       ProcessVmCounters,
+                                       &VmCounters,
+                                       sizeof(VmCounters),
+                                       NULL);
+    if (skip(NT_SUCCESS(Status), "Failed to query quota usage -- %lx\n", Status))
+    {
+        return;
+    }
+
+    /*
+     * Check that nobody else has returned quotas
+     * but only us.
+     */
+    ok_eq_size(VmCounters.QuotaNonPagedPoolUsage, NonPagedUsage);
+    ok_eq_size(VmCounters.QuotaPagedPoolUsage, PagedUsage);
+
+    /* Report the usage again */
+    trace("=== QUOTA USAGE AFTER RETURN ===\n\n");
+    trace("Process paged pool quota usage -- %lu\n", VmCounters.QuotaPagedPoolUsage);
+    trace("Process paged pool quota peak -- %lu\n", VmCounters.QuotaPeakPagedPoolUsage);
+    trace("Process non paged pool quota usage -- %lu\n", VmCounters.QuotaNonPagedPoolUsage);
+    trace("Process non paged pool quota peak -- %lu\n", VmCounters.QuotaPeakNonPagedPoolUsage);
+    trace("Process page file quota usage -- %lu\n", VmCounters.PagefileUsage);
+    trace("Process page file quota peak -- %lu\n\n", VmCounters.PeakPagefileUsage);
+
+    /* We're done, leave the region */
+    KeLeaveGuardedRegion();
+}