mxaccessgw/src/MxGateway.Server/Workers/WorkerFrameReader.cs

using System.Buffers.Binary;
using Google.Protobuf;
using MxGateway.Contracts.Proto;

namespace MxGateway.Server.Workers;

public sealed class WorkerFrameReader
{
    private readonly WorkerFrameProtocolOptions _options;
    private readonly Stream _stream;

    public WorkerFrameReader(
        Stream stream,
        WorkerFrameProtocolOptions options)
    {
        _stream = stream ?? throw new ArgumentNullException(nameof(stream));
        _options = options ?? throw new ArgumentNullException(nameof(options));
    }

    public async ValueTask<WorkerEnvelope> ReadAsync(CancellationToken cancellationToken = default)
    {
        byte[] lengthPrefix = new byte[sizeof(uint)];
        await ReadExactlyOrThrowAsync(lengthPrefix, cancellationToken).ConfigureAwait(false);

        uint payloadLength = BinaryPrimitives.ReadUInt32LittleEndian(lengthPrefix);
        if (payloadLength == 0)
        {
            throw new WorkerFrameProtocolException(
                WorkerFrameProtocolErrorCode.MalformedLength,
                "Worker frame payload length must be greater than zero.");
        }

        if (payloadLength > _options.MaxMessageBytes)
        {
            throw new WorkerFrameProtocolException(
                WorkerFrameProtocolErrorCode.MessageTooLarge,
                $"Worker frame payload length {payloadLength} exceeds the configured maximum of {_options.MaxMessageBytes} bytes.");
        }

        byte[] payload = new byte[payloadLength];
        await ReadExactlyOrThrowAsync(payload, cancellationToken).ConfigureAwait(false);

        WorkerEnvelope envelope;
        try
        {
            envelope = WorkerEnvelope.Parser.ParseFrom(payload);
        }
        catch (InvalidProtocolBufferException exception)
        {
            throw new WorkerFrameProtocolException(
                WorkerFrameProtocolErrorCode.InvalidEnvelope,
                "Worker frame payload is not a valid WorkerEnvelope protobuf message.",
                exception);
        }

        WorkerEnvelopeValidator.Validate(envelope, _options);

        return envelope;
    }

    private async ValueTask ReadExactlyOrThrowAsync(
        Memory<byte> buffer,
        CancellationToken cancellationToken)
    {
        try
        {
            await _stream.ReadExactlyAsync(buffer, cancellationToken).ConfigureAwait(false);
        }
        catch (EndOfStreamException exception)
        {
            throw new WorkerFrameProtocolException(
                WorkerFrameProtocolErrorCode.EndOfStream,
                "Worker frame ended before the expected number of bytes were read.",
                exception);
        }
    }
}