ref.c

Go to the documentation of this file.

/*
 * Process Hacker -
 *   object manager
 *
 * Copyright (C) 2009-2015 wj32
 *
 * This file is part of Process Hacker.
 *
 * Process Hacker is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Process Hacker is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Process Hacker.  If not, see <http://www.gnu.org/licenses/>.
 */

#define _PH_REF_PRIVATE
#include <phbase.h>
#include <phintrnl.h>
#include <refp.h>

PPH_OBJECT_TYPE PhObjectTypeObject = NULL;
SLIST_HEADER PhObjectDeferDeleteListHead;
PH_FREE_LIST PhObjectSmallFreeList;
PPH_OBJECT_TYPE PhAllocType = NULL;

ULONG PhObjectTypeCount = 0;
PPH_OBJECT_TYPE PhObjectTypeTable[PH_OBJECT_TYPE_TABLE_SIZE];

static ULONG PhpAutoPoolTlsIndex;

#ifdef DEBUG
LIST_ENTRY PhDbgObjectListHead;
PH_QUEUED_LOCK PhDbgObjectListLock = PH_QUEUED_LOCK_INIT;
PPH_CREATE_OBJECT_HOOK PhDbgCreateObjectHook = NULL;
#endif

#define REF_STAT_UP(Name) PHLIB_INC_STATISTIC(Name)

NTSTATUS PhInitializeRef(
    VOID
    )
{
    PH_OBJECT_TYPE dummyObjectType;

#ifdef DEBUG
    InitializeListHead(&PhDbgObjectListHead);
#endif

    RtlInitializeSListHead(&PhObjectDeferDeleteListHead);
    PhInitializeFreeList(
        &PhObjectSmallFreeList,
        PhAddObjectHeaderSize(PH_OBJECT_SMALL_OBJECT_SIZE),
        PH_OBJECT_SMALL_OBJECT_COUNT
        );

    // Create the fundamental object type.

    memset(&dummyObjectType, 0, sizeof(PH_OBJECT_TYPE));
    PhObjectTypeObject = &dummyObjectType; // PhCreateObject expects an object type.
    PhObjectTypeTable[0] = &dummyObjectType; // PhCreateObject also expects PhObjectTypeTable[0] to be filled in.
    PhObjectTypeObject = PhCreateObjectType(L"Type", 0, NULL);

    // Now that the fundamental object type exists, fix it up.
    PhObjectToObjectHeader(PhObjectTypeObject)->TypeIndex = PhObjectTypeObject->TypeIndex;
    PhObjectTypeObject->NumberOfObjects = 1;

    // Create the allocated memory object type.
    PhAllocType = PhCreateObjectType(L"Alloc", 0, NULL);

    // Reserve a slot for the auto pool.
    PhpAutoPoolTlsIndex = TlsAlloc();

    if (PhpAutoPoolTlsIndex == TLS_OUT_OF_INDEXES)
        return STATUS_INSUFFICIENT_RESOURCES;

    return STATUS_SUCCESS;
}

_May_raise_ PVOID PhCreateObject(
    _In_ SIZE_T ObjectSize,
    _In_ PPH_OBJECT_TYPE ObjectType
    )
{
    NTSTATUS status = STATUS_SUCCESS;
    PPH_OBJECT_HEADER objectHeader;

    // Allocate storage for the object. Note that this includes the object header followed by the object body.
    objectHeader = PhpAllocateObject(ObjectType, ObjectSize);

    // Object type statistics.
    _InterlockedIncrement((PLONG)&ObjectType->NumberOfObjects);

    // Initialize the object header.
    objectHeader->RefCount = 1;
    objectHeader->TypeIndex = ObjectType->TypeIndex;
    // objectHeader->Flags is set by PhpAllocateObject.

    REF_STAT_UP(RefObjectsCreated);

#ifdef DEBUG
    {
        USHORT capturedFrames;

        capturedFrames = RtlCaptureStackBackTrace(1, 16, objectHeader->StackBackTrace, NULL);
        memset(
            &objectHeader->StackBackTrace[capturedFrames],
            0,
            sizeof(objectHeader->StackBackTrace) - capturedFrames * sizeof(PVOID)
            );
    }

    PhAcquireQueuedLockExclusive(&PhDbgObjectListLock);
    InsertTailList(&PhDbgObjectListHead, &objectHeader->ObjectListEntry);
    PhReleaseQueuedLockExclusive(&PhDbgObjectListLock);

    {
        PPH_CREATE_OBJECT_HOOK dbgCreateObjectHook;

        dbgCreateObjectHook = PhDbgCreateObjectHook;

        if (dbgCreateObjectHook)
        {
            dbgCreateObjectHook(
                PhObjectHeaderToObject(objectHeader),
                ObjectSize,
                0,
                ObjectType
                );
        }
    }
#endif

    return PhObjectHeaderToObject(objectHeader);
}

