// Native resource service bootstrap and lifetime. RSSRV is one critical server with internally owned components.

#include "rssrv_internal.h"
#include <ntlite/exception.h>
#include <ntlite/service.h>
#include <ntlite/service_control.h>
#include <ntlite/syscall.h>
#include <stdint.h>
#include <string.h>

static RSSRV_STATE gRssrv;

// Sets a best-effort diagnostic name for the current RSSRV thread.
void RssrvSetCurrentThreadName(const char* name) {
    SYS_SET_THREAD_NAME_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.ThreadHandle = SYS_HANDLE_INVALID;
    request.Name = SysSvFromCstr(name);
    (void)SysSetThreadName(&request);
}

// Terminates RSSRV without depending on a compiler runtime return path.
void RssrvExit(int code) {
    for (;;) SysExitProcess(code);
}

// Accepts the one private service-manager control connection.
static KSTATUS RssrvAcceptControl(SYS_HANDLE controlListener, SYS_HANDLE* endpointOut) {
    SYS_ACCEPT_PORT_REQUEST accept;
    memset(&accept, 0, sizeof(accept));
    accept.Size = sizeof(accept);
    accept.ListenerHandle = controlListener;
    accept.Flags = SYS_PORT_ACCEPT_FLAG_NONE;
    accept.TimeoutMilliseconds = SYS_PORT_TIMEOUT_INFINITE;
    accept.EndpointHandleOut = endpointOut;
    return SysAcceptPort(&accept);
}

// Sends service readiness and delegates init-only namespace and DWM-session bootstrap capabilities.
static KSTATUS RssrvSendReady(SYS_HANDLE controlEndpoint,
                              const RSSRV_RESOURCE_SERVICE* resource,
                              const RSSRV_SESSION_STATE* session) {
    if (!resource || !session || resource->AdminListener == SYS_HANDLE_INVALID ||
        session->DwmListener == SYS_HANDLE_INVALID)
        return NTLSTATUS_INVALID_PARAMETER;

    NTL_SERVICE_CONTROL_MESSAGE message;
    memset(&message, 0, sizeof(message));
    message.Size = sizeof(message);
    message.Version = NTL_SERVICE_CONTROL_VERSION;
    message.Operation = NTL_SERVICE_CONTROL_READY;
    message.ExportCount = 2u;
    message.Status = NTLSTATUS_SUCCESS;

    SYS_PORT_HANDLE_TRANSFER transfers[2];
    memset(transfers, 0, sizeof(transfers));
    transfers[0].Handle = resource->AdminListener;
    transfers[0].DesiredRights = SYS_RIGHT_QUERY | SYS_PORT_RIGHT_CONNECT;
    transfers[0].Flags = SYS_PORT_TRANSFER_FLAG_NONE;
    transfers[1].Handle = session->DwmListener;
    // Init must transfer this private listener once more into DWM. DWM itself receives only connect/query rights.
    transfers[1].DesiredRights = SYS_RIGHT_DUPLICATE | SYS_RIGHT_TRANSFER | SYS_RIGHT_QUERY | SYS_PORT_RIGHT_CONNECT;
    transfers[1].Flags = SYS_PORT_TRANSFER_FLAG_NONE;

    SYS_SEND_PORT_REQUEST send;
    memset(&send, 0, sizeof(send));
    send.Size = sizeof(send);
    send.EndpointHandle = controlEndpoint;
    send.Flags = SYS_PORT_FLAG_NONE;
    send.Payload = &message;
    send.PayloadLength = sizeof(message);
    send.Handles = transfers;
    send.HandleCount = 2u;
    send.TimeoutMilliseconds = SYS_PORT_TIMEOUT_INFINITE;
    return SysSendPort(&send);
}

