nt-lite/code/kernel/ke_process.c

// ke_process.c - process model
#include <string.h>
#include <intrin.h>
#include <kernel/ke_process.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_event.h>
#include <kernel/ke_userenv.h>
#include <kernel/mm_virt.h>
#include <kernel/ob_object.h>

// ---------------------------------------------------------------------------
// Global process list
// ---------------------------------------------------------------------------
static LIST_ENTRY gProcessList;
static SPINLOCK   gProcessListLock;
static uint32_t   gNextProcessId = 0;
static PROCESS*   gSystemProcess = (PROCESS*)0;

static void KiProcessInit(void) {
    static int inited = 0;
    if (!inited) {
        LIST_INIT(&gProcessList);
        gProcessListLock = (SPINLOCK)SPINLOCK_INIT;
        inited = 1;
    }
}

LIST_ENTRY* KeGetProcessList(void) {
    KiProcessInit();
    return &gProcessList;
}

PROCESS* KeGetCurrentProcess(void) {
    THREAD* currentThread = KeGetCurrentThread();
    if (currentThread) {
        return currentThread->Owner;
    }
    return NULL;
}

PROCESS* KeGetSystemProcess(void) {
    return gSystemProcess;
}

KSTATUS KeCreateSystemProcess(PROCESS** out) {
    if (!out) return STATUS_INVALID_PARAMETER;
    if (gSystemProcess) {
        *out = gSystemProcess;
        return STATUS_SUCCESS;
    }

    PROCESS* proc = (PROCESS*)0;
    KSTATUS s = KeCreateProcess("system", &proc);
    if (!KSUCCESS(s)) {
        *out = (PROCESS*)0;
        return s;
    }

    gSystemProcess = proc;
    *out = proc;
    return STATUS_SUCCESS;
}

size_t KeCopyProcessIds(uint32_t* pidsOut, size_t maxCount) {
    KiProcessInit();

    uintptr_t flags = KeAcquireSpinlock(&gProcessListLock);
    size_t count = 0;
    LIST_ENTRY* head = &gProcessList;
    for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
        PROCESS* p = CONTAINING_RECORD(node, PROCESS, GlobalLink);
        if (pidsOut && count < maxCount) pidsOut[count] = p->ProcessId;
        count++;
    }
    KeReleaseSpinlock(&gProcessListLock, flags);
    return count;
}

// ---------------------------------------------------------------------------
// Destructor (called by ObDereferenceObject when RefCount hits zero)
// ---------------------------------------------------------------------------
static void KiProcessDestroy(OBJECT* obj) {
    PROCESS* proc = CONTAINING_RECORD(obj, PROCESS, Header);

    if (proc == gSystemProcess)
        gSystemProcess = (PROCESS*)0;

    // Release all remaining handles.
    for (int i = 1; i < HANDLE_TABLE_SIZE; i++) {
        if (proc->Handles[i].Object) {
            ObDereferenceObject(proc->Handles[i].Object);
            proc->Handles[i].Object = (OBJECT*)0;
            proc->Handles[i].Rights = 0;
            proc->Handles[i].Flags = 0;
        }
    }

    // Release every user VA region owned by this process. This includes the
    // main image, DLLs, TEBs/stacks that survived thread teardown, process
    // environment strings, and leaked user allocations.
    MmDestroyAllUserRegions(proc);
    proc->ImageBase = 0;
    proc->ImageSize = 0;
    proc->ModuleCount = 0;
    proc->UserPebBase = 0;
    proc->UserPebSize = 0;
    proc->UserImagePathBase = 0;
    proc->UserImagePathSize = 0;
    proc->UserCmdlineBase = 0;
    proc->UserCmdlineSize = 0;
    proc->UserCwdBase = 0;
    proc->UserCwdSize = 0;

    // Release the exit event.
    if (proc->ExitEvent) {
        ObDereferenceObject(&proc->ExitEvent->Header);
        proc->ExitEvent = (EVENT *)0;
    }

    // Remove from global list.
    uintptr_t flags = KeAcquireSpinlock(&gProcessListLock);
    ListRemove(&proc->GlobalLink);
    KeReleaseSpinlock(&gProcessListLock, flags);

    // Memory for the struct itself is freed by ObDereferenceObject after
    // Destroy returns.
}