PVOID PhReferenceObject(
    _In_ PVOID Object
    )
{
    PPH_OBJECT_HEADER objectHeader;

    objectHeader = PhObjectToObjectHeader(Object);
    // Increment the reference count.
    _InterlockedIncrement(&objectHeader->RefCount);

    return Object;
}

_May_raise_ LONG PhReferenceObjectEx(
    _In_ PVOID Object,
    _In_ LONG RefCount
    )
{
    PPH_OBJECT_HEADER objectHeader;
    LONG oldRefCount;

    assert(!(RefCount < 0));

    objectHeader = PhObjectToObjectHeader(Object);
    // Increase the reference count.
    oldRefCount = _InterlockedExchangeAdd(&objectHeader->RefCount, RefCount);

    return oldRefCount + RefCount;
}

BOOLEAN PhReferenceObjectSafe(
    _In_ PVOID Object
    )
{
    PPH_OBJECT_HEADER objectHeader;
    BOOLEAN result;

    objectHeader = PhObjectToObjectHeader(Object);
    // Increase the reference count only if it isn't 0 (atomically).
    result = PhpInterlockedIncrementSafe(&objectHeader->RefCount);

    return result;
}

VOID PhDereferenceObject(
    _In_ PVOID Object
    )
{
    PPH_OBJECT_HEADER objectHeader;
    LONG newRefCount;

    objectHeader = PhObjectToObjectHeader(Object);
    // Decrement the reference count.
    newRefCount = _InterlockedDecrement(&objectHeader->RefCount);
    ASSUME_ASSERT(newRefCount >= 0);

    // Free the object if it has 0 references.
    if (newRefCount == 0)
    {
        PhpFreeObject(objectHeader);
    }
}

BOOLEAN PhDereferenceObjectDeferDelete(
    _In_ PVOID Object
    )
{
    return PhDereferenceObjectEx(Object, 1, TRUE) == 0;
}

_May_raise_ LONG PhDereferenceObjectEx(
    _In_ PVOID Object,
    _In_ LONG RefCount,
    _In_ BOOLEAN DeferDelete
    )
{
    PPH_OBJECT_HEADER objectHeader;
    LONG oldRefCount;
    LONG newRefCount;

    assert(!(RefCount < 0));

    objectHeader = PhObjectToObjectHeader(Object);

    // Decrease the reference count.
    oldRefCount = _InterlockedExchangeAdd(&objectHeader->RefCount, -RefCount);
    newRefCount = oldRefCount - RefCount;

    // Free the object if it has 0 references.
    if (newRefCount == 0)
    {
        if (DeferDelete)
        {
            PhpDeferDeleteObject(objectHeader);
        }
        else
        {
            // Free the object.
            PhpFreeObject(objectHeader);
        }
    }
    else if (newRefCount < 0)
    {
        PhRaiseStatus(STATUS_INVALID_PARAMETER);
    }

    return newRefCount;
}

PPH_OBJECT_TYPE PhGetObjectType(
    _In_ PVOID Object
    )
{
    return PhObjectTypeTable[PhObjectToObjectHeader(Object)->TypeIndex];
}

