[ros-diffs] [ion] 60002: [NTOSKRNL]: Implement ExAllocatePoolWithQuotaTag. However PsChargeProcessPoolQuota is still unimplemented so nothing really gets charged (we do write the owner process, though).

Author: ion
Date: Tue Sep 10 10:02:27 2013
New Revision: 60002

URL: http://svn.reactos.org/svn/reactos?rev=60002&view=rev
Log:
[NTOSKRNL]: Implement ExAllocatePoolWithQuotaTag. However PsChargeProcessPoolQuota is
still unimplemented so nothing really gets charged (we do write the owner process,
though).

Modified:
    trunk/reactos/ntoskrnl/mm/ARM3/expool.c

Modified: trunk/reactos/ntoskrnl/mm/ARM3/expool.c
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/ntoskrnl/mm/ARM3/expool.c?…
==============================================================================
--- trunk/reactos/ntoskrnl/mm/ARM3/expool.c	[iso-8859-1] (original)
+++ trunk/reactos/ntoskrnl/mm/ARM3/expool.c	[iso-8859-1] Tue Sep 10 10:02:27 2013
@@ -2511,11 +2511,114 @@
                            IN SIZE_T NumberOfBytes,
                            IN ULONG Tag)
 {
-    //
-    // Allocate the pool
-    //
-    UNIMPLEMENTED;
-    return ExAllocatePoolWithTag(PoolType, NumberOfBytes, Tag);
+    BOOLEAN Raise = TRUE;
+    PVOID Buffer;
+    PPOOL_HEADER Entry;
+    NTSTATUS Status;
+    PEPROCESS Process = PsGetCurrentProcess();
+
+    //
+    // Check if we should fail intead of raising an exception
+    //
+    if (PoolType & POOL_QUOTA_FAIL_INSTEAD_OF_RAISE)
+    {
+        Raise = FALSE;
+        PoolType &= ~POOL_QUOTA_FAIL_INSTEAD_OF_RAISE;
+    }
+
+    //
+    // Inject the pool quota mask
+    //
+    PoolType += QUOTA_POOL_MASK;
+
+    //
+    // Check if we have enough space to add the quota owner process, as long as
+    // this isn't the system process, which never gets charged quota
+    //
+    ASSERT(NumberOfBytes != 0);
+    if ((NumberOfBytes <= (PAGE_SIZE - sizeof(POOL_BLOCK_SIZE) - sizeof(PVOID)))
&&
+        (Process != PsInitialSystemProcess))
+    {
+        //
+        // Add space for our EPROCESS pointer
+        //
+        NumberOfBytes += sizeof(PEPROCESS);
+    }
+    else
+    {
+        //
+        // We won't be able to store the pointer, so don't use quota for this
+        //
+        PoolType -= QUOTA_POOL_MASK;
+    }
+
+    //
+    // Allocate the pool buffer now
+    //
+    Buffer = ExAllocatePoolWithTag(PoolType, NumberOfBytes, Tag);
+
+    //
+    // If the buffer is page-aligned, this is a large page allocation and we
+    // won't touch it
+    //
+    if (PAGE_ALIGN(Buffer) != Buffer)
+    {
+        //
+        // Also if special pool is enabled, and this was allocated from there,
+        // we won't touch it either
+        //
+        if ((ExpPoolFlags & POOL_FLAG_SPECIAL_POOL) &&
+            (MmIsSpecialPoolAddress(Buffer)))
+        {
+            return Buffer;
+        }
+
+        //
+        // If it wasn't actually allocated with quota charges, ignore it too
+        //
+        if (!(PoolType & QUOTA_POOL_MASK)) return Buffer;
+
+        //
+        // If this is the system process, we don't charge quota, so ignore
+        //
+        if (Process == PsInitialSystemProcess) return Buffer;
+
+        //
+        // Actually go and charge quota for the process now
+        //
+        Entry = POOL_ENTRY(Buffer);
+        Status = PsChargeProcessPoolQuota(Process,
+                                          PoolType & BASE_POOL_TYPE_MASK,
+                                          Entry->BlockSize *
sizeof(POOL_BLOCK_SIZE));
+        if (!NT_SUCCESS(Status))
+        {
+            //
+            // Quota failed, back out the allocation, clear the owner, and fail
+            //
+            *(PVOID*)((ULONG_PTR)POOL_NEXT_BLOCK(Entry) - sizeof(PVOID)) = NULL;
+            ExFreePoolWithTag(Buffer, Tag);
+            if (Raise) RtlRaiseStatus(Status);
+            return NULL;
+        }
+
+        //
+        // Quota worked, write the owner and then reference it before returning
+        //
+        *(PVOID*)((ULONG_PTR)POOL_NEXT_BLOCK(Entry) - sizeof(PVOID)) = Process;
+        ObReferenceObject(Process);
+    }
+    else if (!(Buffer) && (Raise))
+    {
+        //
+        // The allocation failed, raise an error if we are in raise mode
+        //
+        RtlRaiseStatus(STATUS_INSUFFICIENT_RESOURCES);
+    }
+
+    //
+    // Return the allocated buffer
+    //
+    return Buffer;
 }
 
 #if DBG && KDBG