// Native userspace process construction built on capability-relative file open and PE materialization.

#include <ntlite/process_loader.h>
#include <ntlite/path.h>
#include <ntlite/resource.h>
#include <ntlite/syscall.h>
#include <string.h>
#include "module_loader_internal.h"

// Terminates one failed constructing child and releases the parent process handle.
static void NtlDiscardChild(SYS_HANDLE processHandle) {
    if (processHandle == SYS_HANDLE_INVALID)
        return;
    (void)SysKillProcess(processHandle, (int32_t)NTLSTATUS_PROCESS_KILLED);
    (void)SysCloseHandle(processHandle);
}

// Creates and starts one child from an already-open PE file capability.
static KSTATUS NtlCreateProcessFile(const NTL_CREATE_PROCESS_AT_REQUEST* request, SYS_HANDLE fileHandle) {
    if (!request || fileHandle == SYS_HANDLE_INVALID)
        return NTLSTATUS_INVALID_PARAMETER;

    SYS_HANDLE processHandle = SYS_HANDLE_INVALID;
    SYS_CREATE_PROCESS_OBJECT_REQUEST create;
    memset(&create, 0, sizeof(create));
    create.Size = sizeof(create);
    create.Flags = SYS_CREATE_PROCESS_OBJECT_FLAG_NONE;
    create.ImagePath = request->Path;
    create.CommandLine = request->CommandLine;
    create.SecurityContext = request->SecurityContext;
    create.ProcessHandleOut = &processHandle;
    KSTATUS status = SysCreateProcessObject(&create);
    if (!KSUCCESS(status))
        return status;

    NTL_MODULE_CONTEXT* modules = 0;
    status = NtlpCreateModuleContext(processHandle,
                                     request->ModuleRootHandle,
                                     request->ModuleDirectory,
                                     request->ResolveModuleName,
                                     request->ResolveContext,
                                     0,
                                     &modules);
    if (!KSUCCESS(status)) {
        NtlDiscardChild(processHandle);
        return status;
    }

    NTL_PE_IMAGE image;
    memset(&image, 0, sizeof(image));
    status = NtlpLoadMainImage(modules, fileHandle, &image);
    if (!KSUCCESS(status)) {
        NtlpDestroyModuleContext(modules);
        NtlDiscardChild(processHandle);
        return status;
    }

    SYS_STARTUP_CAPABILITY childCapabilities[SYS_STARTUP_CAPABILITY_MAX];
    memset(childCapabilities, 0, sizeof(childCapabilities));
    for (uint32_t i = 0; i < request->StartupCapabilityCount; ++i) {
        const NTL_PROCESS_STARTUP_CAPABILITY* source = &request->StartupCapabilities[i];
        if (!source->Tag || source->SourceHandle == SYS_HANDLE_INVALID) {
            NtlpDestroyModuleContext(modules);
            NtlDiscardChild(processHandle);
            return NTLSTATUS_INVALID_PARAMETER;
        }
        for (uint32_t j = 0; j < i; ++j) {
            if (childCapabilities[j].Tag == source->Tag) {
                NtlpDestroyModuleContext(modules);
                NtlDiscardChild(processHandle);
                return NTLSTATUS_INVALID_PARAMETER;
            }
        }

        SYS_HANDLE childHandle = SYS_HANDLE_INVALID;
        SYS_DUPLICATE_HANDLE_REQUEST duplicate;
        memset(&duplicate, 0, sizeof(duplicate));
        duplicate.Size = sizeof(duplicate);
        duplicate.SourceHandle = source->SourceHandle;
        duplicate.TargetProcessHandle = processHandle;
        duplicate.DesiredRights = source->DesiredRights;
        duplicate.Flags = source->DuplicateFlags;
        duplicate.TargetHandleOut = &childHandle;
        status = SysDuplicateHandle(&duplicate);
        if (!KSUCCESS(status)) {
            NtlpDestroyModuleContext(modules);
            NtlDiscardChild(processHandle);
            return status;
        }
        childCapabilities[i].Tag = source->Tag;
        childCapabilities[i].Handle = childHandle;
    }

    char currentDirectoryBuffer[SYS_PATH_MAX + 1u];
    SYS_STRING_VIEW currentDirectory;
    if (request->CurrentDirectory.Length != 0) {
        size_t required = 0;
        status = NtlRealPath(request->CurrentDirectory, currentDirectoryBuffer,
                             sizeof(currentDirectoryBuffer), &required);
        if (!KSUCCESS(status)) {
            NtlpDestroyModuleContext(modules);
            NtlDiscardChild(processHandle);
            return status;
        }
        currentDirectory = SysSvFromParts(currentDirectoryBuffer, required - 1u);

        SYS_HANDLE cwdHandle = SYS_HANDLE_INVALID;
        status = NtlOpenDir(currentDirectory, SYS_DIR_ACCESS_ENUMERATE | SYS_DIR_ACCESS_OPEN,
                            SYS_DIR_FLAG_NONE, &cwdHandle);
        if (!KSUCCESS(status)) {
            NtlpDestroyModuleContext(modules);
            NtlDiscardChild(processHandle);
            return status;
        }
        (void)SysCloseHandle(cwdHandle);
    } else {
        size_t required = 0;
        status = NtlGetCurrentDirectory(currentDirectoryBuffer, sizeof(currentDirectoryBuffer), &required);
        if (!KSUCCESS(status)) {
            NtlpDestroyModuleContext(modules);
            NtlDiscardChild(processHandle);
            return status;
        }
        currentDirectory = SysSvFromParts(currentDirectoryBuffer, required - 1u);
    }

    void* loaderData = 0;
    status = NtlpBuildLoaderData(modules, currentDirectory, request->StartupData,
                                  request->StartupDataSize, &loaderData);
    if (!KSUCCESS(status)) {
        NtlpDestroyModuleContext(modules);
        NtlDiscardChild(processHandle);
        return status;
    }

    SYS_START_PROCESS_REQUEST startRequest;
    memset(&startRequest, 0, sizeof(startRequest));
    startRequest.Size = sizeof(startRequest);
    startRequest.ProcessHandle = processHandle;
    startRequest.Flags = SYS_START_PROCESS_FLAG_NONE;
    startRequest.MainImageBase = image.Base;
    startRequest.MainImageSize = image.Size;
    startRequest.EntryPoint = image.EntryPoint;
    startRequest.StartupCapabilities = childCapabilities;
    startRequest.StartupCapabilityCount = request->StartupCapabilityCount;
    startRequest.LoaderData = loaderData;
    status = SysStartProcess(&startRequest);
    NtlpDestroyModuleContext(modules);
    if (!KSUCCESS(status)) {
        NtlDiscardChild(processHandle);
        return status;
    }

    *request->ProcessHandleOut = processHandle;
    return NTLSTATUS_SUCCESS;
}

