// Kernel APC infrastructure. This file owns APC queuing, delivery, and user-mode return transitions.

#include <kernel/ke_apc.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_irql.h>
#include <kernel/ke_trap.h>
#include <kernel/ke_debug.h>
#include <kernel/mm_virt.h>
#include <string.h>
#include <intrin.h>

#define KI_USER_CS 0x1BULL
#define KI_USER_SS 0x23ULL
#define KI_RFLAGS_FIXED 0x2ULL
#define KI_RFLAGS_IOPL 0x3000ULL
#define KI_RFLAGS_NT 0x4000ULL
#define KI_RFLAGS_VM 0x20000ULL

static volatile long long gUserApcReturnTokenSeed;

// Returns whether an address is in the canonical user half accepted by IRET.
static bool KiUserAddressValid(uint64_t address) {
    return address >= USER_BASE && address < USER_LIMIT;
}

// Returns whether a saved user RFLAGS value can be restored safely.
static bool KiUserRflagsValid(uint64_t rflags) {
    if ((rflags & KI_RFLAGS_FIXED) == 0)
        return false;
    if ((rflags & (KI_RFLAGS_IOPL | KI_RFLAGS_NT | KI_RFLAGS_VM)) != 0)
        return false;
    return true;
}

// Validates the kernel-owned state that will be restored after a user APC.
static KSTATUS KiValidateUserApcResumeState(uint64_t rip, uint64_t rsp, uint64_t rflags,
                                            uint64_t cs, uint64_t ss) {
    if (!KiUserAddressValid(rip) || !KiUserAddressValid(rsp))
        return STATUS_ACCESS_DENIED;
    if (cs != KI_USER_CS || ss != KI_USER_SS)
        return STATUS_ACCESS_DENIED;
    if (!KiUserRflagsValid(rflags))
        return STATUS_ACCESS_DENIED;
    return STATUS_SUCCESS;
}

// Allocates a nonzero token for one active user APC return frame.
static uint64_t KiNextUserApcReturnToken(THREAD* thread, uintptr_t frameAddress) {
    uint64_t seed = (uint64_t)_InterlockedIncrement64(&gUserApcReturnTokenSeed);
    uint64_t token = seed ^ (uint64_t)(uintptr_t)thread ^ (uint64_t)frameAddress ^ SYS_USER_APC_FRAME_MAGIC;
    if (token == 0)
        token = SYS_USER_APC_FRAME_MAGIC;
    return token;
}

// Initializes apc State.
void KiInitializeApcState(THREAD* thread) {
    if (!thread) return;
    LIST_INIT(&thread->ApcQueue);
    thread->ApcLock = (SPINLOCK)SPINLOCK_INIT;
    thread->KernelApcPending = 0;
    thread->KernelApcDisable = 0;
    thread->ApcDeliveryActive = 0;

    LIST_INIT(&thread->UserApcQueue);
    thread->UserApcLock = (SPINLOCK)SPINLOCK_INIT;
    thread->UserApcPending = 0;
    thread->UserApcDispatcher = 0;
    thread->UserApcDeliverable = 0;
    thread->UserApcDeliveryActive = 0;
    thread->UserApcFrameUserAddress = 0;
    thread->UserApcReturnToken = 0;
    thread->UserApcSavedRip = 0;
    thread->UserApcSavedRsp = 0;
    thread->UserApcSavedRflags = 0;
    thread->UserApcSavedRax = 0;
}

__declspec(dllexport) void KeInitializeApc(KAPC* apc,
                                           THREAD* thread,
                                           PKKERNEL_ROUTINE kernelRoutine,
                                           PKNORMAL_ROUTINE normalRoutine,
                                           void* normalContext) {
    if (!apc) return;

    LIST_INIT(&apc->Link);
    apc->Thread = thread;
    apc->KernelRoutine = kernelRoutine;
    apc->NormalRoutine = normalRoutine;
    apc->RundownRoutine = (PKRUNDOWN_ROUTINE)0;
    apc->NormalContext = normalContext;
    apc->SystemArgument1 = (void*)0;
    apc->SystemArgument2 = (void*)0;
    apc->Inserted = 0;
}

