mxaccessgw/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Workers/WorkerFrameProtocolTests.cs

using System.Buffers.Binary;
using Google.Protobuf;
using ZB.MOM.WW.MxGateway.Contracts;
using ZB.MOM.WW.MxGateway.Contracts.Proto;
using ZB.MOM.WW.MxGateway.Server.Workers;

namespace ZB.MOM.WW.MxGateway.Tests.Gateway.Workers;

public sealed class WorkerFrameProtocolTests
{
    private const string SessionId = "session-1";

    /// <summary>Verifies that writing and reading a valid envelope round-trips the frame correctly.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task WriteAndReadAsync_WithValidEnvelope_RoundTripsFrame()
    {
        WorkerFrameProtocolOptions options = new(SessionId);
        await using MemoryStream stream = new();
        WorkerEnvelope original = CreateEnvelope();

        WorkerFrameWriter writer = new(stream, options);
        await writer.WriteAsync(original);
        stream.Position = 0;

        WorkerFrameReader reader = new(stream, options);
        WorkerEnvelope parsed = await reader.ReadAsync();

        Assert.Equal(original, parsed);
    }

    /// <summary>Verifies that a large payload near the message cap round-trips through the pooled write/read path.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task WriteAndReadAsync_WithLargePayloadNearMax_RoundTripsFrame()
    {
        // A payload comfortably larger than an ArrayPool bucket boundary so the rented buffer is
        // larger than the exact frame length, exercising the length-bounded read and parse.
        const int maxMessageBytes = 1024 * 1024;
        WorkerFrameProtocolOptions options =
            new(SessionId, GatewayContractInfo.WorkerProtocolVersion, maxMessageBytes);
        await using MemoryStream stream = new();

        WorkerEnvelope original = CreateEnvelope();
        original.WorkerHello.WorkerVersion = new string('x', maxMessageBytes - 4096);
        Assert.InRange(original.CalculateSize(), 1, maxMessageBytes);

        WorkerFrameWriter writer = new(stream, options);
        await writer.WriteAsync(original);
        stream.Position = 0;

        WorkerFrameReader reader = new(stream, options);
        WorkerEnvelope parsed = await reader.ReadAsync();

        Assert.Equal(original, parsed);
    }

    /// <summary>
    ///     One reader instance reads many frames in sequence. The reader reuses a single length-prefix
    ///     scratch buffer across calls, so varying the payload length frame to frame proves the
    ///     reused buffer is fully overwritten each time rather than carrying a stale prefix forward.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithMultipleFramesOnOneReader_ParsesEveryFrame()
    {
        const int frameCount = 5;
        WorkerFrameProtocolOptions options = new(SessionId);
        await using MemoryStream stream = new();
        WorkerFrameWriter writer = new(stream, options);

        List<WorkerEnvelope> originals = [];
        for (int index = 1; index <= frameCount; index++)
        {
            WorkerEnvelope envelope = CreateEnvelope();
            envelope.Sequence = (ulong)index;

            // Differing payload lengths so a stale length prefix would misparse rather than pass.
            envelope.WorkerHello.WorkerVersion = new string('v', index * 37);
            originals.Add(envelope);
            await writer.WriteAsync(envelope);
        }

        stream.Position = 0;
        WorkerFrameReader reader = new(stream, options);
        foreach (WorkerEnvelope original in originals)
        {
            WorkerEnvelope parsed = await reader.ReadAsync();
            Assert.Equal(original, parsed);
        }
    }

    /// <summary>Verifies that reading a frame with partial reads reassembles the frame correctly.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithPartialReads_ReassemblesFrame()
    {
        WorkerFrameProtocolOptions options = new(SessionId);
        WorkerEnvelope original = CreateEnvelope();
        byte[] frame = CreateFrame(original);
        await using ChunkedReadStream stream = new(frame, chunkSize: 2);

        WorkerFrameReader reader = new(stream, options);
        WorkerEnvelope parsed = await reader.ReadAsync();

        Assert.Equal(original, parsed);
        Assert.True(stream.ReadCallCount > 2);
    }

    /// <summary>Verifies that reading a frame with zero length throws a malformed length exception.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithZeroLengthFrame_ThrowsMalformedLength()
    {
        WorkerFrameProtocolOptions options = new(SessionId);
        await using MemoryStream stream = new(new byte[sizeof(uint)]);

        WorkerFrameReader reader = new(stream, options);
        WorkerFrameProtocolException exception =
            await Assert.ThrowsAsync<WorkerFrameProtocolException>(
                async () => await reader.ReadAsync());

        Assert.Equal(WorkerFrameProtocolErrorCode.MalformedLength, exception.ErrorCode);
    }

    /// <summary>Verifies that reading a frame with oversized length throws before allocating the payload.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithOversizedLength_ThrowsBeforePayloadAllocation()
    {
        WorkerFrameProtocolOptions options = new(SessionId, GatewayContractInfo.WorkerProtocolVersion, maxMessageBytes: 16);
        byte[] lengthPrefix = new byte[sizeof(uint)];
        BinaryPrimitives.WriteUInt32LittleEndian(lengthPrefix, 17);
        await using MemoryStream stream = new(lengthPrefix);

        WorkerFrameReader reader = new(stream, options);
        WorkerFrameProtocolException exception =
            await Assert.ThrowsAsync<WorkerFrameProtocolException>(
                async () => await reader.ReadAsync());

        Assert.Equal(WorkerFrameProtocolErrorCode.MessageTooLarge, exception.ErrorCode);
    }

