mirror-ac/driver/apc.c

#include "apc.h"

#include "driver.h"
#include "imports.h"

VOID
GetApcContextByIndex(_Out_ PVOID* Context, _In_ INT Index)
{
        AcquireDriverConfigLock();
        *Context = GetApcContextArray()[Index];
        ReleaseDriverConfigLock();
}

VOID
GetApcContext(_Out_ PVOID* Context, _In_ LONG ContextIdentifier)
{
        AcquireDriverConfigLock();

        for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
                PAPC_CONTEXT_HEADER header = GetApcContextArray()[index];

                if (!header)
                        continue;

                if (header->context_id != ContextIdentifier)
                        continue;

                *Context = header;
                goto unlock;
        }
unlock:
        ReleaseDriverConfigLock();
}

/*
 * No need to hold the lock here as the thread freeing the APCs will
 * already hold the configuration lock. We also dont want to release and
 * reclaim the lock before calling this function since we need to ensure
 * we hold the lock during the entire decrement and free process.
 */
BOOLEAN
FreeApcContextStructure(_Out_ PAPC_CONTEXT_HEADER Context)
{
        DEBUG_VERBOSE("All APCs executed, freeing context structure");

        for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
                PUINT64 entry = GetApcContextArray();

                if (entry[index] != Context)
                        continue;

                if (Context->count > 0)
                        return FALSE;

                ImpExFreePoolWithTag(Context, POOL_TAG_APC);
                entry[index] = NULL;
                return TRUE;
        }

        return FALSE;
}

VOID
IncrementApcCount(_In_ LONG ContextId)
{
        PAPC_CONTEXT_HEADER header = NULL;
        GetApcContext(&header, ContextId);

        if (!header)
                return;

        /* i actually dont think we need this lock here */
        AcquireDriverConfigLock();
        header->count += 1;
        ReleaseDriverConfigLock();
}

VOID
FreeApcAndDecrementApcCount(_Inout_ PRKAPC Apc, _In_ LONG ContextId)
{
        PAPC_CONTEXT_HEADER context = NULL;

        ImpExFreePoolWithTag(Apc, POOL_TAG_APC);
        GetApcContext(&context, ContextId);

        if (!context)
                return;

        AcquireDriverConfigLock();
        context->count -= 1;
        ReleaseDriverConfigLock();
}

/*
 * The reason we use a query model rather then checking the count of queued APCs
 * after each APC free and decrement is that the lock will be recursively
 * acquired by freeing threads (i.e executing APCs) rather then APC allocation
 * threads. The reason for this being that freeing threads are executing at a
 * higher IRQL then the APC allocation thread, hence they are granted higher
 * priority by the scheduler when determining which thread will accquire the
 * lock next:
 *
 * [+] Freeing thread -> ApcKernelRoutine IRQL: 1 (APC_LEVEL)
 * [+] Allocation thread -> ValidateThreadViaKernelApcCallback IRQL: 0
 * (PASSIVE_LEVEL)
 *
 * As a result, once an APC is executed and reaches the freeing stage, it will
 * acquire the lock and decrement it. Then, if atleast 1 APC execution thread is
 * waiting on the lock, it will be prioritised due to its higher IRQL and the
 * cycle will continue. Eventually, the count will reach 0 due to recursive
 * acquisition by the executing APC threads and then the function will free the
 * APC context structure. This will then cause a bug check the next time a
 * thread accesses the context structure and hence not good :c.
 *
 * So to combat this, we add in a flag specifying whether or not an allocation
 * of APCs is in progress, and even if the count is 0 we will not free the
 * context structure until the count is 0 and allocation_in_progress is 0. We
 * can then call this function alongside other query callbacks via IOCTL to
 * constantly monitor the status of open APC contexts.
 */
NTSTATUS
QueryActiveApcContextsForCompletion()
{
        for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
                PAPC_CONTEXT_HEADER entry = NULL;
                GetApcContextByIndex(&entry, index);
                AcquireDriverConfigLock();

                if (!entry)
                        goto increment;

                if (entry->count > 0 || entry->allocation_in_progress == TRUE)
                        goto increment;

                switch (entry->context_id) {
                case APC_CONTEXT_ID_STACKWALK:
                        FreeApcStackwalkApcContextInformation(entry);
                        FreeApcContextStructure(entry);
                        break;
                }

        increment:
                ReleaseDriverConfigLock();
        }
        return STATUS_SUCCESS;
}

VOID
InsertApcContext(_In_ PVOID Context)
{
        if (IsDriverUnloading())
                return STATUS_UNSUCCESSFUL;

        AcquireDriverConfigLock();
        PAPC_CONTEXT_HEADER header = Context;

        for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
                PUINT64 entry = GetApcContextArray();

                if (entry[index] == NULL) {
                        entry[index] = Context;
                        goto end;
                }
        }
end:
        ReleaseDriverConfigLock();
}

/*
 * The driver config structure holds an array of pointers to APC context
 * structures. These APC context structures are unique to each APC operation
 * that this driver will perform. For example, a single context will manage all
 * APCs that are used to stackwalk, whilst another context will be used to
 * manage all APCs used to query a threads memory for example.
 *
 * Due to the nature of APCs, its important to keep a total or count of the
 * number of APCs we have allocated and queued to threads. This information is
 * stored in the APC_CONTEXT_HEADER which all APC context structures will
 * contain as the first entry in their structure. It holds the ContextId which
 * is a unique identifier for the type of APC operation it is managing aswell as
 * the number of currently queued APCs.
 *
 * When an APC is allocated a queued, we increment this count. When an APC is
 * completed and freed, we decrement this counter and free the APC itself. If
 * all APCs have been freed and the counter is 0,the following objects will be
 * freed:
 *
 * 1. Any additional allocations used by the APC stored in the context structure
 * 2. The APC context structure for the given APC operation
 * 3. The APC context entry in g_DriverConfig->>apc_contexts will be zero'd.
 *
 * It's important to remember that the driver can unload when pending APC's have
 * not been freed due to the limitations windows places on APCs, however I am in
 * the process of finding a solution for this.
 */
BOOLEAN
DrvUnloadFreeAllApcContextStructures()
{
        AcquireDriverConfigLock();

        for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
                PUINT64 entry = GetApcContextArray();

                if (entry[index] == NULL)
                        continue;

                PAPC_CONTEXT_HEADER context = entry[index];

                if (context->count > 0) {
                        ReleaseDriverConfigLock();
                        return FALSE;
                }

                ImpExFreePoolWithTag(entry, POOL_TAG_APC);
        }
unlock:
        ReleaseDriverConfigLock();
        return TRUE;
}