// Creates and starts one child from a PE file beneath a capability-relative root.
KSTATUS NtlCreateProcessAt(const NTL_CREATE_PROCESS_AT_REQUEST* request) {
    if (!request || request->Size != sizeof(*request) ||
        request->Flags != NTL_PROCESS_LOAD_FLAG_NONE ||
        request->RootHandle == SYS_HANDLE_INVALID ||
        request->ModuleRootHandle == SYS_HANDLE_INVALID ||
        !request->Path.Data || request->Path.Length == 0 ||
        !SysSvIsValid(request->ModuleDirectory) ||
        request->StartupCapabilityCount > SYS_STARTUP_CAPABILITY_MAX ||
        (request->StartupCapabilityCount && !request->StartupCapabilities) ||
        request->StartupDataSize > NTL_PROCESS_STARTUP_DATA_MAX ||
        (request->StartupDataSize && !request->StartupData) ||
        !request->ProcessHandleOut)
        return NTLSTATUS_INVALID_PARAMETER;

    *request->ProcessHandleOut = SYS_HANDLE_INVALID;

    SYS_HANDLE fileHandle = SYS_HANDLE_INVALID;
    SYS_OPEN_FILE_AT_REQUEST open;
    memset(&open, 0, sizeof(open));
    open.Size = sizeof(open);
    open.RootHandle = request->RootHandle;
    open.Path = request->Path;
    open.Access = SYS_FILE_ACCESS_READ | SYS_FILE_ACCESS_QUERY;
    open.Disposition = SYS_FILE_OPEN_EXISTING;
    open.Flags = SYS_FILE_FLAG_NONE;
    open.HandleOut = &fileHandle;
    KSTATUS status = SysOpenFileAt(&open);
    if (!KSUCCESS(status))
        return status;

    status = NtlCreateProcessFile(request, fileHandle);
    (void)SysCloseHandle(fileHandle);
    return status;
}

KSTATUS NtlCreateSystemProcess(const NTL_CREATE_SYSTEM_PROCESS_REQUEST* request) {
    if (!request || request->Size != sizeof(*request) ||
        request->Flags != NTL_PROCESS_LOAD_FLAG_NONE || !request->Path.Data ||
        request->Path.Length == 0 || request->StartupDataSize > NTL_PROCESS_STARTUP_DATA_MAX ||
        (request->StartupDataSize && !request->StartupData) || !request->ProcessHandleOut)
        return NTLSTATUS_INVALID_PARAMETER;

    *request->ProcessHandleOut = SYS_HANDLE_INVALID;

    SYS_HANDLE systemRoot = SYS_HANDLE_INVALID;
    KSTATUS status = NtlResourceOpenSystemRoot(SYS_RIGHT_QUERY | SYS_VOLUME_RIGHT_OPEN, &systemRoot);
    if (!KSUCCESS(status)) return status;

    NTL_CREATE_PROCESS_AT_REQUEST create;
    memset(&create, 0, sizeof(create));
    create.Size = sizeof(create);
    create.Flags = request->Flags;
    create.RootHandle = systemRoot;
    create.Path = request->Path;
    create.CommandLine = request->CommandLine;
    create.CurrentDirectory = request->CurrentDirectory;
    create.SecurityContext = request->SecurityContext;
    create.ModuleRootHandle = systemRoot;
    create.ModuleDirectory = SysSvFromCstr("system");
    create.ResolveModuleName = request->ResolveModuleName;
    create.ResolveContext = request->ResolveContext;
    create.StartupCapabilities = request->StartupCapabilities;
    create.StartupCapabilityCount = request->StartupCapabilityCount;
    create.StartupData = request->StartupData;
    create.StartupDataSize = request->StartupDataSize;
    create.ProcessHandleOut = request->ProcessHandleOut;
    if (request->ImageFile != SYS_HANDLE_INVALID) {
        status = NtlCreateProcessFile(&create, request->ImageFile);
    } else {
        SYS_HANDLE imageFile = SYS_HANDLE_INVALID;
        status = NtlOpenFile(request->Path, SYS_FILE_ACCESS_READ | SYS_FILE_ACCESS_QUERY,
                             SYS_FILE_OPEN_EXISTING, SYS_FILE_FLAG_NONE, &imageFile);
        if (KSUCCESS(status)) {
            status = NtlCreateProcessFile(&create, imageFile);
            (void)SysCloseHandle(imageFile);
        }
    }
    (void)SysCloseHandle(systemRoot);
    return status;
}