    /// <summary>Verifies that reading a frame with wrong protocol version throws a protocol version mismatch exception.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithWrongProtocolVersion_ThrowsProtocolVersionMismatch()
    {
        WorkerFrameProtocolOptions options = new(SessionId);
        WorkerEnvelope envelope = CreateEnvelope();
        envelope.ProtocolVersion++;
        await using MemoryStream stream = new(CreateFrame(envelope));

        WorkerFrameReader reader = new(stream, options);
        WorkerFrameProtocolException exception =
            await Assert.ThrowsAsync<WorkerFrameProtocolException>(
                async () => await reader.ReadAsync());

        Assert.Equal(WorkerFrameProtocolErrorCode.ProtocolVersionMismatch, exception.ErrorCode);
    }

    /// <summary>Verifies that reading a frame with wrong session ID throws a session mismatch exception.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithWrongSessionId_ThrowsSessionMismatch()
    {
        WorkerFrameProtocolOptions options = new(SessionId);
        WorkerEnvelope envelope = CreateEnvelope();
        envelope.SessionId = "different-session";
        await using MemoryStream stream = new(CreateFrame(envelope));

        WorkerFrameReader reader = new(stream, options);
        WorkerFrameProtocolException exception =
            await Assert.ThrowsAsync<WorkerFrameProtocolException>(
                async () => await reader.ReadAsync());

        Assert.Equal(WorkerFrameProtocolErrorCode.SessionMismatch, exception.ErrorCode);
    }

    /// <summary>Verifies that reading a frame with malformed payload throws an invalid envelope exception.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithMalformedPayload_ThrowsInvalidEnvelope()
    {
        WorkerFrameProtocolOptions options = new(SessionId);
        byte[] frame = CreateFrame([0x80]);
        await using MemoryStream stream = new(frame);

        WorkerFrameReader reader = new(stream, options);
        WorkerFrameProtocolException exception =
            await Assert.ThrowsAsync<WorkerFrameProtocolException>(
                async () => await reader.ReadAsync());

        Assert.Equal(WorkerFrameProtocolErrorCode.InvalidEnvelope, exception.ErrorCode);
    }

    /// <summary>Verifies that reading a frame with missing envelope body throws an invalid envelope exception.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithMissingEnvelopeBody_ThrowsInvalidEnvelope()
    {
        WorkerFrameProtocolOptions options = new(SessionId);
        WorkerEnvelope envelope = CreateEnvelope();
        envelope.ClearBody();
        await using MemoryStream stream = new(CreateFrame(envelope));

        WorkerFrameReader reader = new(stream, options);
        WorkerFrameProtocolException exception =
            await Assert.ThrowsAsync<WorkerFrameProtocolException>(
                async () => await reader.ReadAsync());

        Assert.Equal(WorkerFrameProtocolErrorCode.InvalidEnvelope, exception.ErrorCode);
    }

    /// <summary>Verifies that writing an oversized envelope throws a message too large exception.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task WriteAsync_WithOversizedEnvelope_ThrowsMessageTooLarge()
    {
        WorkerFrameProtocolOptions options = new(SessionId, GatewayContractInfo.WorkerProtocolVersion, maxMessageBytes: 8);
        await using MemoryStream stream = new();

        WorkerFrameWriter writer = new(stream, options);
        WorkerFrameProtocolException exception =
            await Assert.ThrowsAsync<WorkerFrameProtocolException>(
                async () => await writer.WriteAsync(CreateEnvelope()));

        Assert.Equal(WorkerFrameProtocolErrorCode.MessageTooLarge, exception.ErrorCode);
        Assert.Equal(0, stream.Length);
    }

    private static WorkerEnvelope CreateEnvelope()
    {
        return new WorkerEnvelope
        {
            ProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
            SessionId = SessionId,
            Sequence = 1,
            CorrelationId = "correlation-1",
            WorkerHello = new WorkerHello
            {
                ProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
                Nonce = "nonce",
                WorkerProcessId = 1234,
                WorkerVersion = "test-worker",
            },
        };
    }

    private static byte[] CreateFrame(IMessage message)
    {
        return CreateFrame(message.ToByteArray());
    }

    private static byte[] CreateFrame(byte[] payload)
    {
        byte[] frame = new byte[sizeof(uint) + payload.Length];
        BinaryPrimitives.WriteUInt32LittleEndian(frame.AsSpan(0, sizeof(uint)), (uint)payload.Length);
        payload.CopyTo(frame.AsSpan(sizeof(uint)));

        return frame;
    }

    private sealed class ChunkedReadStream : MemoryStream
    {
        private readonly int _chunkSize;

        /// <summary>Initializes a new instance of the <see cref="ChunkedReadStream"/> class with chunked reads.</summary>
        /// <param name="buffer">The buffer containing data to read.</param>
        /// <param name="chunkSize">The maximum number of bytes to read per operation.</param>
        public ChunkedReadStream(
            byte[] buffer,
            int chunkSize)
            : base(buffer)
        {
            _chunkSize = chunkSize;
        }

        /// <summary>Gets the number of read calls made to the stream.</summary>
        public int ReadCallCount { get; private set; }

        /// <inheritdoc />
        public override ValueTask<int> ReadAsync(
            Memory<byte> buffer,
            CancellationToken cancellationToken = default)
        {
            ReadCallCount++;
            int requestedCount = Math.Min(buffer.Length, _chunkSize);

            return base.ReadAsync(buffer[..requestedCount], cancellationToken);
        }
    }
}