nt-lite/code/kernel/ke_thread.c

// ke_thread.c - thread model
#include <string.h>
#include <stddef.h>
#include <intrin.h>
#include <ke_thread.h>
#include <ke_sched.h>
#include <ke_irql.h>
#include <ke_event.h>
#include <ke_apc.h>
#include <ke_userenv.h>
#include <mm_virt.h>
#include <ob_object.h>

#define FPU_FCW_DEFAULT   0x037Fu
#define FPU_MXCSR_DEFAULT 0x1F80u

// These offsets are consumed directly by ke_switch.asm.  Keep the asserts
// here so future THREAD layout edits fail at build time instead of corrupting
// context switches at runtime.
#define STATIC_ASSERT_OFFSET(type, field, expected) \
    _Static_assert(offsetof(type, field) == (expected), \
                   #type "." #field " offset must match assembly")

STATIC_ASSERT_OFFSET(THREAD, SavedContext, 88);
STATIC_ASSERT_OFFSET(THREAD, TssRsp0, 104);
STATIC_ASSERT_OFFSET(THREAD, Teb, 112);
STATIC_ASSERT_OFFSET(THREAD, FpuContext, 168);

static uint32_t gNextThreadId = 1;
extern void KeThreadEntry(void);

static void KiThreadDestroy(OBJECT* obj) {
    THREAD* thread = CONTAINING_RECORD(obj, THREAD, Header);

    KiRundownApcState(thread);

    if (thread->ExitEvent) {
        ObDereferenceObject(&thread->ExitEvent->Header);
        thread->ExitEvent = (EVENT*)0;
    }

    if (thread->Owner) {
        uintptr_t flags = KeAcquireSpinlock(&thread->Owner->HandleLock);
        ListRemove(&thread->ProcessLink);
        thread->Owner->ThreadCount--;
        KeReleaseSpinlock(&thread->Owner->HandleLock, flags);
        ObDereferenceObject(&thread->Owner->Header);
        thread->Owner = (PROCESS *)0;
    }

    if (thread->KernelStackBase) {
        MmFreePool(thread->KernelStackBase);
        thread->KernelStackBase = (void *)0;
    }
    
    if (thread->Teb && thread->TebSize) {
        MmFreeUserMemory(thread->Teb);
        thread->Teb = (void*)0;
        thread->TebSize = 0;
    }

    if (thread->UserStackBase) {
        MmFreeUserMemory(thread->UserStackBase);
        thread->UserStackBase = (void *)0;
    }
}

__declspec(dllexport) KSTATUS KeCreateThread(PROCESS* owner, uintptr_t entry, uintptr_t arg,
                       THREAD_PRIORITY priority, THREAD** out) {
    if (!owner || !out) return STATUS_INVALID_PARAMETER;
    if (priority >= PriorityCount) return STATUS_INVALID_PARAMETER;

    THREAD* thread = (THREAD *)ObAllocateObject(ObjectTypeThread, sizeof(THREAD));
    if (!thread) return STATUS_NO_MEMORY;

    thread->Header.Destroy   = KiThreadDestroy;
    thread->ThreadId         = _InterlockedExchangeAdd((volatile long *)&gNextThreadId, 1);
    thread->State            = ThreadReady;
    thread->Priority         = priority;
    thread->QuantaConsumed   = 0;
    thread->WaitObject       = (OBJECT *)0;
    thread->WaitStatus       = STATUS_SUCCESS;
    thread->WaitListHead     = (LIST_ENTRY*)0;
    thread->AlertableWait    = 0;
    thread->SchedLockFlags   = 0x202ULL;
    KiInitializeApcState(thread);
    thread->TrapFrame = (KTRAP_FRAME*)0;

    void* kstack = MmAllocatePool(KSTACK_SIZE);
    if (!kstack) {
        ObDereferenceObject(&thread->Header);
        return STATUS_NO_MEMORY;
    }
    memset(kstack, 0, KSTACK_SIZE);
    thread->KernelStackBase = kstack;
    thread->KernelStackTop  = (uintptr_t)kstack + KSTACK_SIZE;
    thread->TssRsp0         = thread->KernelStackTop;

    void* ustack = MmAllocateUserMemory(USTACK_SIZE);
    if (!ustack) {
        MmFreePool(kstack);
        thread->KernelStackBase = (void*)0;
        ObDereferenceObject(&thread->Header);
        return STATUS_NO_MEMORY;
    }
    thread->UserStackBase = ustack;
    uintptr_t userStackTop = (uintptr_t)ustack + USTACK_SIZE - 8; // i don't know why, i don't want to know why
    // i shouldn't have to wonder why but the stack is off by 8 bytes unless we do this terribleness
    // this magic number will hunt me in my dreams

    SYS_TEB* teb = (SYS_TEB*)MmAllocateUserMemory(sizeof(SYS_TEB));
    if (!teb) {
        MmFreeUserMemory(ustack);
        thread->UserStackBase = (void*)0;
        MmFreePool(kstack);
        thread->KernelStackBase = (void*)0;
        ObDereferenceObject(&thread->Header);
        return STATUS_NO_MEMORY;
    }
    memset(teb, 0, sizeof(SYS_TEB));
    teb->Peb = (SYS_PEB*)owner->UserPebBase;
    teb->StackBase = (void*)userStackTop;
    teb->StackLimit = ustack;
    teb->Self = teb;
    teb->ThreadId = thread->ThreadId;
    teb->Magic = SYS_TEB_MAGIC;

    thread->Teb = teb;
    thread->TebSize = sizeof(SYS_TEB);

    thread->UserEntry = entry;
    thread->UserArg   = arg;

    uintptr_t* ksp = (uintptr_t *)thread->KernelStackTop;
    *(--ksp) = userStackTop;
    *(--ksp) = 0x23ULL;          
    *(--ksp) = userStackTop;     
    *(--ksp) = 0x202ULL;         
    *(--ksp) = 0x1BULL;          
    *(--ksp) = entry;            
    *(--ksp) = arg;              

    ksp -= sizeof(KCONTEXT) / sizeof(uintptr_t);
    KCONTEXT* ctx = (KCONTEXT *)ksp;
    memset(ctx, 0, sizeof(KCONTEXT));
    ctx->Rip = (uint64_t)KeThreadEntry;
    thread->SavedContext = ctx;

    ObReferenceObject(&owner->Header);
    thread->Owner = owner;
    
    // Allocate the exit event before the thread becomes visible. Manual-reset
    // so multiple joiners (if any) all wake up when the thread exits.
    EVENT* exitEvent;
    KSTATUS s = KeEventCreate(EventManualReset, 0, &exitEvent);
    if (!KSUCCESS(s)) {
        if (thread->Teb && thread->TebSize) MmFreeUserMemory(thread->Teb);
        thread->Teb = (void*)0;
        thread->TebSize = 0;
        MmFreeUserMemory(ustack);
        thread->UserStackBase = (void*)0;
        MmFreePool(kstack);
        thread->KernelStackBase = (void*)0;
        ObDereferenceObject(&owner->Header);
        thread->Owner = (PROCESS*)0;
        ObDereferenceObject(&thread->Header);
        return s;
    }
    thread->ExitEvent = exitEvent;

    uintptr_t flags = KeAcquireSpinlock(&owner->HandleLock);
    ListInsertTail(&owner->ThreadList, &thread->ProcessLink);
    owner->ThreadCount++;
    KeReleaseSpinlock(&owner->HandleLock, flags);

    KeInitThreadFpuContext(thread);
    KeReadyThread(thread);
    *out = thread;
    return STATUS_SUCCESS;
}

__declspec(dllexport) void KeExitThread(THREAD* thread, int code) {
    (void)code;
    if (!thread) return;

    KiRundownApcState(thread);

    // Wake any pthread_join waiters before the thread is marked dead and
    // potentially freed. Manual-reset means all concurrent waiters unblock.
    if (thread->ExitEvent)
        KeSignalEvent(thread->ExitEvent);

    thread->State = ThreadDead;

    KeYield();
}

// ---------------------------------------------------------------------------
// KeCreateKernelThread
// ---------------------------------------------------------------------------
// Kernel threads never go to ring 3: they have no user stack, no TEB, and
// their entry stub (KeKernelThreadEntry) calls the function directly then
// calls KeExitThread instead of IRETQ-ing.
//
// Stack layout built here (from KernelStackTop downward):
//   [ksp - 8 ]  arg          <-- popped by KeKernelThreadEntry as r12
//   [ksp - 16]  fn           <-- popped by KeKernelThreadEntry as rbx
//   [ksp - 16 - sizeof(KCONTEXT)]  KCONTEXT { .Rip = KeKernelThreadEntry }
// ---------------------------------------------------------------------------
extern void KeKernelThreadEntry(void);

__declspec(dllexport) KSTATUS KeCreateKernelThread(uintptr_t fn, uintptr_t arg,
                              THREAD_PRIORITY priority, THREAD** out) {
    // Kernel threads are owned by the idle process. We need a PROCESS* for
    // the owner field; grab it from the idle/init thread that already exists.
    THREAD* idleThread = KeGetCurrentThread(); // safe to call before scheduler starts
    PROCESS* owner = idleThread ? idleThread->Owner : (PROCESS*)0;
    if (!owner) return STATUS_INVALID_PARAMETER;
    if (priority >= PriorityCount) return STATUS_INVALID_PARAMETER;

    THREAD* thread = (THREAD*)ObAllocateObject(ObjectTypeThread, sizeof(THREAD));
    if (!thread) return STATUS_NO_MEMORY;

    thread->Header.Destroy  = KiThreadDestroy;
    thread->ThreadId        = _InterlockedExchangeAdd((volatile long*)&gNextThreadId, 1);
    thread->State           = ThreadReady;
    thread->Priority        = priority;
    thread->QuantaConsumed  = 0;
    thread->WaitObject      = (OBJECT*)0;
    thread->WaitStatus      = STATUS_SUCCESS;
    thread->WaitListHead    = (LIST_ENTRY*)0;
    thread->AlertableWait   = 0;
    thread->SchedLockFlags  = 0x202ULL;
    KiInitializeApcState(thread);
    thread->TrapFrame       = (KTRAP_FRAME*)0;
    thread->UserStackBase   = (void*)0;   // no user stack
    thread->Teb             = (void*)0;   // no TEB; GS_BASE will be 0 in kernel ctx
    thread->TebSize         = 0;

    void* kstack = MmAllocatePool(KSTACK_SIZE);
    if (!kstack) {
        ObDereferenceObject(&thread->Header);
        return STATUS_NO_MEMORY;
    }
    memset(kstack, 0, KSTACK_SIZE);
    thread->KernelStackBase = kstack;
    thread->KernelStackTop  = (uintptr_t)kstack + KSTACK_SIZE;
    thread->TssRsp0         = thread->KernelStackTop;

    // Build the initial stack: arg, fn, then KCONTEXT.
    uintptr_t* ksp = (uintptr_t*)thread->KernelStackTop;
    *(--ksp) = arg;   // popped second  (-> r12)
    *(--ksp) = fn;    // popped first   (-> rbx)

    ksp -= sizeof(KCONTEXT) / sizeof(uintptr_t);
    KCONTEXT* ctx = (KCONTEXT*)ksp;
    memset(ctx, 0, sizeof(KCONTEXT));
    ctx->Rip = (uint64_t)KeKernelThreadEntry;
    thread->SavedContext = ctx;

    // Register with the owning process.
    ObReferenceObject(&owner->Header);
    thread->Owner = owner;

    uintptr_t flags = KeAcquireSpinlock(&owner->HandleLock);
    ListInsertTail(&owner->ThreadList, &thread->ProcessLink);
    owner->ThreadCount++;
    KeReleaseSpinlock(&owner->HandleLock, flags);

    if (out) *out = thread;

    KeInitThreadFpuContext(thread);
    KeReadyThread(thread);
    return STATUS_SUCCESS;
}

void KePrepareBlockThread(THREAD* thread, OBJECT* waitObject, uint8_t alertable) {
    // Called with the synchronization object's spinlock held, so the caller is
    // normally at DISPATCH_LEVEL here.  This function only marks intent to
    // block; the actual scheduling point is KeCommitBlockThread at PASSIVE.
    KeAssertIrqlAtOrBelow(DISPATCH_LEVEL);
    if (!thread) return;

    thread->State         = ThreadWaiting;
    thread->WaitObject    = waitObject;
    thread->WaitStatus    = STATUS_PENDING;
    thread->WaitListHead  = (LIST_ENTRY*)0;
    thread->AlertableWait = alertable ? 1 : 0;
}


void KeSetThreadWaitList(THREAD* thread, LIST_ENTRY* waitListHead) {
    KeAssertIrqlAtOrBelow(DISPATCH_LEVEL);
    if (!thread) return;
    thread->WaitListHead = waitListHead;
}

static void KiFinishWait(THREAD* thread, KSTATUS status) {
    thread->WaitObject    = (OBJECT*)0;
    thread->WaitStatus    = status;
    thread->WaitListHead  = (LIST_ENTRY*)0;
    thread->AlertableWait = 0;
}

KSTATUS KeCommitBlockThread(THREAD* thread, uint8_t alertable) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!thread) return STATUS_INVALID_PARAMETER;

    if (alertable)
        thread->AlertableWait = 1;

    // A signal/APC may have arrived after the waiter was placed on the
    // object's wait list but before this thread reached the scheduling point.
    // In that case KeUnblockThreadWithStatus completed the wait in-place and
    // there is nothing left to block on.
    if (thread->State != ThreadWaiting) {
        KSTATUS status = thread->WaitStatus;
        thread->AlertableWait = 0;
        return status;
    }

    // If a user APC was already pending before the thread fully blocked,
    // consume the alertable wait immediately instead of going to sleep.
    if (alertable && thread->UserApcPending && thread->UserApcDispatcher) {
        if (thread->WaitListHead)
            ListRemove(&thread->SchedulerLink);
        KiFinishWait(thread, STATUS_USER_APC);
        thread->State = ThreadRunning;
        return STATUS_USER_APC;
    }

    KeYield();

    // Kernel APCs are delivered before an alertable wait returns
    // STATUS_KERNEL_APC to its caller.
    KiDeliverKernelApcs();

    thread->AlertableWait = 0;
    return thread->WaitStatus;
}