// ---------------------------------------------------------------------------
// KeCreateProcess
// ---------------------------------------------------------------------------
KSTATUS KeCreateProcess(const char* name, PROCESS **out) {
    KiProcessInit();

    if (!name || !out) return STATUS_INVALID_PARAMETER;

    PROCESS* proc = (PROCESS *)ObAllocateObject(ObjectTypeProcess, sizeof(PROCESS));
    if (!proc) return STATUS_NO_MEMORY;

    proc->Header.Destroy = KiProcessDestroy;

    proc->ProcessId  = _InterlockedExchangeAdd((volatile long *)&gNextProcessId, 1);
    strncpy(proc->Name, name, PROCESS_NAME_MAX - 1);
    proc->Name[PROCESS_NAME_MAX - 1] = '\0';

    LIST_INIT(&proc->ThreadList);
    proc->ThreadCount = 0;
    proc->HandleLock  = (SPINLOCK)SPINLOCK_INIT;
    proc->ExitCode    = 0;
    proc->Terminating = 0;
    proc->Exited      = 0;
    proc->ImageBase   = 0;
    proc->ImageSize   = 0;
    proc->UserPebBase = 0;
    proc->UserPebSize = 0;
    proc->UserImagePathBase = 0;
    proc->UserImagePathSize = 0;
    proc->UserCmdlineBase = 0;
    proc->UserCmdlineSize = 0;
    proc->UserCwdBase = 0;
    proc->UserCwdSize = 0;
    proc->ModuleCount = 0;
    proc->CurrentDirectory[0] = '\0';
    proc->CurrentDirectoryLength = 0;
    proc->CurrentDrive = 0;
    proc->CwdLock = (SPINLOCK)SPINLOCK_INIT;
    LIST_INIT(&proc->UserRegions);
    proc->UserRegionLock = (SPINLOCK)SPINLOCK_INIT;

    // Slot 0 is the invalid handle - leave it NULL.
    memset(proc->Handles, 0, sizeof(proc->Handles));

    LIST_INIT(&proc->GlobalLink);

    // Create the manual-reset exit event (initially unsignaled).
    if (!KSUCCESS(KeEventCreate(EventManualReset, 0, &proc->ExitEvent))) {
        ObDereferenceObject(&proc->Header);
        return STATUS_NO_MEMORY;
    }

    uintptr_t flags = KeAcquireSpinlock(&gProcessListLock);
    ListInsertTail(&gProcessList, &proc->GlobalLink);
    KeReleaseSpinlock(&gProcessListLock, flags);

    *out = proc;
    return STATUS_SUCCESS;
}


// ---------------------------------------------------------------------------
// KeSetupUserProcessEnvironment
// ---------------------------------------------------------------------------
static KSTATUS KiCopyEnvironmentStringView(PROCESS* proc,
                                           SYS_STRING_VIEW src,
                                           uintptr_t* baseOut,
                                           size_t* sizeOut) {
    if (!baseOut || !sizeOut) return STATUS_INVALID_PARAMETER;
    *baseOut = 0;
    *sizeOut = 0;

    if (!SysSvIsValid(src)) return STATUS_INVALID_PARAMETER;
    if (SysSvContainsByte(src, '\0')) return STATUS_INVALID_PARAMETER;
    if (src.Length == (size_t)-1) return STATUS_BUFFER_TOO_SMALL;

    size_t len = src.Length + 1;
    char* dst = (char*)MmAllocateUserMemoryForProcess(proc, len);
    if (!dst) return STATUS_NO_MEMORY;

    if (src.Length) memcpy(dst, src.Data, src.Length);
    dst[src.Length] = '\0';

    *baseOut = (uintptr_t)dst;
    *sizeOut = len;
    return STATUS_SUCCESS;
}