PPH_OBJECT_TYPE PhCreateObjectType(
    _In_ PWSTR Name,
    _In_ ULONG Flags,
    _In_opt_ PPH_TYPE_DELETE_PROCEDURE DeleteProcedure
    )
{
    return PhCreateObjectTypeEx(
        Name,
        Flags,
        DeleteProcedure,
        NULL
        );
}

PPH_OBJECT_TYPE PhCreateObjectTypeEx(
    _In_ PWSTR Name,
    _In_ ULONG Flags,
    _In_opt_ PPH_TYPE_DELETE_PROCEDURE DeleteProcedure,
    _In_opt_ PPH_OBJECT_TYPE_PARAMETERS Parameters
    )
{
    NTSTATUS status = STATUS_SUCCESS;
    PPH_OBJECT_TYPE objectType;

    // Check the flags.
    if ((Flags & PH_OBJECT_TYPE_VALID_FLAGS) != Flags) /* Valid flag mask */
        PhRaiseStatus(STATUS_INVALID_PARAMETER_3);
    if ((Flags & PH_OBJECT_TYPE_USE_FREE_LIST) && !Parameters)
        PhRaiseStatus(STATUS_INVALID_PARAMETER_MIX);

    // Create the type object.
    objectType = PhCreateObject(sizeof(PH_OBJECT_TYPE), PhObjectTypeObject);

    // Initialize the type object.
    objectType->Flags = (USHORT)Flags;
    objectType->TypeIndex = (USHORT)_InterlockedIncrement(&PhObjectTypeCount) - 1;
    objectType->NumberOfObjects = 0;
    objectType->DeleteProcedure = DeleteProcedure;
    objectType->Name = Name;

    if (objectType->TypeIndex < PH_OBJECT_TYPE_TABLE_SIZE)
        PhObjectTypeTable[objectType->TypeIndex] = objectType;
    else
        PhRaiseStatus(STATUS_UNSUCCESSFUL);

    if (Parameters)
    {
        if (Flags & PH_OBJECT_TYPE_USE_FREE_LIST)
        {
            PhInitializeFreeList(
                &objectType->FreeList,
                PhAddObjectHeaderSize(Parameters->FreeListSize),
                Parameters->FreeListCount
                );
        }
    }

    return objectType;
}

VOID PhGetObjectTypeInformation(
    _In_ PPH_OBJECT_TYPE ObjectType,
    _Out_ PPH_OBJECT_TYPE_INFORMATION Information
    )
{
    Information->Name = ObjectType->Name;
    Information->NumberOfObjects = ObjectType->NumberOfObjects;
    Information->Flags = ObjectType->Flags;
    Information->TypeIndex = ObjectType->TypeIndex;
}

PPH_OBJECT_HEADER PhpAllocateObject(
    _In_ PPH_OBJECT_TYPE ObjectType,
    _In_ SIZE_T ObjectSize
    )
{
    PPH_OBJECT_HEADER objectHeader;

    if (ObjectType->Flags & PH_OBJECT_TYPE_USE_FREE_LIST)
    {
        assert(ObjectType->FreeList.Size == PhAddObjectHeaderSize(ObjectSize));

        objectHeader = PhAllocateFromFreeList(&ObjectType->FreeList);
        objectHeader->Flags = PH_OBJECT_FROM_TYPE_FREE_LIST;
        REF_STAT_UP(RefObjectsAllocatedFromTypeFreeList);
    }
    else if (ObjectSize <= PH_OBJECT_SMALL_OBJECT_SIZE)
    {
        objectHeader = PhAllocateFromFreeList(&PhObjectSmallFreeList);
        objectHeader->Flags = PH_OBJECT_FROM_SMALL_FREE_LIST;
        REF_STAT_UP(RefObjectsAllocatedFromSmallFreeList);
    }
    else
    {
        objectHeader = PhAllocate(PhAddObjectHeaderSize(ObjectSize));
        objectHeader->Flags = 0;
        REF_STAT_UP(RefObjectsAllocated);
    }

    return objectHeader;
}

