mxaccessgw/src/ZB.MOM.WW.MxGateway.Worker.Tests/Ipc/WorkerFrameProtocolTests.cs

using System;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using ZB.MOM.WW.MxGateway.Contracts;
using ZB.MOM.WW.MxGateway.Contracts.Proto;
using ZB.MOM.WW.MxGateway.Worker.Ipc;
using ZB.MOM.WW.MxGateway.Worker.Tests.TestSupport;

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

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

    /// <summary>Verifies that valid envelopes round-trip through write and read.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task WriteAndReadAsync_WithValidEnvelope_RoundTripsFrame()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        using MemoryStream stream = new();
        WorkerEnvelope original = CreateGatewayHelloEnvelope();

        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 wrong protocol version throws mismatch error.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithWrongProtocolVersion_ThrowsProtocolVersionMismatch()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        WorkerEnvelope envelope = CreateGatewayHelloEnvelope();
        envelope.ProtocolVersion++;
        using MemoryStream stream = new(WorkerFrameTestHelpers.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 wrong session ID throws mismatch error.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithWrongSessionId_ThrowsSessionMismatch()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        WorkerEnvelope envelope = CreateGatewayHelloEnvelope();
        envelope.SessionId = "different-session";
        using MemoryStream stream = new(WorkerFrameTestHelpers.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 a frame whose length prefix is zero is rejected before the
    /// payload buffer is allocated. <c>docs/WorkerFrameProtocol.md</c> states the
    /// reader rejects zero-length payloads as a malformed-length error. The
    /// length prefix is the leading four bytes of the stream, so a four-zero-byte
    /// stream is exactly a frame declaring a zero-length payload.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithZeroLengthPayload_ThrowsMalformedLength()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        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 a frame whose length prefix exceeds the configured maximum
    /// is rejected before the payload buffer is allocated. <c>docs/WorkerFrameProtocol.md</c>
    /// states the reader rejects oversized payloads as a message-too-large error.
    /// A small maximum is configured so the rejection is asserted without
    /// allocating a multi-megabyte buffer.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithPayloadAboveConfiguredMaximum_ThrowsMessageTooLarge()
    {
        const int maxMessageBytes = 64;
        WorkerFrameProtocolOptions options = new(
            SessionId,
            GatewayContractInfo.WorkerProtocolVersion,
            Nonce,
            maxMessageBytes);
        byte[] frame = new byte[sizeof(uint)];
        WorkerFrameTestHelpers.WriteUInt32LittleEndian(frame, maxMessageBytes + 1);
        using MemoryStream stream = new(frame);

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

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

    /// <summary>Verifies that malformed payload throws invalid envelope error.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithMalformedPayload_ThrowsInvalidEnvelope()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        using MemoryStream stream = new(WorkerFrameTestHelpers.CreateFrame(new byte[] { 0x80 }));

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

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

    /// <summary>
    ///     Pins the <c>EndOfStream</c> branch of
    ///     <c>WorkerFrameReader.ReadExactlyOrThrowAsync</c>. The gateway
    ///     closing its end of the pipe during a partial-frame read is the
    ///     most common production transport failure; the reader must
    ///     surface this as <c>WorkerFrameProtocolErrorCode.EndOfStream</c>
    ///     so the worker session can fault deterministically rather than
    ///     spinning on a partial buffer. The stream here declares a 100-byte
    ///     payload but only supplies 50 bytes, so the inner read loop sees
    ///     <c>bytesRead == 0</c> mid-frame.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WhenStreamEndsMidFrame_ThrowsEndOfStream()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        byte[] frame = new byte[sizeof(uint) + 50];
        WorkerFrameTestHelpers.WriteUInt32LittleEndian(frame, 100);
        using MemoryStream stream = new(frame);

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

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

    /// <summary>
    ///     Pins the writer-side <c>MessageTooLarge</c> branch. A session that
    ///     constructs an envelope whose serialised size exceeds
    ///     <c>MaxMessageBytes</c> must be rejected by the writer before any
    ///     bytes are sent down the pipe, so a misbehaving producer cannot
    ///     push the receiver past its bounds. A small <c>MaxMessageBytes</c>
    ///     is configured so a modest <c>GatewayHello</c> payload — with its
    ///     nonce padded out to several hundred bytes — exceeds the limit
    ///     without allocating anything large.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task WriteAsync_WithEnvelopeAboveConfiguredMaximum_ThrowsMessageTooLarge()
    {
        const int maxMessageBytes = 64;
        WorkerFrameProtocolOptions options = new(
            SessionId,
            GatewayContractInfo.WorkerProtocolVersion,
            Nonce,
            maxMessageBytes);
        using MemoryStream stream = new();
        WorkerFrameWriter writer = new(stream, options);

        WorkerEnvelope envelope = CreateGatewayHelloEnvelope();
        envelope.GatewayHello.GatewayVersion = new string('x', 1024);

        WorkerFrameProtocolException exception =
            await Assert.ThrowsAsync<WorkerFrameProtocolException>(
                async () => await writer.WriteAsync(envelope));

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

    /// <summary>
    ///     Documents that the writer-side <c>InvalidEnvelope</c> branch
    ///     (raised when <c>WorkerEnvelope.CalculateSize()</c> returns 0) is
    ///     unreachable through public API. <c>WorkerEnvelopeValidator.Validate</c>
    ///     (run before the size check in <c>WorkerFrameWriter.WriteAsync</c>)
    ///     rejects any envelope whose <c>BodyCase</c> is <c>None</c> with
    ///     <c>InvalidEnvelope</c>; a body-less envelope is therefore
    ///     intercepted before the empty-payload branch can fire. Any
    ///     envelope carrying a typed body serialises at least the field
    ///     tag bytes, so <c>CalculateSize()</c> is strictly positive. This
    ///     test exercises the body-less path and asserts the same
    ///     <c>InvalidEnvelope</c> error code reaches the caller, pinning
    ///     the contract that "no body" is rejected before any size check.
    ///     The defensive zero-length branch in <c>WriteAsync</c> is left
    ///     in place because the cost is one comparison and removing it
    ///     would weaken the writer against future serialisation
    ///     regressions; this test makes its rationale visible.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task WriteAsync_WithEmptyEnvelope_ThrowsInvalidEnvelopeFromValidator()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        using MemoryStream stream = new();
        WorkerFrameWriter writer = new(stream, options);

        WorkerEnvelope envelope = new()
        {
            ProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
            SessionId = SessionId,
            Sequence = 1,
            // No body — BodyCase == None, validator rejects.
        };

        WorkerFrameProtocolException exception =
            await Assert.ThrowsAsync<WorkerFrameProtocolException>(
                async () => await writer.WriteAsync(envelope));

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

    /// <summary>Verifies that concurrent writes produce complete serialized frames.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task WriteAsync_WithConcurrentCalls_SerializesCompleteFrames()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        using MemoryStream stream = new();
        WorkerFrameWriter writer = new(stream, options);

        await Task.WhenAll(
            writer.WriteAsync(CreateGatewayHelloEnvelope(sequence: 1)),
            writer.WriteAsync(CreateGatewayHelloEnvelope(sequence: 2)),
            writer.WriteAsync(CreateGatewayHelloEnvelope(sequence: 3)));

        stream.Position = 0;
        WorkerFrameReader reader = new(stream, options);

        WorkerEnvelope first = await reader.ReadAsync();
        WorkerEnvelope second = await reader.ReadAsync();
        WorkerEnvelope third = await reader.ReadAsync();

        Assert.Equal(new ulong[] { 1, 2, 3 }, new[] { first.Sequence, second.Sequence, third.Sequence }.OrderBy(sequence => sequence));
    }

    /// <summary>
    /// The reader rents its payload buffer from a shared pool, so a rented
    /// buffer can be larger than the current frame and may carry bytes from
    /// a previous, larger frame. Reading frames of differing sizes
    /// back-to-back through one reader must parse each frame using only its
    /// own payload length, never trailing pooled bytes.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ReadAsync_WithVaryingFrameSizes_ParsesEachFrameExactly()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        using MemoryStream stream = new();
        WorkerFrameWriter writer = new(stream, options);

        // A large-payload frame followed by a small-payload frame: if the
        // reader reused a pooled buffer without honouring the second frame's
        // length, the small frame would parse with stale trailing bytes.
        WorkerEnvelope large = CreateGatewayHelloEnvelope(sequence: 1);
        large.GatewayHello.GatewayVersion = new string('x', 4096);
        WorkerEnvelope small = CreateGatewayHelloEnvelope(sequence: 2);

        await writer.WriteAsync(large);
        await writer.WriteAsync(small);
        stream.Position = 0;

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

        Assert.Equal(large, firstParsed);
        Assert.Equal(small, secondParsed);
    }

    private static WorkerFrameProtocolOptions CreateOptions()
    {
        return new WorkerFrameProtocolOptions(
            SessionId,
            GatewayContractInfo.WorkerProtocolVersion,
            Nonce);
    }

    /// <summary>
    /// Verifies that under concurrent writers every frame receives a distinct, gap-free sequence in
    /// strictly increasing on-wire order — the sequence is stamped by the writer at write time, so the
    /// wire order and the stamped sequence always agree (WRK-04).
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task WriteAsync_UnderConcurrentCalls_StampsGapFreeMonotonicSequence()
    {
        const int frameCount = 50;
        WorkerFrameProtocolOptions options = CreateOptions();
        using MemoryStream stream = new();
        WorkerFrameWriter writer = new(stream, options);

        await Task.WhenAll(
            Enumerable.Range(0, frameCount).Select(_ => writer.WriteAsync(CreateEventEnvelope())));

        stream.Position = 0;
        WorkerFrameReader reader = new(stream, options);
        ulong[] sequences = new ulong[frameCount];
        for (int index = 0; index < frameCount; index++)
        {
            sequences[index] = (await reader.ReadAsync()).Sequence;
        }

        // On-wire order is strictly increasing 1..frameCount with no gaps or duplicates.
        Assert.Equal(Enumerable.Range(1, frameCount).Select(value => (ulong)value), sequences);
    }

    /// <summary>
    /// Verifies that when a control frame and an event frame are both queued behind an in-progress
    /// write, the draining lock-holder writes the control frame first even though the event was queued
    /// earlier (WRK-07).
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task WriteAsync_WhenControlAndEventQueuedTogether_WritesControlFirst()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        using GatedWriteStream stream = new();
        WorkerFrameWriter writer = new(stream, options);

        // First write occupies the writer and blocks inside the stream, holding the write lock.
        Task firstWrite = writer.WriteAsync(CreateGatewayHelloEnvelope(), WorkerFrameWritePriority.Control);
        await AwaitWithTimeoutAsync(stream.FirstWriteStarted);

        // Queue an event first, then a control frame, while the writer is blocked. Both wait for the lock.
        Task eventWrite = writer.WriteAsync(CreateEventEnvelope(), WorkerFrameWritePriority.Event);
        Task controlWrite = writer.WriteAsync(CreateGatewayHelloEnvelope(), WorkerFrameWritePriority.Control);
        await Task.Delay(50);

        stream.ReleaseFirstWrite();
        await AwaitWithTimeoutAsync(Task.WhenAll(firstWrite, eventWrite, controlWrite));

        stream.Position = 0;
        WorkerFrameReader reader = new(stream, options);
        WorkerEnvelope frame1 = await reader.ReadAsync();
        WorkerEnvelope frame2 = await reader.ReadAsync();
        WorkerEnvelope frame3 = await reader.ReadAsync();

        Assert.Equal(WorkerEnvelope.BodyOneofCase.GatewayHello, frame1.BodyCase);
        // The control frame jumped ahead of the earlier-queued event.
        Assert.Equal(WorkerEnvelope.BodyOneofCase.GatewayHello, frame2.BodyCase);
        Assert.Equal(WorkerEnvelope.BodyOneofCase.WorkerEvent, frame3.BodyCase);
    }

    /// <summary>Verifies a zero negotiated frame maximum keeps the constructor default (IPC-02).</summary>
    [Fact]
    public void AdoptNegotiatedMaxMessageBytes_WithZero_KeepsDefault()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        int original = options.MaxMessageBytes;
        options.AdoptNegotiatedMaxMessageBytes(0);
        Assert.Equal(original, options.MaxMessageBytes);
    }

    /// <summary>Verifies an in-range negotiated frame maximum is adopted (IPC-02).</summary>
    [Fact]
    public void AdoptNegotiatedMaxMessageBytes_WithInRangeValue_Adopts()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        options.AdoptNegotiatedMaxMessageBytes(4 * 1024 * 1024);
        Assert.Equal(4 * 1024 * 1024, options.MaxMessageBytes);
    }

    /// <summary>Verifies a negotiated frame maximum above the worker ceiling is rejected (IPC-02).</summary>
    [Fact]
    public void AdoptNegotiatedMaxMessageBytes_AboveCeiling_Throws()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        WorkerFrameProtocolException exception = Assert.Throws<WorkerFrameProtocolException>(
            () => options.AdoptNegotiatedMaxMessageBytes((uint)WorkerFrameProtocolOptions.MaxNegotiableFrameBytes + 1));
        Assert.Equal(WorkerFrameProtocolErrorCode.InvalidConfiguration, exception.ErrorCode);
    }

    private static WorkerEnvelope CreateGatewayHelloEnvelope(ulong sequence = 1)
    {
        return new WorkerEnvelope
        {
            ProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
            SessionId = SessionId,
            Sequence = sequence,
            GatewayHello = new GatewayHello
            {
                SupportedProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
                Nonce = Nonce,
                GatewayVersion = "test-gateway",
            },
        };
    }

    // net48 has no Task.WaitAsync(TimeSpan); fail the test rather than hang if the writer misbehaves.
    private static async Task AwaitWithTimeoutAsync(Task task)
    {
        Task completed = await Task.WhenAny(task, Task.Delay(TimeSpan.FromSeconds(5)));
        if (completed != task)
        {
            throw new TimeoutException("Timed out waiting for the frame writer.");
        }

        await task;
    }

    private static WorkerEnvelope CreateEventEnvelope()
    {
        return new WorkerEnvelope
        {
            ProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
            SessionId = SessionId,
            WorkerEvent = new WorkerEvent
            {
                Event = new MxEvent { SessionId = SessionId },
            },
        };
    }

    // A MemoryStream whose first WriteAsync blocks until released, so a test can queue additional frames
    // behind an in-progress write and observe the writer's priority ordering.
    private sealed class GatedWriteStream : MemoryStream
    {
        private readonly SemaphoreSlim _release = new SemaphoreSlim(0);
        private readonly TaskCompletionSource<bool> _firstWriteStarted =
            new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
        private int _writeCount;

        public Task FirstWriteStarted => _firstWriteStarted.Task;

        public void ReleaseFirstWrite() => _release.Release();

        public override async Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
        {
            if (Interlocked.Increment(ref _writeCount) == 1)
            {
                _firstWriteStarted.TrySetResult(true);
                await _release.WaitAsync(cancellationToken);
            }

            await base.WriteAsync(buffer, offset, count, cancellationToken);
        }

        protected override void Dispose(bool disposing)
        {
            if (disposing)
            {
                _release.Dispose();
            }

            base.Dispose(disposing);
        }
    }
}