__declspec(dllexport) void KeSetApcRundownRoutine(KAPC* apc,
                                                    PKRUNDOWN_ROUTINE rundownRoutine) {
    if (!apc) return;
    apc->RundownRoutine = rundownRoutine;
}

__declspec(dllexport) bool KeInsertQueueApc(KAPC* apc,
                                            void* systemArgument1,
                                            void* systemArgument2) {
    if (!apc || !apc->Thread)
        return false;

    THREAD* thread = apc->Thread;
    if (thread->State == ThreadDead)
        return false;

    if (_InterlockedCompareExchange(&apc->Inserted, 1, 0) != 0)
        return false;

    apc->SystemArgument1 = systemArgument1;
    apc->SystemArgument2 = systemArgument2;

    uintptr_t flags = KeAcquireSpinlock(&thread->ApcLock);
    ListInsertTail(&thread->ApcQueue, &apc->Link);
    thread->KernelApcPending = 1;
    uint8_t wakeAlertable = (thread->State == ThreadWaiting &&
                            thread->AlertableWait &&
                            thread->KernelApcDisable == 0);
    KeReleaseSpinlock(&thread->ApcLock, flags);

    // If the target opted into an alertable wait, make the APC itself the
    // wait completion. The target thread will deliver the APC before
    // KeWaitForEvent(..., true) returns STATUS_KERNEL_APC.
    if (wakeAlertable)
        KiWakeThreadForKernelApc(thread);

    // If the caller queued an APC to itself at PASSIVE_LEVEL, try to deliver
    // immediately. Cross-thread APCs run the next time the target thread
    // returns to PASSIVE_LEVEL.
    if (thread == KeGetCurrentThread())
        KiDeliverKernelApcs();

    return true;
}

// Implements the enter Critical Region operation.
__declspec(dllexport) void KeEnterCriticalRegion(void) {
    THREAD* thread = KeGetCurrentThread();
    if (!thread) return;

    uintptr_t flags = KeAcquireSpinlock(&thread->ApcLock);
    thread->KernelApcDisable++;
    KeReleaseSpinlock(&thread->ApcLock, flags);
}

// Implements the leave Critical Region operation.
__declspec(dllexport) void KeLeaveCriticalRegion(void) {
    THREAD* thread = KeGetCurrentThread();
    if (!thread) return;

    uintptr_t flags = KeAcquireSpinlock(&thread->ApcLock);
    if (thread->KernelApcDisable > 0)
        thread->KernelApcDisable--;
    int deliver = (thread->KernelApcDisable == 0 && thread->KernelApcPending);
    KeReleaseSpinlock(&thread->ApcLock, flags);

    if (deliver)
        KiDeliverKernelApcs();
}

// Implements the are Kernel Apcs Disabled operation.
__declspec(dllexport) bool KeAreKernelApcsDisabled(void) {
    THREAD* thread = KeGetCurrentThread();
    if (!thread) return true;
    return thread->KernelApcDisable != 0;
}

// Implements the pop Kernel Apc operation.
static KAPC* KiPopKernelApc(THREAD* thread) {
    uintptr_t flags = KeAcquireSpinlock(&thread->ApcLock);
    LIST_ENTRY* node = ListPopHead(&thread->ApcQueue);
    if (LIST_IS_EMPTY(&thread->ApcQueue))
        thread->KernelApcPending = 0;
    KeReleaseSpinlock(&thread->ApcLock, flags);

    if (!node)
        return (KAPC*)0;
    return CONTAINING_RECORD(node, KAPC, Link);
}

