// Object manager. This file owns reference-counted kernel objects and their terminal lifetime transitions.

#include <ntlite/atomic.h>
#include <stdbool.h>
#include <string.h>
#include <intrin.h>
#include <kernel/ob_object.h>
#include <kernel/ke_debug.h>
#include <kernel/mm_virt.h>

// Returns a diagnostic pool tag for the object type.
static uint32_t KiObjectPoolTag(OBJECT_TYPE type) {
    switch (type) {
    case ObjectTypeProcess:     return MM_POOL_TAG('p', 'r', 'o', 'c');
    case ObjectTypeThread:      return MM_POOL_TAG('t', 'h', 'r', 'd');
    case ObjectTypeMutex:       return MM_POOL_TAG('m', 't', 'x', ' ');
    case ObjectTypeSemaphore:   return MM_POOL_TAG('s', 'e', 'm', ' ');
    case ObjectTypeEvent:       return MM_POOL_TAG('e', 'v', 'n', 't');
    case ObjectTypeQueue:       return MM_POOL_TAG('q', 'u', 'e', 'u');
    case ObjectTypeFile:        return MM_POOL_TAG('f', 'i', 'l', 'e');
    case ObjectTypeDir:         return MM_POOL_TAG('d', 'i', 'r', ' ');
    case ObjectTypeAsyncIo:     return MM_POOL_TAG('a', 'i', 'o', ' ');
    case ObjectTypeSocket:      return MM_POOL_TAG('s', 'o', 'c', 'k');
    case ObjectTypeDevice:      return MM_POOL_TAG('d', 'e', 'v', ' ');
    case ObjectTypeDriver:      return MM_POOL_TAG('d', 'r', 'v', 'r');
    case ObjectTypeUsbDevice:   return MM_POOL_TAG('u', 's', 'b', 'd');
    case ObjectTypeUsbPipe:     return MM_POOL_TAG('u', 's', 'b', 'p');
    case ObjectTypeFileMapping: return MM_POOL_TAG('m', 'a', 'p', 'f');
    case ObjectTypeSurface:     return MM_POOL_TAG('s', 'u', 'r', 'f');
    case ObjectTypeImageSection: return MM_POOL_TAG('i', 'm', 's', 'e');
    case ObjectTypeTimer:       return MM_POOL_TAG('t', 'm', 'r', ' ');
    case ObjectTypeDisplayClock: return MM_POOL_TAG('d', 'c', 'l', 'k');
    case ObjectTypeAudioEndpoint: return MM_POOL_TAG('a', 'e', 'p', 't');
    case ObjectTypeAudioStream:   return MM_POOL_TAG('a', 's', 't', 'r');
    case ObjectTypeSecurityContext: return MM_POOL_TAG('s', 'e', 'c', 'x');
    case ObjectTypePort:       return MM_POOL_TAG('p', 'o', 'r', 't');
    case ObjectTypeVolume:     return MM_POOL_TAG('v', 'o', 'l', 'm');
    case ObjectTypeVolumeManager: return MM_POOL_TAG('v', 'm', 'g', 'r');
    case ObjectTypeInputManager: return MM_POOL_TAG('i', 'm', 'g', 'r');
    case ObjectTypeDisplayManager: return MM_POOL_TAG('d', 'm', 'g', 'r');
    case ObjectTypePipeEndpoint: return MM_POOL_TAG('p', 'i', 'p', 'e');
    case ObjectTypeServiceRegistration: return MM_POOL_TAG('s', 'r', 'e', 'g');
    default:                    return MM_POOL_TAG('o', 'b', 'j', ' ');
    }
}

// Initializes a caller-owned object header in the live, inserted state.
__declspec(dllexport) void ObInitializeObjectHeader(OBJECT* object, OBJECT_TYPE type) {
    if (!object)
        return;

    object->Type = type;
    object->RefCount = 1;
    object->LifetimeState = ObjectAlive;
    object->Flags = OBJECT_FLAG_INSERTED;
    object->Destroy = (void*)0;
}

// Creates a zeroed, tagged object with one initial reference.
__declspec(dllexport) OBJECT* ObAllocateObject(OBJECT_TYPE type, size_t size) {
    if (size < sizeof(OBJECT)) size = sizeof(OBJECT);

    OBJECT* object = (OBJECT*)MmAllocatePoolTag(size, KiObjectPoolTag(type));
    if (!object) return (OBJECT*)0;

    memset(object, 0, size);
    ObInitializeObjectHeader(object, type);
    return object;
}

// Acquires a reference without allowing zero-to-one resurrection.
__declspec(dllexport) bool ObTryReferenceObject(OBJECT* object) {
    if (!object)
        return false;

    while (true) {
        if (NtlAtomicLoad32(&object->LifetimeState) != ObjectAlive)
            return false;
        if ((NtlAtomicLoadU32(&object->Flags) & OBJECT_FLAG_INSERTED) == 0)
            return false;

        int32_t refs = NtlAtomicLoad32(&object->RefCount);
        if (refs <= 0 || refs == 0x7fffffffL)
            return false;

        if (NtlAtomicCompareExchange32(&object->RefCount, refs + 1, refs) == refs)
            return true;
    }
}

// Acquires a required reference and catches stale pointers at the caller site.
__declspec(dllexport) void ObReferenceObject(OBJECT* object) {
    if (!object)
        return;
    if (ObTryReferenceObject(object))
        return;

    KePanic(KE_PANIC_CODE_GENERIC,
            "object: stale reference object=0x%016llX type=%u refs=%ld state=%ld flags=0x%08x\n",
            (uint64_t)(uintptr_t)object,
            (unsigned)object->Type,
            object->RefCount,
            object->LifetimeState,
            (unsigned)object->Flags);
}

// Releases a reference and performs the one-way Alive-to-Dying-to-Dead transition.
__declspec(dllexport) void ObDereferenceObject(OBJECT* object) {
    if (!object)
        return;

    int32_t refs = NtlAtomicDecrement32(&object->RefCount);
    if (refs > 0)
        return;
    if (refs < 0) {
        KePanic(KE_PANIC_CODE_GENERIC,
                "object: reference underflow object=0x%016llX type=%u refs=%ld state=%ld flags=0x%08x\n",
                (uint64_t)(uintptr_t)object,
                (unsigned)object->Type,
                refs,
                object->LifetimeState,
                (unsigned)object->Flags);
    }

    if (NtlAtomicCompareExchange32(&object->LifetimeState, ObjectDying, ObjectAlive) != ObjectAlive) {
        KePanic(KE_PANIC_CODE_GENERIC,
                "object: duplicate terminal transition object=0x%016llX type=%u state=%ld flags=0x%08x\n",
                (uint64_t)(uintptr_t)object,
                (unsigned)object->Type,
                object->LifetimeState,
                (unsigned)object->Flags);
    }

    NtlAtomicExchangeU32(&object->Flags, OBJECT_FLAG_UNLINKED);
    if (object->Destroy)
        object->Destroy(object);

    NtlAtomicExchange32(&object->LifetimeState, ObjectDead);
    MmFreePool(object);
}