VOID PhpFreeObject(
    _In_ PPH_OBJECT_HEADER ObjectHeader
    )
{
    PPH_OBJECT_TYPE objectType;

    objectType = PhObjectTypeTable[ObjectHeader->TypeIndex];

    // Object type statistics.
    _InterlockedDecrement(&objectType->NumberOfObjects);

#ifdef DEBUG
    PhAcquireQueuedLockExclusive(&PhDbgObjectListLock);
    RemoveEntryList(&ObjectHeader->ObjectListEntry);
    PhReleaseQueuedLockExclusive(&PhDbgObjectListLock);
#endif

    REF_STAT_UP(RefObjectsDestroyed);

    // Call the delete procedure if we have one.
    if (objectType->DeleteProcedure)
    {
        objectType->DeleteProcedure(PhObjectHeaderToObject(ObjectHeader), 0);
    }

    if (ObjectHeader->Flags & PH_OBJECT_FROM_TYPE_FREE_LIST)
    {
        PhFreeToFreeList(&objectType->FreeList, ObjectHeader);
        REF_STAT_UP(RefObjectsFreedToTypeFreeList);
    }
    else if (ObjectHeader->Flags & PH_OBJECT_FROM_SMALL_FREE_LIST)
    {
        PhFreeToFreeList(&PhObjectSmallFreeList, ObjectHeader);
        REF_STAT_UP(RefObjectsFreedToSmallFreeList);
    }
    else
    {
        PhFree(ObjectHeader);
        REF_STAT_UP(RefObjectsFreed);
    }
}

VOID PhpDeferDeleteObject(
    _In_ PPH_OBJECT_HEADER ObjectHeader
    )
{
    PSLIST_ENTRY oldFirstEntry;

    oldFirstEntry = RtlFirstEntrySList(&PhObjectDeferDeleteListHead);
    RtlInterlockedPushEntrySList(&PhObjectDeferDeleteListHead, &ObjectHeader->DeferDeleteListEntry);
    REF_STAT_UP(RefObjectsDeleteDeferred);

    // Was the to-free list empty before? If so, we need to queue a work item.
    if (!oldFirstEntry)
    {
        PhQueueItemGlobalWorkQueue(PhpDeferDeleteObjectRoutine, NULL);
    }
}

NTSTATUS PhpDeferDeleteObjectRoutine(
    _In_ PVOID Parameter
    )
{
    PSLIST_ENTRY listEntry;
    PPH_OBJECT_HEADER objectHeader;

    // Clear the list and obtain the first object to free.
    listEntry = RtlInterlockedFlushSList(&PhObjectDeferDeleteListHead);

    while (listEntry)
    {
        objectHeader = CONTAINING_RECORD(listEntry, PH_OBJECT_HEADER, DeferDeleteListEntry);
        listEntry = listEntry->Next;

        PhpFreeObject(objectHeader);
    }

    return STATUS_SUCCESS;
}

PVOID PhCreateAlloc(
    _In_ SIZE_T Size
    )
{
    return PhCreateObject(Size, PhAllocType);
}

FORCEINLINE PPH_AUTO_POOL PhpGetCurrentAutoPool(
    VOID
    )
{
    return (PPH_AUTO_POOL)TlsGetValue(PhpAutoPoolTlsIndex);
}

_May_raise_ FORCEINLINE VOID PhpSetCurrentAutoPool(
    _In_ PPH_AUTO_POOL AutoPool
    )
{
    if (!TlsSetValue(PhpAutoPoolTlsIndex, AutoPool))
        PhRaiseStatus(STATUS_UNSUCCESSFUL);

#ifdef DEBUG
    {
        PPHP_BASE_THREAD_DBG dbg;

        dbg = (PPHP_BASE_THREAD_DBG)TlsGetValue(PhDbgThreadDbgTlsIndex);

        if (dbg)
        {
            dbg->CurrentAutoPool = AutoPool;
        }
    }
#endif
}