// Implements the deliver Kernel Apcs operation.
void KiDeliverKernelApcs(void) {
    THREAD* thread = KeGetCurrentThread();
    if (!thread)
        return;

    if (KeGetCurrentIrql() != PASSIVE_LEVEL)
        return;

    if (thread->ApcDeliveryActive || thread->KernelApcDisable != 0 || !thread->KernelApcPending)
        return;

    thread->ApcDeliveryActive = 1;

    for (;;) {
        if (KeGetCurrentIrql() != PASSIVE_LEVEL || thread->KernelApcDisable != 0)
            break;

        KAPC* apc = KiPopKernelApc(thread);
        if (!apc)
            break;

        apc->Inserted = 0;

        PKNORMAL_ROUTINE normalRoutine = apc->NormalRoutine;
        void* normalContext = apc->NormalContext;
        void* arg1 = apc->SystemArgument1;
        void* arg2 = apc->SystemArgument2;

        if (apc->KernelRoutine)
            apc->KernelRoutine(apc,
                               (void**)&normalRoutine,
                               &normalContext,
                               &arg1,
                               &arg2);

        // After the kernel routine has had a chance to adjust the callout,
        // stop touching APC storage before invoking the normal routine. The
        // normal routine is allowed to free the APC's containing allocation.
        apc->SystemArgument1 = (void*)0;
        apc->SystemArgument2 = (void*)0;

        if (normalRoutine && KeGetCurrentIrql() == PASSIVE_LEVEL)
            normalRoutine(normalContext, arg1, arg2);
    }

    thread->ApcDeliveryActive = 0;
}

