// RSSRV mounted-volume publication and /Device/HarddiskVolumeN synchronization.

#include "rssrv_internal.h"
#include <stdlib.h>
#include <string.h>

struct RSSRV_VOLUME_PUBLICATION {
    RSSRV_VOLUME_PUBLICATION* Next;
    uint32_t InstanceId;
    uint32_t NamespaceId;
    uint32_t PathLength;
    char Path[RSSRV_VOLUME_PATH_CAPACITY];
};

// Finds one published mounted volume by kernel instance identity.
static RSSRV_VOLUME_PUBLICATION* RssrvFindPublishedVolume(RSSRV_VOLUME_STATE* volumes, uint32_t instanceId) {
    if (!volumes) return (RSSRV_VOLUME_PUBLICATION*)0;
    for (RSSRV_VOLUME_PUBLICATION* publication = volumes->Published; publication; publication = publication->Next) {
        if (publication->InstanceId == instanceId)
            return publication;
    }
    return (RSSRV_VOLUME_PUBLICATION*)0;
}

// Returns whether one kernel volume instance is present in the current manager snapshot.
static int RssrvSnapshotContainsVolume(const SYS_VOLUME_INFO* snapshot, uint32_t count, uint32_t instanceId) {
    for (uint32_t i = 0; i < count; ++i) {
        if (snapshot[i].InstanceId == instanceId)
            return 1;
    }
    return 0;
}

// Builds one RSSRV-local /Device/HarddiskVolumeN name without introducing a formatting-runtime dependency.
static KSTATUS RssrvBuildHarddiskVolumePath(uint32_t namespaceId,
                                            char* path,
                                            uint32_t capacity,
                                            uint32_t* lengthOut) {
    static const char prefix[] = "/Device/HarddiskVolume";
    if (!namespaceId || !path || !lengthOut)
        return NTLSTATUS_INVALID_PARAMETER;

    uint32_t prefixLength = (uint32_t)(sizeof(prefix) - 1u);
    char digits[10];
    uint32_t digitCount = 0;
    uint32_t value = namespaceId;
    do {
        digits[digitCount++] = (char)('0' + (value % 10u));
        value /= 10u;
    } while (value && digitCount < (uint32_t)sizeof(digits));

    if (prefixLength + digitCount + 1u > capacity)
        return NTLSTATUS_BUFFER_TOO_SMALL;
    memcpy(path, prefix, prefixLength);
    for (uint32_t i = 0; i < digitCount; ++i)
        path[prefixLength + i] = digits[digitCount - 1u - i];
    path[prefixLength + digitCount] = '\0';
    *lengthOut = prefixLength + digitCount;
    return NTLSTATUS_SUCCESS;
}