KSTATUS KeBlockThread(THREAD* thread, OBJECT* waitObject, uint8_t alertable) {
    KePrepareBlockThread(thread, waitObject, alertable);
    return KeCommitBlockThread(thread, alertable);
}


void KeUnblockThreadWithStatus(THREAD* thread, KSTATUS status) {
    if (!thread) return;

    // If the selected waiter is still the currently executing thread, the
    // signal landed in the prepare->commit window.  Complete the wait in-place
    // instead of inserting the running thread into the ready queue.
    if (thread == KeGetCurrentThread()) {
        KiFinishWait(thread, status);
        if (thread->State == ThreadWaiting)
            thread->State = ThreadRunning;
        return;
    }

    // The same object can be signaled after a timeout/cancel path exists, or a
    // buggy caller can try to wake a non-waiting thread.  Do not enqueue such a
    // thread twice.
    if (thread->State != ThreadWaiting) {
        thread->WaitStatus = status;
        return;
    }

    KiFinishWait(thread, status);
    thread->State = ThreadReady;
    
    if (thread->Priority < PriorityHigh) {
        thread->Priority = (THREAD_PRIORITY)(thread->Priority + 1);
        thread->QuantaConsumed = 0;
    }
    
    KeReadyThread(thread);
}

void KeUnblockThread(THREAD* thread) {
    KeUnblockThreadWithStatus(thread, STATUS_SUCCESS);
}