// Updates user Apc Dispatcher.
__declspec(dllexport) KSTATUS KeSetUserApcDispatcher(void* dispatcher) {
    THREAD* thread = KeGetCurrentThread();
    if (!thread) return STATUS_FAILURE;
    if (!dispatcher) return STATUS_INVALID_PARAMETER;
    if (!KiUserAddressValid((uint64_t)(uintptr_t)dispatcher)) return STATUS_ACCESS_DENIED;

    KSTATUS s = MmProbeUserRead(dispatcher, 1);
    if (!KSUCCESS(s)) return s;

    thread->UserApcDispatcher = (uintptr_t)dispatcher;
    return STATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS KeQueueUserApc(THREAD* thread,
                                             void* routine,
                                             void* normalContext,
                                             void* systemArgument1,
                                             void* systemArgument2) {
    if (!thread || !routine)
        return STATUS_INVALID_PARAMETER;
    if (!KiUserAddressValid((uint64_t)(uintptr_t)routine))
        return STATUS_ACCESS_DENIED;
    if (thread->State == ThreadDead)
        return STATUS_OBJECT_DEAD;

    KUSER_APC* apc = (KUSER_APC*)MmAllocatePoolTag(sizeof(KUSER_APC), MM_POOL_TAG('a', 'p', 'c', ' '));
    if (!apc)
        return STATUS_NO_MEMORY;

    memset(apc, 0, sizeof(*apc));
    LIST_INIT(&apc->Link);
    apc->Routine = routine;
    apc->NormalContext = normalContext;
    apc->SystemArgument1 = systemArgument1;
    apc->SystemArgument2 = systemArgument2;

    uintptr_t flags = KeAcquireSpinlock(&thread->UserApcLock);
    ListInsertTail(&thread->UserApcQueue, &apc->Link);
    thread->UserApcPending = 1;
    uint8_t wakeAlertable = (thread->State == ThreadWaiting &&
                             thread->AlertableWait &&
                             thread->UserApcDispatcher != 0);
    KeReleaseSpinlock(&thread->UserApcLock, flags);

    if (wakeAlertable)
        KiWakeThreadForUserApc(thread);

    return STATUS_SUCCESS;
}

// Implements the pop User Apc operation.
static KUSER_APC* KiPopUserApc(THREAD* thread) {
    uintptr_t flags = KeAcquireSpinlock(&thread->UserApcLock);
    LIST_ENTRY* node = ListPopHead(&thread->UserApcQueue);
    if (LIST_IS_EMPTY(&thread->UserApcQueue))
        thread->UserApcPending = 0;
    KeReleaseSpinlock(&thread->UserApcLock, flags);

    if (!node)
        return (KUSER_APC*)0;
    return CONTAINING_RECORD(node, KUSER_APC, Link);
}

// Implements the check For User Apc Delivery operation.
void KiCheckForUserApcDelivery(KTRAP_FRAME* frame) {
    if (!frame || !KeTrapFrameFromUser(frame))
        return;

    THREAD* thread = KeGetCurrentThread();
    if (!thread)
        return;

    if (!thread->UserApcDeliverable || thread->UserApcDeliveryActive ||
        !thread->UserApcPending || !thread->UserApcDispatcher)
        return;

    if (KeGetCurrentIrql() != PASSIVE_LEVEL)
        return;

    KUSER_APC* apc = KiPopUserApc(thread);
    if (!apc)
        return;

    KSTATUS s = KiValidateUserApcResumeState(frame->Rip, frame->Rsp, frame->Rflags,
                                             frame->Cs, frame->Ss);
    if (!KSUCCESS(s)) {
        DbgPrint("USER/APC: rejected saved return state rip=0x%016llX rsp=0x%016llX rflags=0x%016llX cs=0x%016llX ss=0x%016llX status=%u\n",
                 frame->Rip, frame->Rsp, frame->Rflags, frame->Cs, frame->Ss, s);
        uintptr_t flags = KeAcquireSpinlock(&thread->UserApcLock);
        ListInsertTail(&thread->UserApcQueue, &apc->Link);
        thread->UserApcPending = 1;
        KeReleaseSpinlock(&thread->UserApcLock, flags);
        return;
    }

    uintptr_t oldRsp = (uintptr_t)frame->Rsp;
    uintptr_t newRsp = oldRsp - sizeof(SYS_USER_APC_FRAME_K);
    newRsp &= ~(uintptr_t)0xFULL;
    newRsp -= 8; // Make injected dispatcher entry look like a normal call target.

    SYS_USER_APC_FRAME_K userFrame;
    memset(&userFrame, 0, sizeof(userFrame));
    userFrame.Magic = SYS_USER_APC_FRAME_MAGIC;
    userFrame.ReturnToken = KiNextUserApcReturnToken(thread, newRsp);
    userFrame.Routine = (uint64_t)(uintptr_t)apc->Routine;
    userFrame.NormalContext = (uint64_t)(uintptr_t)apc->NormalContext;
    userFrame.SystemArgument1 = (uint64_t)(uintptr_t)apc->SystemArgument1;
    userFrame.SystemArgument2 = (uint64_t)(uintptr_t)apc->SystemArgument2;

    s = CopyToUser((void*)newRsp, &userFrame, sizeof(userFrame));
    if (!KSUCCESS(s)) {
        uintptr_t flags = KeAcquireSpinlock(&thread->UserApcLock);
        ListInsertTail(&thread->UserApcQueue, &apc->Link);
        thread->UserApcPending = 1;
        KeReleaseSpinlock(&thread->UserApcLock, flags);
        return;
    }

    thread->UserApcFrameUserAddress = newRsp;
    thread->UserApcReturnToken = userFrame.ReturnToken;
    thread->UserApcSavedRip = frame->Rip;
    thread->UserApcSavedRsp = frame->Rsp;
    thread->UserApcSavedRflags = frame->Rflags;
    thread->UserApcSavedRax = STATUS_USER_APC;

    frame->Rip = thread->UserApcDispatcher;
    frame->Rsp = newRsp;
    frame->Rax = STATUS_USER_APC;

    thread->UserApcDeliverable = 0;
    thread->UserApcDeliveryActive = 1;

    MmFreePool(apc);
}

// Implements the rundown Apc State operation.
void KiRundownApcState(THREAD* thread) {
    if (!thread) return;

    for (;;) {
        uintptr_t flags = KeAcquireSpinlock(&thread->ApcLock);
        LIST_ENTRY* node = ListPopHead(&thread->ApcQueue);
        if (LIST_IS_EMPTY(&thread->ApcQueue))
            thread->KernelApcPending = 0;
        KeReleaseSpinlock(&thread->ApcLock, flags);

        if (!node) break;
        KAPC* apc = CONTAINING_RECORD(node, KAPC, Link);
        apc->Inserted = 0;
        if (apc->RundownRoutine)
            apc->RundownRoutine(apc);
    }

    for (;;) {
        uintptr_t flags = KeAcquireSpinlock(&thread->UserApcLock);
        LIST_ENTRY* node = ListPopHead(&thread->UserApcQueue);
        if (LIST_IS_EMPTY(&thread->UserApcQueue))
            thread->UserApcPending = 0;
        KeReleaseSpinlock(&thread->UserApcLock, flags);

        if (!node) break;
        KUSER_APC* apc = CONTAINING_RECORD(node, KUSER_APC, Link);
        MmFreePool(apc);
    }

    thread->UserApcDispatcher = 0;
    thread->UserApcDeliverable = 0;
    thread->UserApcDeliveryActive = 0;
    thread->UserApcFrameUserAddress = 0;
    thread->UserApcReturnToken = 0;
    thread->UserApcSavedRip = 0;
    thread->UserApcSavedRsp = 0;
    thread->UserApcSavedRflags = 0;
    thread->UserApcSavedRax = 0;
}

// Implements the user Apc Return operation.
__declspec(dllexport) uint64_t KeUserApcReturn(const SYS_USER_APC_FRAME_K* userFrame) {
    THREAD* thread = KeGetCurrentThread();
    if (!thread)
        return STATUS_FAILURE;

    if (!thread->UserApcDeliveryActive || !userFrame)
        return STATUS_INVALID_PARAMETER;

    SYS_USER_APC_FRAME_K frameCopy;
    KSTATUS s = CopyFromUser(&frameCopy, userFrame, sizeof(frameCopy));
    if (!KSUCCESS(s))
        return s;

    if (frameCopy.Magic != SYS_USER_APC_FRAME_MAGIC ||
        frameCopy.ReturnToken != thread->UserApcReturnToken ||
        (uintptr_t)userFrame != thread->UserApcFrameUserAddress) {
        DbgPrint("USER/APC: rejected return frame user=0x%016llX active=0x%016llX token=0x%016llX expected=0x%016llX\n",
                 (uint64_t)(uintptr_t)userFrame,
                 (uint64_t)thread->UserApcFrameUserAddress,
                 frameCopy.ReturnToken,
                 thread->UserApcReturnToken);
        return STATUS_INVALID_PARAMETER;
    }

    KTRAP_FRAME* frame = KeGetCurrentTrapFrame();
    if (!frame || !KeTrapFrameFromUser(frame) || frame->Kind != KTRAP_KIND_SYSCALL)
        return STATUS_FAILURE;

    s = KiValidateUserApcResumeState(thread->UserApcSavedRip,
                                     thread->UserApcSavedRsp,
                                     thread->UserApcSavedRflags,
                                     frame->Cs,
                                     frame->Ss);
    if (!KSUCCESS(s)) {
        DbgPrint("USER/APC: rejected kernel-owned resume state rip=0x%016llX rsp=0x%016llX rflags=0x%016llX status=%u\n",
                 thread->UserApcSavedRip,
                 thread->UserApcSavedRsp,
                 thread->UserApcSavedRflags,
                 s);
        return s;
    }

    uint64_t savedRax = thread->UserApcSavedRax;
    frame->Rip = thread->UserApcSavedRip;
    frame->Rsp = thread->UserApcSavedRsp;
    frame->Rflags = thread->UserApcSavedRflags;
    frame->Rax = savedRax;

    thread->UserApcDeliveryActive = 0;
    thread->UserApcDeliverable = 0;
    thread->UserApcFrameUserAddress = 0;
    thread->UserApcReturnToken = 0;
    thread->UserApcSavedRip = 0;
    thread->UserApcSavedRsp = 0;
    thread->UserApcSavedRflags = 0;
    thread->UserApcSavedRax = 0;
    return savedRax;
}