// Waits for the service manager's stop notification or control-port disconnect.
static void RssrvControlLoop(SYS_HANDLE controlEndpoint) {
    for (;;) {
        NTL_SERVICE_CONTROL_MESSAGE message;
        SYS_PORT_MESSAGE_INFO info;
        memset(&message, 0, sizeof(message));
        memset(&info, 0, sizeof(info));
        info.Size = sizeof(info);
        SYS_RECEIVE_PORT_REQUEST receive;
        memset(&receive, 0, sizeof(receive));
        receive.Size = sizeof(receive);
        receive.EndpointHandle = controlEndpoint;
        receive.Payload = &message;
        receive.PayloadCapacity = sizeof(message);
        receive.InfoOut = &info;
        receive.TimeoutMilliseconds = SYS_PORT_TIMEOUT_INFINITE;
        KSTATUS status = SysReceivePort(&receive);
        if (status == NTLSTATUS_PORT_DISCONNECTED)
            return;
        if (!KSUCCESS(status))
            return;
        if (info.Type != SYS_PORT_MESSAGE_NOTIFICATION || info.HandleCount != 0 ||
            info.PayloadLength != sizeof(message) || message.Size != sizeof(message) ||
            message.Version != NTL_SERVICE_CONTROL_VERSION || message.ExportCount != 0 ||
            message.Reserved[0] != 0 || message.Reserved[1] != 0 || message.Reserved[2] != 0)
            continue;
        if (message.Operation == NTL_SERVICE_CONTROL_STOP)
            return;
    }
}