// Takes a stable manager snapshot, growing the userspace buffer if the mounted-volume count changes.
static KSTATUS RssrvQueryVolumes(SYS_HANDLE manager,
                                 SYS_VOLUME_INFO** volumesOut,
                                 uint32_t* countOut,
                                 uint64_t* generationOut) {
    if (!volumesOut || !countOut || !generationOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *volumesOut = (SYS_VOLUME_INFO*)0;
    *countOut = 0;
    *generationOut = 0;

    SYS_VOLUME_INFO* snapshot = (SYS_VOLUME_INFO*)0;
    uint32_t capacity = 0;
    for (;;) {
        uint32_t count = 0;
        uint64_t generation = 0;
        SYS_QUERY_VOLUMES_REQUEST query;
        memset(&query, 0, sizeof(query));
        query.Size = sizeof(query);
        query.ManagerHandle = manager;
        query.Flags = SYS_VOLUME_QUERY_FLAG_NONE;
        query.Capacity = capacity;
        query.Volumes = snapshot;
        query.CountOut = &count;
        query.GenerationOut = &generation;
        KSTATUS status = SysQueryVolumes(&query);
        if (status == NTLSTATUS_BUFFER_TOO_SMALL) {
            if (snapshot) free(snapshot);
            snapshot = (SYS_VOLUME_INFO*)0;
            capacity = count;
            if (capacity) {
                snapshot = (SYS_VOLUME_INFO*)malloc((size_t)capacity * sizeof(*snapshot));
                if (!snapshot)
                    return NTLSTATUS_NO_MEMORY;
                memset(snapshot, 0, (size_t)capacity * sizeof(*snapshot));
            }
            continue;
        }
        if (!KSUCCESS(status)) {
            if (snapshot) free(snapshot);
            return status;
        }
        for (uint32_t i = 0; i < count; ++i) {
            if (snapshot[i].Size != sizeof(SYS_VOLUME_INFO) ||
                snapshot[i].Version != SYS_VOLUME_ENUM_VERSION || !snapshot[i].InstanceId ||
                snapshot[i].Flags != 0) {
                if (snapshot) free(snapshot);
                return NTLSTATUS_INVALID_PARAMETER;
            }
        }
        *volumesOut = snapshot;
        *countOut = count;
        *generationOut = generation;
        return NTLSTATUS_SUCCESS;
    }
}

// Keeps automatic /Device/HarddiskVolumeN bindings synchronized with mounted IO_VOLUME instances.
static void RssrvVolumeWorker(uintptr_t argument) {
    RSSRV_VOLUME_STATE* volumes = (RSSRV_VOLUME_STATE*)argument;
    if (!volumes) {
        SysExitThread(1);
        return;
    }
    RssrvSetCurrentThreadName("rssrv.volume");
    for (;;) {
        uint64_t generation = 0;
        KSTATUS status = RssrvVolumesSynchronize(volumes, &generation);
        if (!KSUCCESS(status))
            RssrvExit(12);
        status = SysWaitVolumeChange(volumes->Manager, generation, SYS_WAIT_TIMEOUT_INFINITE);
        if (!KSUCCESS(status))
            RssrvExit(12);
    }
}

// Initializes mounted-volume publication state against one volume-manager capability.
KSTATUS RssrvVolumesInitialize(RSSRV_VOLUME_STATE* volumes,
                               RSSRV_NAMESPACE* namespaceState,
                               SYS_HANDLE manager) {
    if (!volumes || !namespaceState || manager == SYS_HANDLE_INVALID)
        return NTLSTATUS_INVALID_PARAMETER;
    memset(volumes, 0, sizeof(*volumes));
    volumes->Namespace = namespaceState;
    volumes->Manager = manager;
    volumes->NextHarddiskVolumeId = 1u;
    return NTLSTATUS_SUCCESS;
}

// Reconciles RSSRV volume bindings with the kernel's mounted-volume snapshot.
KSTATUS RssrvVolumesSynchronize(RSSRV_VOLUME_STATE* volumes, uint64_t* generationOut) {
    if (!volumes || !volumes->Namespace || volumes->Manager == SYS_HANDLE_INVALID || !generationOut)
        return NTLSTATUS_INVALID_PARAMETER;

    SYS_VOLUME_INFO* snapshot = (SYS_VOLUME_INFO*)0;
    uint32_t count = 0;
    uint64_t generation = 0;
    KSTATUS status = RssrvQueryVolumes(volumes->Manager, &snapshot, &count, &generation);
    if (!KSUCCESS(status))
        return status;

    RSSRV_VOLUME_PUBLICATION** link = &volumes->Published;
    while (*link) {
        RSSRV_VOLUME_PUBLICATION* publication = *link;
        if (RssrvSnapshotContainsVolume(snapshot, count, publication->InstanceId)) {
            link = &publication->Next;
            continue;
        }
        status = RssrvNamespaceUnbind(volumes->Namespace, publication->Path, publication->PathLength);
        if (!KSUCCESS(status) && status != NTLSTATUS_NOT_FOUND) {
            free(snapshot);
            return status;
        }
        *link = publication->Next;
        free(publication);
    }

    for (uint32_t i = 0; i < count; ++i) {
        if (RssrvFindPublishedVolume(volumes, snapshot[i].InstanceId))
            continue;

        SYS_HANDLE volume = SYS_HANDLE_INVALID;
        status = SysOpenVolume(volumes->Manager, snapshot[i].InstanceId,
                               SYS_RIGHT_QUERY | SYS_RIGHT_TRANSFER |
                               SYS_VOLUME_RIGHT_OPEN | SYS_VOLUME_RIGHT_MODIFY,
                               &volume);
        if (status == NTLSTATUS_NOT_FOUND || status == NTLSTATUS_OBJECT_DEAD)
            continue;
        if (!KSUCCESS(status)) {
            free(snapshot);
            return status;
        }

        RSSRV_VOLUME_PUBLICATION* publication =
            (RSSRV_VOLUME_PUBLICATION*)malloc(sizeof(*publication));
        if (!publication) {
            (void)SysCloseHandle(volume);
            free(snapshot);
            return NTLSTATUS_NO_MEMORY;
        }
        memset(publication, 0, sizeof(*publication));
        publication->InstanceId = snapshot[i].InstanceId;
        publication->NamespaceId = volumes->NextHarddiskVolumeId;
        status = RssrvBuildHarddiskVolumePath(publication->NamespaceId,
                                              publication->Path,
                                              (uint32_t)sizeof(publication->Path),
                                              &publication->PathLength);
        if (!KSUCCESS(status)) {
            free(publication);
            (void)SysCloseHandle(volume);
            free(snapshot);
            return status;
        }

        status = RssrvNamespaceBindOwned(volumes->Namespace,
                                         publication->Path,
                                         publication->PathLength,
                                         NTL_RESOURCE_KIND_VOLUME,
                                         volume,
                                         SYS_RIGHT_QUERY | SYS_VOLUME_RIGHT_OPEN | SYS_VOLUME_RIGHT_MODIFY);
        if (!KSUCCESS(status)) {
            free(publication);
            (void)SysCloseHandle(volume);
            free(snapshot);
            return status;
        }
        publication->Next = volumes->Published;
        volumes->Published = publication;
        volumes->NextHarddiskVolumeId++;
        if (!volumes->NextHarddiskVolumeId) {
            free(snapshot);
            return NTLSTATUS_LIMIT_REACHED;
        }
    }

    if (snapshot) free(snapshot);
    *generationOut = generation;
    return NTLSTATUS_SUCCESS;
}

// Copies the namespace path assigned to one mounted kernel volume instance.
KSTATUS RssrvVolumesQueryPublishedPath(RSSRV_VOLUME_STATE* volumes,
                                       uint32_t instanceId,
                                       char* pathOut,
                                       uint32_t capacity,
                                       uint32_t* lengthOut) {
    if (!volumes || !instanceId || !pathOut || !capacity || !lengthOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *lengthOut = 0;
    RSSRV_VOLUME_PUBLICATION* publication = RssrvFindPublishedVolume(volumes, instanceId);
    if (!publication) return NTLSTATUS_NOT_FOUND;
    if (publication->PathLength + 1u > capacity)
        return NTLSTATUS_BUFFER_TOO_SMALL;
    memcpy(pathOut, publication->Path, publication->PathLength + 1u);
    *lengthOut = publication->PathLength;
    return NTLSTATUS_SUCCESS;
}

// Starts the mounted-volume namespace publication worker.
KSTATUS RssrvVolumesStartWorker(RSSRV_VOLUME_STATE* volumes) {
    if (!volumes || volumes->Manager == SYS_HANDLE_INVALID)
        return NTLSTATUS_INVALID_PARAMETER;
    SYS_CREATE_THREAD_REQUEST create;
    memset(&create, 0, sizeof(create));
    create.Size = sizeof(create);
    create.Entry = (uintptr_t)RssrvVolumeWorker;
    create.Argument = (uintptr_t)volumes;
    create.Priority = 1;
    create.Flags = SYS_CREATE_THREAD_FLAG_NONE;
    return SysCreateThread(&create);
}

// Returns whether a path belongs to the RSSRV-owned automatic mounted-volume naming family.
int RssrvVolumesOwnPath(const char* path, uint32_t length) {
    static const char prefix[] = "/Device/HarddiskVolume";
    uint32_t prefixLength = (uint32_t)(sizeof(prefix) - 1u);
    return path && length >= prefixLength && memcmp(path, prefix, prefixLength) == 0;
}