static KSTATUS KiCopyEnvironmentStringViewFixed(PROCESS* proc,
                                                SYS_STRING_VIEW src,
                                                size_t capacity,
                                                uintptr_t* baseOut,
                                                size_t* sizeOut) {
    if (!baseOut || !sizeOut || capacity == 0) return STATUS_INVALID_PARAMETER;
    *baseOut = 0;
    *sizeOut = 0;

    if (!SysSvIsValid(src)) return STATUS_INVALID_PARAMETER;
    if (SysSvContainsByte(src, '\0')) return STATUS_INVALID_PARAMETER;
    if (src.Length + 1 > capacity) return STATUS_BUFFER_TOO_SMALL;

    char* dst = (char*)MmAllocateUserMemoryForProcess(proc, capacity);
    if (!dst) return STATUS_NO_MEMORY;
    memset(dst, 0, capacity);

    if (src.Length) memcpy(dst, src.Data, src.Length);
    dst[src.Length] = '\0';

    *baseOut = (uintptr_t)dst;
    *sizeOut = capacity;
    return STATUS_SUCCESS;
}

static void KiFreePartialProcessEnvironment(PROCESS* proc) {
    if (!proc) return;

    if (proc->UserPebBase && proc->UserPebSize) {
        MmFreeUserMemoryForProcess(proc, (void*)proc->UserPebBase);
        proc->UserPebBase = 0;
        proc->UserPebSize = 0;
    }
    if (proc->UserImagePathBase && proc->UserImagePathSize) {
        MmFreeUserMemoryForProcess(proc, (void*)proc->UserImagePathBase);
        proc->UserImagePathBase = 0;
        proc->UserImagePathSize = 0;
    }
    if (proc->UserCmdlineBase && proc->UserCmdlineSize) {
        MmFreeUserMemoryForProcess(proc, (void*)proc->UserCmdlineBase);
        proc->UserCmdlineBase = 0;
        proc->UserCmdlineSize = 0;
    }
    if (proc->UserCwdBase && proc->UserCwdSize) {
        MmFreeUserMemoryForProcess(proc, (void*)proc->UserCwdBase);
        proc->UserCwdBase = 0;
        proc->UserCwdSize = 0;
    }
}

static KSTATUS KiFormatNativeCwdV(const PROCESS* proc,
                                  SYS_STRING_VIEW currentDirectory,
                                  char* out,
                                  size_t outSize,
                                  size_t* lengthOut) {
    if (lengthOut) *lengthOut = 0;
    if (!out || outSize == 0) return STATUS_INVALID_PARAMETER;
    out[0] = '\0';

    if (!SysSvIsValid(currentDirectory)) return STATUS_INVALID_PARAMETER;
    if (SysSvContainsByte(currentDirectory, '\0')) return STATUS_INVALID_PARAMETER;

    uint8_t drive = proc ? proc->CurrentDrive : 0;
    if (drive > 7) drive = 0;

    if (currentDirectory.Length + 4 > outSize) return STATUS_BUFFER_TOO_SMALL;

    out[0] = (char)('0' + drive);
    out[1] = ':';
    out[2] = '/';
    if (currentDirectory.Length)
        memcpy(out + 3, currentDirectory.Data, currentDirectory.Length);
    out[3 + currentDirectory.Length] = '\0';

    if (lengthOut) *lengthOut = 3 + currentDirectory.Length;
    return STATUS_SUCCESS;
}