// Queries the mounted-volume identity carried by one volume or directory root capability.
static KSTATUS RssrvQueryRootVolumeId(SYS_HANDLE root, uint32_t* volumeIdOut) {
    if (root == SYS_HANDLE_INVALID || !volumeIdOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *volumeIdOut = 0;
    SYS_FILE_STAT stat;
    memset(&stat, 0, sizeof(stat));
    SYS_STAT_AT_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.RootHandle = root;
    request.Path = SYS_SV_NULL;
    request.StatOut = &stat;
    KSTATUS status = SysStatAt(&request);
    if (!KSUCCESS(status)) return status;
    if (stat.Size != sizeof(stat) || stat.VolumeId == 0 || stat.VolumeId > UINT32_MAX)
        return NTLSTATUS_FAILURE;
    *volumeIdOut = (uint32_t)stat.VolumeId;
    return NTLSTATUS_SUCCESS;
}

// Initializes the fixed native namespace roots before any dynamic publication starts.
static KSTATUS RssrvInitializeNamespaceRoots(RSSRV_NAMESPACE* namespaceState) {
    static const char* roots[] = {
        "/Device",
        "/Sessions",
        "/Services",
    };
    for (uint32_t i = 0; i < (uint32_t)(sizeof(roots) / sizeof(roots[0])); ++i) {
        KSTATUS status = RssrvNamespaceCreateDirectory(namespaceState,
                                                       roots[i],
                                                       (uint32_t)strlen(roots[i]));
        if (!KSUCCESS(status)) return status;
    }
    return NTLSTATUS_SUCCESS;
}

// Starts RSSRV and owns the global native resource namespace until init requests shutdown.
void _start(void) {
    KSTATUS status = NtlInitializeExceptionRuntime();
    if (!KSUCCESS(status)) RssrvExit(1);
    RssrvSetCurrentThreadName("rssrv.main");

    SYS_SECURITY_CONTEXT_INFO security;
    memset(&security, 0, sizeof(security));
    security.Size = sizeof(security);
    status = SysQuerySecurityContext(SYS_SECURITY_CONTEXT_CURRENT, &security);
    if (!KSUCCESS(status) || security.EffectivePrivileges != SYS_PRIV_SYSTEM_SERVICE_REGISTER ||
        security.DelegablePrivileges != 0 || security.ServiceClaims != SYS_SERVICE_CLAIM_RESOURCE ||
        security.DelegableServiceClaims != 0)
        RssrvExit(2);

    SYS_HANDLE systemVolume = SysGetStartupCapability(SYS_STARTUP_CAP_SYSTEM_VOLUME_NAMESPACE);
    SYS_HANDLE controlListener = SysGetStartupCapability(SYS_STARTUP_CAP_SERVICE_CONTROL);
    SYS_HANDLE volumeManager = SysGetStartupCapability(SYS_STARTUP_CAP_VOLUME_MANAGER);
    SYS_HANDLE inputManager = SysGetStartupCapability(SYS_STARTUP_CAP_INPUT_MANAGER);
    SYS_HANDLE displayManager = SysGetStartupCapability(SYS_STARTUP_CAP_DISPLAY_MANAGER);
    if (systemVolume == SYS_HANDLE_INVALID || controlListener == SYS_HANDLE_INVALID ||
        volumeManager == SYS_HANDLE_INVALID || inputManager == SYS_HANDLE_INVALID ||
        displayManager == SYS_HANDLE_INVALID)
        RssrvExit(3);

    uint32_t systemVolumeId = 0;
    status = RssrvQueryRootVolumeId(systemVolume, &systemVolumeId);
    if (!KSUCCESS(status)) RssrvExit(3);

    memset(&gRssrv, 0, sizeof(gRssrv));
    status = RssrvNamespaceInitialize(&gRssrv.Namespace);
    if (!KSUCCESS(status)) RssrvExit(3);
    status = RssrvInitializeNamespaceRoots(&gRssrv.Namespace);
    if (!KSUCCESS(status)) RssrvExit(4);

    status = RssrvVolumesInitialize(&gRssrv.Volumes, &gRssrv.Namespace, volumeManager);
    if (!KSUCCESS(status)) RssrvExit(4);
    status = RssrvDisplayInitialize(&gRssrv.Display, &gRssrv.Namespace, displayManager);
    if (!KSUCCESS(status)) RssrvExit(4);
    status = RssrvDisplayPublishDevice(&gRssrv.Display);
    if (!KSUCCESS(status)) RssrvExit(4);
    status = RssrvSessionInitialize(&gRssrv.Session, &gRssrv.Namespace, &gRssrv.Display);
    if (!KSUCCESS(status)) RssrvExit(4);
    status = RssrvInputInitialize(&gRssrv.Input, &gRssrv.Namespace, &gRssrv.Session, inputManager);
    if (!KSUCCESS(status)) RssrvExit(4);
    status = RssrvConsoleInitialize(&gRssrv.Console, &gRssrv.Namespace,
                                    &gRssrv.Display, &gRssrv.Session);
    if (!KSUCCESS(status)) RssrvExit(4);
    uint64_t volumeGeneration = 0;
    status = RssrvVolumesSynchronize(&gRssrv.Volumes, &volumeGeneration);
    if (!KSUCCESS(status)) RssrvExit(4);

    char systemVolumePath[RSSRV_VOLUME_PATH_CAPACITY];
    uint32_t systemVolumePathLength = 0;
    status = RssrvVolumesQueryPublishedPath(&gRssrv.Volumes,
                                            systemVolumeId,
                                            systemVolumePath,
                                            (uint32_t)sizeof(systemVolumePath),
                                            &systemVolumePathLength);
    if (!KSUCCESS(status)) RssrvExit(4);

    // /SystemRoot is a real namespace symbolic link to the independently published system volume.
    status = RssrvNamespaceCreateSymbolicLink(&gRssrv.Namespace,
                                              "/SystemRoot",
                                              sizeof("/SystemRoot") - 1u,
                                              systemVolumePath,
                                              systemVolumePathLength);
    if (!KSUCCESS(status)) RssrvExit(4);
    (void)SysCloseHandle(systemVolume);

    SYS_HANDLE controlEndpoint = SYS_HANDLE_INVALID;
    status = RssrvAcceptControl(controlListener, &controlEndpoint);
    (void)SysCloseHandle(controlListener);
    if (!KSUCCESS(status)) RssrvExit(5);

    status = RssrvResourceInitialize(&gRssrv.Resource, &gRssrv.Namespace);
    if (!KSUCCESS(status)) RssrvExit(6);
    status = RssrvResourceStart(&gRssrv.Resource);
    if (!KSUCCESS(status)) RssrvExit(7);
    status = RssrvInputStart(&gRssrv.Input);
    if (!KSUCCESS(status)) RssrvExit(8);
    status = RssrvConsoleStart(&gRssrv.Console);
    if (!KSUCCESS(status)) RssrvExit(9);
    status = RssrvSessionStart(&gRssrv.Session, &gRssrv.Console);
    if (!KSUCCESS(status)) RssrvExit(10);
    status = RssrvVolumesStartWorker(&gRssrv.Volumes);
    if (!KSUCCESS(status)) RssrvExit(11);

    status = SysRegisterSystemService(SYS_SERVICE_RESOURCE, gRssrv.Resource.PublicListener);
    if (!KSUCCESS(status)) RssrvExit(12);
    status = RssrvSendReady(controlEndpoint, &gRssrv.Resource, &gRssrv.Session);
    if (!KSUCCESS(status)) RssrvExit(13);

    RssrvControlLoop(controlEndpoint);
    (void)SysUnregisterSystemService(SYS_SERVICE_RESOURCE, gRssrv.Resource.PublicListener);
    (void)SysDisconnectPort(controlEndpoint);
    (void)SysCloseHandle(controlEndpoint);
    RssrvExit(0);
}