reactos/ntoskrnl/ps/quota.c

   1 /*
   2  * COPYRIGHT:       See COPYING in the top level directory
   3  * PROJECT:         ReactOS kernel
   4  * FILE:            ntoskrnl/ps/quota.c
   5  * PURPOSE:         Process Pool Quotas
   6  *
   7  * PROGRAMMERS:     Alex Ionescu (alex@relsoft.net)
   8  *                  Mike Nordell
   9  */
  10
  11 /* INCLUDES **************************************************************/
  12
  13 #include <ntoskrnl.h>
  14 #include <ntintsafe.h>
  15 #define NDEBUG
  16 #include <debug.h>
  17
  18 EPROCESS_QUOTA_BLOCK PspDefaultQuotaBlock;
  19 static LIST_ENTRY PspQuotaBlockList = {&PspQuotaBlockList, &PspQuotaBlockList};
  20 static KSPIN_LOCK PspQuotaLock;
  21
  22 #define TAG_QUOTA_BLOCK 'bQsP'
  23 #define VALID_QUOTA_FLAGS (QUOTA_LIMITS_HARDWS_MIN_ENABLE | \
  24                            QUOTA_LIMITS_HARDWS_MIN_DISABLE | \
  25                            QUOTA_LIMITS_HARDWS_MAX_ENABLE | \
  26                            QUOTA_LIMITS_HARDWS_MAX_DISABLE)
  27
  28 /* PRIVATE FUNCTIONS *******************************************************/
  29
  30 /*
  31  * Private helper to charge the specified process quota.
  32  * ReturnsSTATUS_QUOTA_EXCEEDED on quota limit check failure.
  33  * Updates QuotaPeak as needed for specified PoolIndex.
  34  * TODO: Research and possibly add (the undocumented) enum type PS_QUOTA_TYPE
  35  *       to replace UCHAR for 'PoolIndex'.
  36  * Notes: Conceptually translation unit local/private.
  37  */
  38 NTSTATUS
  39 NTAPI
  40 PspChargeProcessQuotaSpecifiedPool(IN PEPROCESS Process,
  41                                    IN UCHAR     PoolIndex,
  42                                    IN SIZE_T    Amount)
  43 {
  44     ASSERT(Process);
  45     ASSERT(Process != PsInitialSystemProcess);
  46     ASSERT(PoolIndex <= 2);
  47     ASSERT(Process->QuotaBlock);
  48
  49     /* Note: Race warning. TODO: Needs to add/use lock for this */
  50     if (Process->QuotaUsage[PoolIndex] + Amount >
  51         Process->QuotaBlock->QuotaEntry[PoolIndex].Limit)
  52     {
  53         DPRINT1("Quota exceeded, but ROS will let it slide...\n");
  54         return STATUS_SUCCESS;
  55         //return STATUS_QUOTA_EXCEEDED; /* caller raises the exception */
  56     }
  57
  58     InterlockedExchangeAdd((LONG*)&Process->QuotaUsage[PoolIndex], Amount);
  59
  60     /* Note: Race warning. TODO: Needs to add/use lock for this */
  61     if (Process->QuotaPeak[PoolIndex] < Process->QuotaUsage[PoolIndex])
  62     {
  63         Process->QuotaPeak[PoolIndex] = Process->QuotaUsage[PoolIndex];
  64     }
  65
  66     return STATUS_SUCCESS;
  67 }
  68
  69 /*
  70  * Private helper to remove quota charge from the specified process quota.
  71  * TODO: Research and possibly add (the undocumented) enum type PS_QUOTA_TYPE
  72  *       to replace UCHAR for 'PoolIndex'.
  73  * Notes: Conceptually translation unit local/private.
  74  */
  75 VOID
  76 NTAPI
  77 PspReturnProcessQuotaSpecifiedPool(IN PEPROCESS Process,
  78                                    IN UCHAR     PoolIndex,
  79                                    IN SIZE_T    Amount)
  80 {
  81     ASSERT(Process);
  82     ASSERT(Process != PsInitialSystemProcess);
  83     ASSERT(PoolIndex <= 2);
  84     ASSERT(!(Amount & 0x80000000)); /* we need to be able to negate it */
  85     if (Process->QuotaUsage[PoolIndex] < Amount)
  86     {
  87         DPRINT1("WARNING: Process->QuotaUsage sanity check failed.\n");
  88     }
  89     else
  90     {
  91         InterlockedExchangeAdd((LONG*)&Process->QuotaUsage[PoolIndex],
  92                                -(LONG)Amount);
  93     }
  94 }
  95
  96 /* FUNCTIONS ***************************************************************/
  97
  98 VOID
  99 NTAPI
 100 INIT_FUNCTION
 101 PsInitializeQuotaSystem(VOID)
 102 {
 103     RtlZeroMemory(&PspDefaultQuotaBlock, sizeof(PspDefaultQuotaBlock));
 104     PspDefaultQuotaBlock.QuotaEntry[PagedPool].Limit = (SIZE_T)-1;
 105     PspDefaultQuotaBlock.QuotaEntry[NonPagedPool].Limit = (SIZE_T)-1;
 106     PspDefaultQuotaBlock.QuotaEntry[2].Limit = (SIZE_T)-1; /* Page file */
 107     PsGetCurrentProcess()->QuotaBlock = &PspDefaultQuotaBlock;
 108 }
 109
 110 VOID
 111 NTAPI
 112 PspInheritQuota(PEPROCESS Process, PEPROCESS ParentProcess)
 113 {
 114     if (ParentProcess != NULL)
 115     {
 116         PEPROCESS_QUOTA_BLOCK QuotaBlock = ParentProcess->QuotaBlock;
 117
 118         ASSERT(QuotaBlock != NULL);
 119
 120         (void)InterlockedIncrementUL(&QuotaBlock->ReferenceCount);
 121
 122         Process->QuotaBlock = QuotaBlock;
 123     }
 124     else
 125         Process->QuotaBlock = &PspDefaultQuotaBlock;
 126 }
 127
 128 VOID
 129 NTAPI
 130 PspInsertQuotaBlock(
 131     PEPROCESS_QUOTA_BLOCK QuotaBlock)
 132 {
 133     KIRQL OldIrql;
 134
 135     KeAcquireSpinLock(&PspQuotaLock, &OldIrql);
 136     InsertTailList(&PspQuotaBlockList, &QuotaBlock->QuotaList);
 137     KeReleaseSpinLock(&PspQuotaLock, OldIrql);
 138 }
 139
 140 VOID
 141 NTAPI
 142 PspDestroyQuotaBlock(PEPROCESS Process)
 143 {
 144     PEPROCESS_QUOTA_BLOCK QuotaBlock = Process->QuotaBlock;
 145     KIRQL OldIrql;
 146
 147     if (QuotaBlock != &PspDefaultQuotaBlock &&
 148         InterlockedDecrementUL(&QuotaBlock->ReferenceCount) == 0)
 149     {
 150         KeAcquireSpinLock(&PspQuotaLock, &OldIrql);
 151         RemoveEntryList(&QuotaBlock->QuotaList);
 152         KeReleaseSpinLock(&PspQuotaLock, OldIrql);
 153         ExFreePool(QuotaBlock);
 154     }
 155 }
 156
 157 /*
 158  * @implemented
 159  */
 160 NTSTATUS
 161 NTAPI
 162 PsChargeProcessPageFileQuota(IN PEPROCESS Process,
 163                              IN SIZE_T Amount)
 164 {
 165     /* Don't do anything for the system process */
 166     if (Process == PsInitialSystemProcess) return STATUS_SUCCESS;
 167
 168     return PspChargeProcessQuotaSpecifiedPool(Process, 2, Amount);
 169 }
 170
 171 /*
 172  * @implemented
 173  */
 174 VOID
 175 NTAPI
 176 PsChargePoolQuota(IN PEPROCESS Process,
 177                   IN POOL_TYPE PoolType,
 178                   IN SIZE_T Amount)
 179 {
 180     NTSTATUS Status;
 181     ASSERT(KeGetCurrentIrql() < DISPATCH_LEVEL);
 182
 183     /* Don't do anything for the system process */
 184     if (Process == PsInitialSystemProcess) return;
 185
 186     /* Charge the usage */
 187     Status = PsChargeProcessPoolQuota(Process, PoolType, Amount);
 188     if (!NT_SUCCESS(Status)) ExRaiseStatus(Status);
 189 }
 190
 191 /*
 192  * @implemented
 193  */
 194 NTSTATUS
 195 NTAPI
 196 PsChargeProcessNonPagedPoolQuota(IN PEPROCESS Process,
 197                                  IN SIZE_T Amount)
 198 {
 199     /* Call the general function */
 200     return PsChargeProcessPoolQuota(Process, NonPagedPool, Amount);
 201 }
 202
 203 /*
 204  * @implemented
 205  */
 206 NTSTATUS
 207 NTAPI
 208 PsChargeProcessPagedPoolQuota(IN PEPROCESS Process,
 209                               IN SIZE_T Amount)
 210 {
 211     /* Call the general function */
 212     return PsChargeProcessPoolQuota(Process, PagedPool, Amount);
 213 }
 214
 215 /*
 216  * @implemented
 217  */
 218 NTSTATUS
 219 NTAPI
 220 PsChargeProcessPoolQuota(IN PEPROCESS Process,
 221                          IN POOL_TYPE PoolType,
 222                          IN SIZE_T Amount)
 223 {
 224     /* Don't do anything for the system process */
 225     if (Process == PsInitialSystemProcess) return STATUS_SUCCESS;
 226
 227     return PspChargeProcessQuotaSpecifiedPool(Process,
 228                                               (PoolType & PAGED_POOL_MASK),
 229                                               Amount);
 230 }
 231
 232 /*
 233  * @implemented
 234  */
 235 VOID
 236 NTAPI
 237 PsReturnPoolQuota(IN PEPROCESS Process,
 238                   IN POOL_TYPE PoolType,
 239                   IN SIZE_T Amount)
 240 {
 241     /* Don't do anything for the system process */
 242     if (Process == PsInitialSystemProcess) return;
 243
 244     PspReturnProcessQuotaSpecifiedPool(Process,
 245                                        (PoolType & PAGED_POOL_MASK),
 246                                        Amount);
 247 }
 248
 249 /*
 250  * @implemented
 251  */
 252 VOID
 253 NTAPI
 254 PsReturnProcessNonPagedPoolQuota(IN PEPROCESS Process,
 255                                  IN SIZE_T Amount)
 256 {
 257     /* Don't do anything for the system process */
 258     if (Process == PsInitialSystemProcess) return;
 259
 260     PsReturnPoolQuota(Process, NonPagedPool, Amount);
 261 }
 262
 263 /*
 264  * @implemented
 265  */
 266 VOID
 267 NTAPI
 268 PsReturnProcessPagedPoolQuota(IN PEPROCESS Process,
 269                               IN SIZE_T Amount)
 270 {
 271     /* Don't do anything for the system process */
 272     if (Process == PsInitialSystemProcess) return;
 273
 274     PsReturnPoolQuota(Process, PagedPool, Amount);
 275 }
 276
 277 /*
 278  * @implemented
 279  */
 280 NTSTATUS
 281 NTAPI
 282 PsReturnProcessPageFileQuota(IN PEPROCESS Process,
 283                              IN SIZE_T Amount)
 284 {
 285     /* Don't do anything for the system process */
 286     if (Process == PsInitialSystemProcess) return STATUS_SUCCESS;
 287
 288     PspReturnProcessQuotaSpecifiedPool(Process, 2, Amount);
 289     return STATUS_SUCCESS;
 290 }
 291
 292 NTSTATUS
 293 NTAPI
 294 PspSetQuotaLimits(
 295     _In_ PEPROCESS Process,
 296     _In_ ULONG Unused,
 297     _In_ PVOID QuotaLimits,
 298     _In_ ULONG QuotaLimitsLength,
 299     _In_ KPROCESSOR_MODE PreviousMode)
 300 {
 301     QUOTA_LIMITS_EX CapturedQuotaLimits;
 302     PEPROCESS_QUOTA_BLOCK QuotaBlock, OldQuotaBlock;
 303     BOOLEAN IncreaseOkay;
 304     KAPC_STATE SavedApcState;
 305     NTSTATUS Status;
 306
 307     UNREFERENCED_PARAMETER(Unused);
 308
 309     _SEH2_TRY
 310     {
 311         ProbeForRead(QuotaLimits, QuotaLimitsLength, sizeof(ULONG));
 312
 313         /* Check if we have the basic or extended structure */
 314         if (QuotaLimitsLength == sizeof(QUOTA_LIMITS))
 315         {
 316             /* Copy the basic structure, zero init the remaining fields */
 317             RtlCopyMemory(&CapturedQuotaLimits, QuotaLimits, sizeof(QUOTA_LIMITS));
 318             CapturedQuotaLimits.WorkingSetLimit = 0;
 319             CapturedQuotaLimits.Reserved2 = 0;
 320             CapturedQuotaLimits.Reserved3 = 0;
 321             CapturedQuotaLimits.Reserved4 = 0;
 322             CapturedQuotaLimits.CpuRateLimit.RateData = 0;
 323             CapturedQuotaLimits.Flags = 0;
 324         }
 325         else if (QuotaLimitsLength == sizeof(QUOTA_LIMITS_EX))
 326         {
 327             /* Copy the full structure */
 328             RtlCopyMemory(&CapturedQuotaLimits, QuotaLimits, sizeof(QUOTA_LIMITS_EX));
 329
 330             /* Verify that the caller passed valid flags */
 331             if ((CapturedQuotaLimits.Flags & ~VALID_QUOTA_FLAGS) ||
 332                 ((CapturedQuotaLimits.Flags & QUOTA_LIMITS_HARDWS_MIN_ENABLE) &&
 333                  (CapturedQuotaLimits.Flags & QUOTA_LIMITS_HARDWS_MIN_DISABLE)) ||
 334                 ((CapturedQuotaLimits.Flags & QUOTA_LIMITS_HARDWS_MAX_ENABLE) &&
 335                  (CapturedQuotaLimits.Flags & QUOTA_LIMITS_HARDWS_MAX_DISABLE)))
 336             {
 337                 DPRINT1("Invalid quota flags: 0x%lx\n", CapturedQuotaLimits.Flags);
 338                 return STATUS_INVALID_PARAMETER;
 339             }
 340
 341             /* Verify that the caller didn't pass reserved values */
 342             if ((CapturedQuotaLimits.WorkingSetLimit != 0) ||
 343                 (CapturedQuotaLimits.Reserved2 != 0) ||
 344                 (CapturedQuotaLimits.Reserved3 != 0) ||
 345                 (CapturedQuotaLimits.Reserved4 != 0) ||
 346                 (CapturedQuotaLimits.CpuRateLimit.RateData != 0))
 347             {
 348                 DPRINT1("Invalid value: (%lx,%lx,%lx,%lx,%lx)\n",
 349                         CapturedQuotaLimits.WorkingSetLimit,
 350                         CapturedQuotaLimits.Reserved2,
 351                         CapturedQuotaLimits.Reserved3,
 352                         CapturedQuotaLimits.Reserved4,
 353                         CapturedQuotaLimits.CpuRateLimit.RateData);
 354                 return STATUS_INVALID_PARAMETER;
 355             }
 356         }
 357         else
 358         {
 359             DPRINT1("Invalid quota size: 0x%lx\n", QuotaLimitsLength);
 360             return STATUS_INFO_LENGTH_MISMATCH;
 361         }
 362     }
 363     _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
 364     {
 365         DPRINT1("Exception while copying data\n");
 366         return _SEH2_GetExceptionCode();
 367     }
 368     _SEH2_END;
 369
 370     /* Check the caller changes the working set size limits */
 371     if ((CapturedQuotaLimits.MinimumWorkingSetSize != 0) &&
 372         (CapturedQuotaLimits.MaximumWorkingSetSize != 0))
 373     {
 374         /* Check for special case: trimming the WS */
 375         if ((CapturedQuotaLimits.MinimumWorkingSetSize == SIZE_T_MAX) &&
 376             (CapturedQuotaLimits.MaximumWorkingSetSize == SIZE_T_MAX))
 377         {
 378             /* No increase allowed */
 379             IncreaseOkay = FALSE;
 380         }
 381         else
 382         {
 383             /* Check if the caller has the required privilege */
 384             IncreaseOkay = SeSinglePrivilegeCheck(SeIncreaseQuotaPrivilege,
 385                                                   PreviousMode);
 386         }
 387
 388         /* Attach to the target process and disable APCs */
 389         KeStackAttachProcess(&Process->Pcb, &SavedApcState);
 390         KeEnterGuardedRegion();
 391
 392         /* Call Mm to adjust the process' working set size */
 393         Status = MmAdjustWorkingSetSize(CapturedQuotaLimits.MinimumWorkingSetSize,
 394                                         CapturedQuotaLimits.MaximumWorkingSetSize,
 395                                         0,
 396                                         IncreaseOkay);
 397
 398         /* Bring back APCs and detach from the process */
 399         KeLeaveGuardedRegion();
 400         KeUnstackDetachProcess(&SavedApcState);
 401     }
 402     else if (Process->QuotaBlock == &PspDefaultQuotaBlock)
 403     {
 404         /* Check if the caller has the required privilege */
 405         if (!SeSinglePrivilegeCheck(SeIncreaseQuotaPrivilege, PreviousMode))
 406         {
 407             return STATUS_PRIVILEGE_NOT_HELD;
 408         }
 409
 410         /* Allocate a new quota block */
 411         QuotaBlock = ExAllocatePoolWithTag(NonPagedPool,
 412                                            sizeof(EPROCESS_QUOTA_BLOCK),
 413                                            TAG_QUOTA_BLOCK);
 414         if (QuotaBlock == NULL)
 415         {
 416             ObDereferenceObject(Process);
 417             return STATUS_NO_MEMORY;
 418         }
 419
 420         /* Initialize the quota block */
 421         QuotaBlock->ReferenceCount = 1;
 422         QuotaBlock->ProcessCount = 1;
 423         QuotaBlock->QuotaEntry[0].Peak = Process->QuotaPeak[0];
 424         QuotaBlock->QuotaEntry[1].Peak = Process->QuotaPeak[1];
 425         QuotaBlock->QuotaEntry[2].Peak = Process->QuotaPeak[2];
 426         QuotaBlock->QuotaEntry[0].Limit = PspDefaultQuotaBlock.QuotaEntry[0].Limit;
 427         QuotaBlock->QuotaEntry[1].Limit = PspDefaultQuotaBlock.QuotaEntry[1].Limit;
 428         QuotaBlock->QuotaEntry[2].Limit = PspDefaultQuotaBlock.QuotaEntry[2].Limit;
 429
 430         /* Try to exchange the quota block, if that failed, just drop it */
 431         OldQuotaBlock = InterlockedCompareExchangePointer(&Process->QuotaBlock,
 432                                                           QuotaBlock,
 433                                                           &PspDefaultQuotaBlock);
 434         if (OldQuotaBlock == &PspDefaultQuotaBlock)
 435         {
 436             /* Success, insert the new quota block */
 437             PspInsertQuotaBlock(QuotaBlock);
 438         }
 439         else
 440         {
 441             /* Failed, free the quota block and ignore it */
 442             ExFreePoolWithTag(QuotaBlock, TAG_QUOTA_BLOCK);
 443         }
 444
 445         Status = STATUS_SUCCESS;
 446     }
 447
 448     return Status;
 449 }
 450
 451
 452 /* EOF */