mxaccessgw/src/ZB.MOM.WW.MxGateway.Worker/Ipc/WorkerFrameReader.cs

using System;
using System.Buffers;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
using Google.Protobuf;
using ZB.MOM.WW.MxGateway.Contracts.Proto;

namespace ZB.MOM.WW.MxGateway.Worker.Ipc;

/// <summary>Reads length-prefixed WorkerEnvelope protobuf frames from a stream.</summary>
/// <remarks>
///     <see cref="ReadAsync"/> is not reentrant: the reader keeps a per-instance length-prefix scratch
///     buffer, so exactly one consumer may be inside a read at a time. That matches how the reader is
///     used — a single read loop per <c>WorkerPipeSession</c>, with the startup handshake read
///     completing before the loop starts.
/// </remarks>
public sealed class WorkerFrameReader
{
    private readonly WorkerFrameProtocolOptions _options;
    private readonly Stream _stream;

    // Reused across frames rather than allocated per read (GWC-30). Safe because ReadAsync is
    // single-consumer by construction; the prefix is fully overwritten by every read.
    //
    // The single-consumer invariant has to survive teardown as well as steady state, because a read
    // abandoned by WorkerPipeSession's message loop still owns this buffer (and its rented payload
    // buffer) until it faults. Nothing may call ReadAsync again after that point: a second read
    // would race the abandoned one for the prefix, and could hand a pooled payload buffer back to
    // ArrayPool twice. The loop guarantees it structurally — it only ever issues a read after
    // awaiting the previous one, and it never re-enters after unwinding — and teardown only awaits
    // the abandoned read, never reissues it (WorkerPipeSession.ObserveAbandonedPipeReadAsync).
    private readonly byte[] _lengthPrefix = new byte[sizeof(uint)];

    /// <summary>Initializes the reader with a stream and protocol options.</summary>
    /// <param name="stream">Stream to read frames from.</param>
    /// <param name="options">Protocol options for frame validation.</param>
    public WorkerFrameReader(
        Stream stream,
        WorkerFrameProtocolOptions options)
    {
        _stream = stream ?? throw new ArgumentNullException(nameof(stream));
        _options = options ?? throw new ArgumentNullException(nameof(options));
    }

    /// <summary>Reads and validates a single length-prefixed frame from the stream.</summary>
    /// <param name="cancellationToken">Token to cancel the asynchronous operation.</param>
    /// <returns>The validated <see cref="WorkerEnvelope"/> read from the stream.</returns>
    public async Task<WorkerEnvelope> ReadAsync(CancellationToken cancellationToken = default)
    {
        await ReadExactlyOrThrowAsync(_lengthPrefix, sizeof(uint), cancellationToken).ConfigureAwait(false);

        uint payloadLength = 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.");
        }

        // Rent the payload buffer from the shared pool rather than allocating
        // a fresh byte[] per frame. ParseFrom copies whatever it needs into
        // the parsed message, so the rented buffer can be returned as soon as
        // parsing completes.
        int length = checked((int)payloadLength);
        byte[] payload = ArrayPool<byte>.Shared.Rent(length);
        WorkerEnvelope envelope;
        try
        {
            await ReadExactlyOrThrowAsync(payload, length, cancellationToken).ConfigureAwait(false);

            try
            {
                envelope = WorkerEnvelope.Parser.ParseFrom(payload, 0, length);
            }
            catch (InvalidProtocolBufferException exception)
            {
                throw new WorkerFrameProtocolException(
                    WorkerFrameProtocolErrorCode.InvalidEnvelope,
                    "Worker frame payload is not a valid WorkerEnvelope protobuf message.",
                    exception);
            }
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(payload);
        }

        WorkerEnvelopeValidator.Validate(envelope, _options);

        return envelope;
    }

    private static uint ReadUInt32LittleEndian(byte[] buffer)
    {
        return (uint)buffer[0]
            | ((uint)buffer[1] << 8)
            | ((uint)buffer[2] << 16)
            | ((uint)buffer[3] << 24);
    }

    private async Task ReadExactlyOrThrowAsync(
        byte[] buffer,
        int count,
        CancellationToken cancellationToken)
    {
        int offset = 0;
        while (offset < count)
        {
            // The token is forwarded but is NOT a bound on a pipe read: on .NET Framework 4.8
            // NamedPipeClientStream.ReadAsync accepts a CancellationToken and never wires it to the
            // overlapped I/O, so a read waiting on gateway bytes ignores cancellation entirely. Only
            // closing the handle ends it (WorkerPipeSession disposes the transport at teardown for
            // exactly this reason). It is still passed because non-pipe streams — the
            // MemoryStream-backed unit tests, and any future transport — do honor it.
            int bytesRead = await _stream
                .ReadAsync(buffer, offset, count - offset, cancellationToken)
                .ConfigureAwait(false);

            if (bytesRead == 0)
            {
                throw new WorkerFrameProtocolException(
                    WorkerFrameProtocolErrorCode.EndOfStream,
                    "Worker frame ended before the expected number of bytes were read.");
            }

            offset += bytesRead;
        }
    }
}