Mstsc-Script-Hook/Shared/DvcServerIo.h

#pragma once

#include <windows.h>
#include <wtsapi32.h>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <vector>

namespace hook::dvc
{
    // Antworten, die das Client-Plugin ueber IWTSVirtualChannel::Write sendet,
    // werden auf der Serverseite bei WTSVirtualChannelRead bei DVCs mit einem
    // CHANNEL_PDU_HEADER (UINT32 length + UINT32 flags) geliefert. Fuer kleine
    // Nachrichten kommt normalerweise genau ein PDU; die Funktion kann aber auch
    // mehrere Chunks zusammensetzen.
    inline bool ReadMessage(
        HANDLE channel,
        DWORD timeoutMs,
        std::vector<char>& payload,
        DWORD& win32Error)
    {
        payload.clear();
        win32Error = ERROR_SUCCESS;

        constexpr std::size_t kHeaderSize = sizeof(std::uint32_t) * 2;
        constexpr std::uint32_t kFlagFirst = 0x00000001u;
        constexpr std::uint32_t kFlagLast = 0x00000002u;
        constexpr std::size_t kMaxMessageSize = 1024 * 1024;

        bool firstRead = true;
        std::uint32_t expectedLength = 0;

        for (;;)
        {
            char buffer[4096]{};
            ULONG bytesRead = 0;
            const BOOL ok = WTSVirtualChannelRead(
                channel,
                timeoutMs,
                buffer,
                static_cast<ULONG>(sizeof(buffer)),
                &bytesRead);

            if (!ok)
            {
                win32Error = GetLastError();
                return false;
            }

            if (bytesRead == 0)
            {
                win32Error = ERROR_NO_DATA;
                return false;
            }

            // Kompatibilitaets-Fallback fuer Implementierungen, die das Payload
            // ohne CHANNEL_PDU_HEADER zurueckgeben. Unsere Protokollnachrichten
            // beginnen immer mit "PLANDENT_".
            if (firstRead &&
                bytesRead >= 9 &&
                std::memcmp(buffer, "PLANDENT_", 9) == 0)
            {
                payload.assign(buffer, buffer + bytesRead);
                return true;
            }

            if (bytesRead < kHeaderSize)
            {
                win32Error = ERROR_INVALID_DATA;
                return false;
            }

            std::uint32_t length = 0;
            std::uint32_t flags = 0;
            std::memcpy(&length, buffer, sizeof(length));
            std::memcpy(&flags, buffer + sizeof(length), sizeof(flags));

            const std::size_t chunkPayload = bytesRead - kHeaderSize;

            if (length > kMaxMessageSize || chunkPayload > length)
            {
                win32Error = ERROR_INVALID_DATA;
                return false;
            }

            if (firstRead)
            {
                expectedLength = length;
                if ((flags & kFlagFirst) == 0)
                {
                    win32Error = ERROR_INVALID_DATA;
                    return false;
                }
                payload.reserve(expectedLength);
                firstRead = false;
            }
            else if (length != expectedLength)
            {
                win32Error = ERROR_INVALID_DATA;
                return false;
            }

            payload.insert(
                payload.end(),
                buffer + static_cast<std::ptrdiff_t>(kHeaderSize),
                buffer + bytesRead);

            if (payload.size() > expectedLength)
            {
                win32Error = ERROR_INVALID_DATA;
                return false;
            }

            if ((flags & kFlagLast) != 0)
            {
                if (payload.size() != expectedLength)
                {
                    win32Error = ERROR_INVALID_DATA;
                    return false;
                }
                return true;
            }

            // Nach dem ersten Chunk gilt das Timeout erneut. Unsere Nachrichten
            // sind sehr klein; dieser Pfad ist nur fuer Vollstaendigkeit vorhanden.
            timeoutMs = std::max<DWORD>(timeoutMs, 1000);
        }
    }
}