From 36f1e2f5d1571d8d6eb7003382598be12c964bcb Mon Sep 17 00:00:00 2001
From: Alex Ionescu <aionescu@gmail.com>
Date: Tue, 10 Sep 2013 10:02:27 +0000
Subject: [PATCH] [NTOSKRNL]: Implement ExAllocatePoolWithQuotaTag. However
 PsChargeProcessPoolQuota is still unimplemented so nothing really gets
 charged (we do write the owner process, though).

svn path=/trunk/; revision=60002
---
 reactos/ntoskrnl/mm/ARM3/expool.c | 109 +++++++++++++++++++++++++++++-
 1 file changed, 106 insertions(+), 3 deletions(-)

diff --git a/reactos/ntoskrnl/mm/ARM3/expool.c b/reactos/ntoskrnl/mm/ARM3/expool.c
index 5a5f6a8236f..d645bef95a4 100644
--- a/reactos/ntoskrnl/mm/ARM3/expool.c
+++ b/reactos/ntoskrnl/mm/ARM3/expool.c
@@ -2511,11 +2511,114 @@ ExAllocatePoolWithQuotaTag(IN POOL_TYPE PoolType,
                            IN SIZE_T NumberOfBytes,
                            IN ULONG Tag)
 {
+    BOOLEAN Raise = TRUE;
+    PVOID Buffer;
+    PPOOL_HEADER Entry;
+    NTSTATUS Status;
+    PEPROCESS Process = PsGetCurrentProcess();
+
     //
-    // Allocate the pool
+    // Check if we should fail intead of raising an exception
     //
-    UNIMPLEMENTED;
-    return ExAllocatePoolWithTag(PoolType, NumberOfBytes, Tag);
+    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