KSTATUS KeSetupUserProcessEnvironment(PROCESS* proc,
                                       SYS_STRING_VIEW imagePath,
                                       SYS_STRING_VIEW commandLine,
                                       SYS_STRING_VIEW currentDirectory) {
    if (!proc || imagePath.Length == 0) return STATUS_INVALID_PARAMETER;
    if (!SysSvIsValid(imagePath) || !SysSvIsValid(commandLine) || !SysSvIsValid(currentDirectory))
        return STATUS_INVALID_PARAMETER;
    if (SysSvContainsByte(imagePath, '\0') ||
        SysSvContainsByte(commandLine, '\0') ||
        SysSvContainsByte(currentDirectory, '\0'))
        return STATUS_INVALID_PARAMETER;

    SYS_STRING_VIEW finalCommandLine = commandLine.Length ? commandLine : imagePath;
    SYS_STRING_VIEW finalCwdRel = currentDirectory;
    if (finalCwdRel.Length == 0 && proc->CurrentDirectoryLength)
        finalCwdRel = SysSvFromParts(proc->CurrentDirectory, proc->CurrentDirectoryLength);

    char finalCwd[SYS_CWD_MAX];
    size_t finalCwdLength = 0;
    KSTATUS s = KiFormatNativeCwdV(proc, finalCwdRel, finalCwd, sizeof(finalCwd), &finalCwdLength);
    if (!KSUCCESS(s)) return s;

    KiFreePartialProcessEnvironment(proc);

    s = KiCopyEnvironmentStringView(proc, imagePath, &proc->UserImagePathBase, &proc->UserImagePathSize);
    if (!KSUCCESS(s)) goto fail;

    s = KiCopyEnvironmentStringView(proc, finalCommandLine, &proc->UserCmdlineBase, &proc->UserCmdlineSize);
    if (!KSUCCESS(s)) goto fail;

    s = KiCopyEnvironmentStringViewFixed(proc,
                                         SysSvFromParts(finalCwd, finalCwdLength),
                                         SYS_CWD_MAX,
                                         &proc->UserCwdBase,
                                         &proc->UserCwdSize);
    if (!KSUCCESS(s)) goto fail;

    SYS_PEB* peb = (SYS_PEB*)MmAllocateUserMemoryForProcess(proc, sizeof(SYS_PEB));
    if (!peb) { s = STATUS_NO_MEMORY; goto fail; }
    memset(peb, 0, sizeof(SYS_PEB));

    peb->Magic = SYS_PEB_MAGIC;
    peb->Version = SYS_ENV_VERSION;
    peb->ProcessId = proc->ProcessId;
    peb->ImagePath = SysSvFromParts((const char*)proc->UserImagePathBase, imagePath.Length);
    peb->CommandLine = SysSvFromParts((const char*)proc->UserCmdlineBase, finalCommandLine.Length);
    peb->CurrentDirectory = SysSvFromParts((const char*)proc->UserCwdBase, finalCwdLength);

    proc->UserPebBase = (uintptr_t)peb;
    proc->UserPebSize = sizeof(SYS_PEB);
    return STATUS_SUCCESS;

fail:
    KiFreePartialProcessEnvironment(proc);
    return s;
}

// ---------------------------------------------------------------------------
// Process exit / termination
// ---------------------------------------------------------------------------
static void KiCompleteProcessExitIfLastThread(PROCESS* process) {
    if (!process) return;

    uintptr_t flags = KeAcquireSpinlock(&process->HandleLock);
    uint32_t threadCount = process->ThreadCount;
    KeReleaseSpinlock(&process->HandleLock, flags);

    if (threadCount == 0 && _InterlockedExchange(&process->Exited, 1) == 0) {
        if (process->ExitEvent) KeSignalEvent(process->ExitEvent);
    }
}

void KeExitProcess(PROCESS* process, int exitCode) {
    if (!process) return;
    process->ExitCode = exitCode;
    _InterlockedExchange(&process->Terminating, 1);

    // The calling thread exits immediately after this in SysExitProcess. Do
    // not drop the process object here: each thread owns one process reference,
    // and KiThreadDestroy releases that reference after unlinking the thread.
    KiCompleteProcessExitIfLastThread(process);
}

KSTATUS KeTerminateProcess(PROCESS* process, int exitCode) {
    if (!process) return STATUS_INVALID_PARAMETER;

    process->ExitCode = exitCode;
    _InterlockedExchange(&process->Terminating, 1);

    // Wake alertable waiters so they can notice Terminating at the syscall/
    // wait boundary. This is intentionally conservative: fully asynchronous
    // forced thread teardown needs object-specific wait-list locking and is a
    // larger process-lifecycle project.
    uintptr_t flags = KeAcquireSpinlock(&process->HandleLock);
    LIST_ENTRY* head = &process->ThreadList;
    for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
        THREAD* t = CONTAINING_RECORD(node, THREAD, ProcessLink);
        if (t != KeGetCurrentThread())
            (void)KeAlertThread(t);
    }
    KeReleaseSpinlock(&process->HandleLock, flags);

    KiCompleteProcessExitIfLastThread(process);
    return STATUS_SUCCESS;
}

int KeIsProcessTerminating(PROCESS* process) {
    return process && process->Terminating != 0;
}