VOID PhInitializeAutoPool(
    _Out_ PPH_AUTO_POOL AutoPool
    )
{
    AutoPool->StaticCount = 0;
    AutoPool->DynamicCount = 0;
    AutoPool->DynamicAllocated = 0;
    AutoPool->DynamicObjects = NULL;

    // Add the pool to the stack.
    AutoPool->NextPool = PhpGetCurrentAutoPool();
    PhpSetCurrentAutoPool(AutoPool);

    REF_STAT_UP(RefAutoPoolsCreated);
}

_May_raise_ VOID PhDeleteAutoPool(
    _Inout_ PPH_AUTO_POOL AutoPool
    )
{
    PhDrainAutoPool(AutoPool);

    if (PhpGetCurrentAutoPool() != AutoPool)
        PhRaiseStatus(STATUS_UNSUCCESSFUL);

    // Remove the pool from the stack.
    PhpSetCurrentAutoPool(AutoPool->NextPool);

    // Free the dynamic array if it hasn't been freed yet.
    if (AutoPool->DynamicObjects)
        PhFree(AutoPool->DynamicObjects);

    REF_STAT_UP(RefAutoPoolsDestroyed);
}

VOID PhDrainAutoPool(
    _In_ PPH_AUTO_POOL AutoPool
    )
{
    ULONG i;

    for (i = 0; i < AutoPool->StaticCount; i++)
        PhDereferenceObject(AutoPool->StaticObjects[i]);

    AutoPool->StaticCount = 0;

    if (AutoPool->DynamicObjects)
    {
        for (i = 0; i < AutoPool->DynamicCount; i++)
        {
            PhDereferenceObject(AutoPool->DynamicObjects[i]);
        }

        AutoPool->DynamicCount = 0;

        if (AutoPool->DynamicAllocated > PH_AUTO_POOL_DYNAMIC_BIG_SIZE)
        {
            AutoPool->DynamicAllocated = 0;
            PhFree(AutoPool->DynamicObjects);
            AutoPool->DynamicObjects = NULL;
        }
    }
}

_May_raise_ PVOID PhAutoDereferenceObject(
    _In_opt_ PVOID Object
    )
{
    PPH_AUTO_POOL autoPool = PhpGetCurrentAutoPool();

#ifdef DEBUG
    // If we don't have an auto-dereference pool, we don't want to leak the
    // object (unlike what Apple does with NSAutoreleasePool).
    if (!autoPool)
        PhRaiseStatus(STATUS_UNSUCCESSFUL);
#endif

    if (!Object)
        return NULL;

    // See if we can use the static array.
    if (autoPool->StaticCount < PH_AUTO_POOL_STATIC_SIZE)
    {
        autoPool->StaticObjects[autoPool->StaticCount++] = Object;
        return Object;
    }

    // Use the dynamic array.

    // Allocate the array if we haven't already.
    if (!autoPool->DynamicObjects)
    {
        autoPool->DynamicAllocated = 64;
        autoPool->DynamicObjects = PhAllocate(
            sizeof(PVOID) * autoPool->DynamicAllocated
            );
        REF_STAT_UP(RefAutoPoolsDynamicAllocated);
    }

    // See if we need to resize the array.
    if (autoPool->DynamicCount == autoPool->DynamicAllocated)
    {
        autoPool->DynamicAllocated *= 2;
        autoPool->DynamicObjects = PhReAllocate(
            autoPool->DynamicObjects,
            sizeof(PVOID) * autoPool->DynamicAllocated
            );
        REF_STAT_UP(RefAutoPoolsDynamicResized);
    }

    autoPool->DynamicObjects[autoPool->DynamicCount++] = Object;

    return Object;
}

VOID PhaDereferenceObject(PVOID Object)
{
    if (!Object)
        PhRaiseStatus(STATUS_INVALID_PARAMETER);

    PhAutoDereferenceObject(Object);
}