static KSTATUS KiAlertThreadWithStatus(THREAD* thread, KSTATUS status) {
    if (!thread) return STATUS_INVALID_PARAMETER;

    if (thread->State != ThreadWaiting || !thread->AlertableWait)
        return STATUS_PENDING;

    if (thread->WaitListHead)
        ListRemove(&thread->SchedulerLink);

    KeUnblockThreadWithStatus(thread, status);
    return STATUS_SUCCESS;
}

KSTATUS KeAlertThread(THREAD* thread) {
    return KiAlertThreadWithStatus(thread, STATUS_ALERTED);
}

void KiWakeThreadForKernelApc(THREAD* thread) {
    (void)KiAlertThreadWithStatus(thread, STATUS_KERNEL_APC);
}

void KiWakeThreadForUserApc(THREAD* thread) {
    (void)KiAlertThreadWithStatus(thread, STATUS_USER_APC);
}

void KeInitThreadFpuContext(THREAD* thread)
{
    uintptr_t raw =
        (uintptr_t)thread->FpuStorage;

    uintptr_t aligned =
        (raw + 15) & ~(uintptr_t)15;

    thread->FpuContext = (void*)aligned;

    memset(thread->FpuContext, 0, 512);

    *(uint16_t*)((uint8_t*)thread->FpuContext + 0x00) = 0x037F;
    *(uint32_t*)((uint8_t*)thread->FpuContext + 0x18) = 0x1F80;
}