// ---------------------------------------------------------------------------
// KeGetProcessExitCode
// ---------------------------------------------------------------------------
KSTATUS KeGetProcessExitCode(PROCESS* process, int* exitCodeOut) {
    if (!process || !exitCodeOut) return STATUS_INVALID_PARAMETER;
    *exitCodeOut = process->ExitCode;
    return STATUS_SUCCESS;
}

// ---------------------------------------------------------------------------
// Handle table
// ---------------------------------------------------------------------------
KSTATUS KeAllocateHandleWithRights(PROCESS* process, OBJECT* object, uint32_t rights, uint32_t* handleOut) {
    if (!process || !object || !handleOut || rights == 0) return STATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&process->HandleLock);

    for (int i = 1; i < HANDLE_TABLE_SIZE; i++) {
        if (!process->Handles[i].Object) {
            ObReferenceObject(object);
            process->Handles[i].Object = object;
            process->Handles[i].Rights = rights;
            process->Handles[i].Flags = 0;
            *handleOut = (uint32_t)i;
            KeReleaseSpinlock(&process->HandleLock, flags);
            return STATUS_SUCCESS;
        }
    }

    KeReleaseSpinlock(&process->HandleLock, flags);
    return STATUS_NO_MEMORY;
}

KSTATUS KeAllocateHandle(PROCESS* process, OBJECT* object, uint32_t* handleOut) {
    return KeAllocateHandleWithRights(process, object, HANDLE_RIGHT_ALL, handleOut);
}

void KeCloseHandle(PROCESS* process, uint32_t handle) {
    if (!process || handle == 0 || handle >= HANDLE_TABLE_SIZE) return;

    uintptr_t flags = KeAcquireSpinlock(&process->HandleLock);
    OBJECT* obj = process->Handles[handle].Object;
    if (obj) {
        process->Handles[handle].Object = (OBJECT*)0;
        process->Handles[handle].Rights = 0;
        process->Handles[handle].Flags = 0;
        KeReleaseSpinlock(&process->HandleLock, flags);
        ObDereferenceObject(obj);
        return;
    }
    KeReleaseSpinlock(&process->HandleLock, flags);
}

OBJECT* KeLookupHandleWithRights(PROCESS* process, uint32_t handle, uint32_t requiredRights) {
    if (!process || handle == 0 || handle >= HANDLE_TABLE_SIZE)
        return (OBJECT*)0;

    uintptr_t flags = KeAcquireSpinlock(&process->HandleLock);
    HANDLE_ENTRY* entry = &process->Handles[handle];
    OBJECT* obj = entry->Object;
    if (obj && requiredRights && ((entry->Rights & requiredRights) != requiredRights))
        obj = (OBJECT*)0;
    if (obj) ObReferenceObject(obj);
    KeReleaseSpinlock(&process->HandleLock, flags);
    return obj;
}

OBJECT* KeLookupHandle(PROCESS* process, uint32_t handle) {
    return KeLookupHandleWithRights(process, handle, 0);
}

uint32_t KeGetHandleRights(PROCESS* process, uint32_t handle) {
    if (!process || handle == 0 || handle >= HANDLE_TABLE_SIZE) return 0;

    uintptr_t flags = KeAcquireSpinlock(&process->HandleLock);
    uint32_t rights = process->Handles[handle].Object ? process->Handles[handle].Rights : 0;
    KeReleaseSpinlock(&process->HandleLock, flags);
    return rights;
}

PROCESS* KeGetProcessById(uint32_t pid) {
    KiProcessInit();
    uintptr_t flags = KeAcquireSpinlock(&gProcessListLock);
    LIST_ENTRY* head = &gProcessList;
    LIST_ENTRY* node = head->Fwd;
    while (node != head) {
        PROCESS* p = CONTAINING_RECORD(node, PROCESS, GlobalLink);
        if (p->ProcessId == pid) {
            ObReferenceObject(&p->Header);
            KeReleaseSpinlock(&gProcessListLock, flags);
            return p;
        }
        node = node->Fwd;
    }
    KeReleaseSpinlock(&gProcessListLock, flags);
    return (PROCESS*)0;
}