lmxopcua/tests/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Modbus.Tests/ModbusTcpReconnectTests.cs

using System.Net;
using System.Net.Sockets;
using Shouldly;
using Xunit;

namespace ZB.MOM.WW.OtOpcUa.Driver.Modbus.Tests;

/// <summary>
///     Exercises <see cref="ModbusTcpTransport"/> against a real TCP listener that can close
///     its socket mid-session on demand. Verifies the PR 53 reconnect-on-drop behavior: after
///     the "first" socket is forcibly torn down, the next SendAsync must re-establish the
///     connection and complete the PDU without bubbling an error to the caller.
/// </summary>
[Trait("Category", "Unit")]
public sealed class ModbusTcpReconnectTests
{
    /// <summary>
    ///     Minimal in-process Modbus-TCP stub. Accepts one TCP connection at a time, reads an
    ///     MBAP + PDU, replies with a canned FC03 response echoing the request quantity of
    ///     zeroed bytes, then optionally closes the socket to simulate a NAT/firewall drop.
    /// </summary>
    private sealed class FlakeyModbusServer : IAsyncDisposable
    {
        private readonly TcpListener _listener;
        /// <summary>Gets the TCP port the server is listening on.</summary>
        public int Port => ((IPEndPoint)_listener.LocalEndpoint).Port;
        /// <summary>Gets or sets the number of transactions to complete before closing the connection.</summary>
        public int DropAfterNTransactions { get; set; } = int.MaxValue;
        private readonly CancellationTokenSource _stop = new();
        private int _txCount;

        // --- Scriptable STAB-3 desync behaviors (each fires for the first N transactions, then
        //     the server behaves normally). Decremented atomically because an old (torn-down) serve
        //     task and a fresh one can briefly coexist. ---
        private int _stallRemaining;
        private int _wrongTxRemaining;
        private int _invalidLenRemaining;
        private int _acceptedConnections;

        /// <summary>The first N responses are withheld entirely — the client blocks until its per-op timeout / caller cancellation fires.</summary>
        public int StallFirstNResponses { set => _stallRemaining = value; }
        /// <summary>The first N responses echo a corrupted MBAP transaction id (framing violation).</summary>
        public int WrongTxIdFirstNResponses { set => _wrongTxRemaining = value; }
        /// <summary>The first N responses claim an MBAP length field of 0 (truncated-length framing violation).</summary>
        public int InvalidLengthFirstNResponses { set => _invalidLenRemaining = value; }
        /// <summary>Gets the number of TCP connections the listener has accepted (a reconnect increments this).</summary>
        public int AcceptedConnectionCount => Volatile.Read(ref _acceptedConnections);

        /// <summary>Initializes a new instance and starts listening on a loopback port.</summary>
        public FlakeyModbusServer()
        {
            _listener = new TcpListener(IPAddress.Loopback, 0);
            _listener.Start();
            _ = Task.Run(AcceptLoopAsync);
        }

        private async Task AcceptLoopAsync()
        {
            while (!_stop.IsCancellationRequested)
            {
                TcpClient? client = null;
                try { client = await _listener.AcceptTcpClientAsync(_stop.Token); }
                catch { return; }

                Interlocked.Increment(ref _acceptedConnections);
                _ = Task.Run(() => ServeAsync(client!));
            }
        }

        private async Task ServeAsync(TcpClient client)
        {
            try
            {
                using var _ = client;
                var stream = client.GetStream();
                while (!_stop.IsCancellationRequested && client.Connected)
                {
                    var header = new byte[7];
                    if (!await ReadExactly(stream, header)) return;
                    var len = (ushort)((header[4] << 8) | header[5]);
                    var pdu = new byte[len - 1];
                    if (!await ReadExactly(stream, pdu)) return;

                    // Stall: read the request but never answer it — simulate a hung / unreachable
                    // unit so the client's per-op CancelAfter (or caller token) fires.
                    if (Interlocked.Decrement(ref _stallRemaining) >= 0)
                        continue;

                    var fc = pdu[0];
                    var qty = (ushort)((pdu[3] << 8) | pdu[4]);

                    // Truncated-length framing violation: a 7-byte header whose length field is 0
                    // (< the mandatory 1 unit-id byte) — trips the transport's `respLen < 1` guard.
                    if (Interlocked.Decrement(ref _invalidLenRemaining) >= 0)
                    {
                        var bad = new byte[7];
                        bad[0] = header[0]; bad[1] = header[1];
                        bad[4] = 0; bad[5] = 0; // length field = 0
                        bad[6] = header[6];
                        await stream.WriteAsync(bad);
                        await stream.FlushAsync();
                        continue;
                    }

                    var respPdu = new byte[2 + qty * 2];
                    respPdu[0] = fc;
                    respPdu[1] = (byte)(qty * 2);
                    // data bytes stay 0

                    var respLen = (ushort)(1 + respPdu.Length);
                    var adu = new byte[7 + respPdu.Length];
                    adu[0] = header[0]; adu[1] = header[1];

                    // TxId-mismatch framing violation: emit an otherwise valid, full frame but with
                    // the transaction id flipped, so the transport reads the whole (unexpected)
                    // response and rejects it on the TxId guard.
                    if (Interlocked.Decrement(ref _wrongTxRemaining) >= 0)
                    {
                        adu[0] = (byte)~header[0];
                        adu[1] = (byte)~header[1];
                    }

                    adu[4] = (byte)(respLen >> 8); adu[5] = (byte)(respLen & 0xFF);
                    adu[6] = header[6];
                    Buffer.BlockCopy(respPdu, 0, adu, 7, respPdu.Length);
                    await stream.WriteAsync(adu);
                    await stream.FlushAsync();

                    _txCount++;
                    if (_txCount >= DropAfterNTransactions)
                    {
                        // Simulate NAT/firewall silent close: slam the socket without a
                        // protocol-level goodbye, which is what DL260 + an intermediate
                        // middlebox would look like from the client's perspective.
                        client.Client.Shutdown(SocketShutdown.Both);
                        client.Close();
                        return;
                    }
                }
            }
            catch { /* best-effort */ }
        }

        private static async Task<bool> ReadExactly(NetworkStream s, byte[] buf)
        {
            var read = 0;
            while (read < buf.Length)
            {
                var n = await s.ReadAsync(buf.AsMemory(read));
                if (n == 0) return false;
                read += n;
            }
            return true;
        }

        /// <summary>Stops the server and releases resources.</summary>
        public async ValueTask DisposeAsync()
        {
            _stop.Cancel();
            _listener.Stop();
            await Task.CompletedTask;
        }
    }

    /// <summary>Verifies that the transport recovers from a mid-session socket drop with auto-reconnect enabled.</summary>
    [Fact]
    public async Task Transport_recovers_from_mid_session_drop_and_retries_successfully()
    {
        await using var server = new FlakeyModbusServer { DropAfterNTransactions = 1 };
        await using var transport = new ModbusTcpTransport("127.0.0.1", server.Port, TimeSpan.FromSeconds(2), autoReconnect: true);
        await transport.ConnectAsync(TestContext.Current.CancellationToken);

        // First transaction succeeds; server then closes the socket.
        var pdu = new byte[] { 0x03, 0x00, 0x00, 0x00, 0x01 };
        var first = await transport.SendAsync(unitId: 1, pdu, TestContext.Current.CancellationToken);
        first[0].ShouldBe((byte)0x03);

        // Second transaction: the connection is dead, but auto-reconnect must transparently
        // spin up a new socket, resend, and produce a valid response. Before PR 53 this would
        // surface as EndOfStreamException / IOException to the caller.
        var second = await transport.SendAsync(unitId: 1, pdu, TestContext.Current.CancellationToken);
        second[0].ShouldBe((byte)0x03);
    }

    /// <summary>Verifies that socket drops propagate to the caller when auto-reconnect is disabled.</summary>
    [Fact]
    public async Task Transport_without_AutoReconnect_propagates_drop_to_caller()
    {
        await using var server = new FlakeyModbusServer { DropAfterNTransactions = 1 };
        await using var transport = new ModbusTcpTransport("127.0.0.1", server.Port, TimeSpan.FromSeconds(2), autoReconnect: false);
        await transport.ConnectAsync(TestContext.Current.CancellationToken);

        var pdu = new byte[] { 0x03, 0x00, 0x00, 0x00, 0x01 };
        _ = await transport.SendAsync(unitId: 1, pdu, TestContext.Current.CancellationToken);

        await Should.ThrowAsync<Exception>(async () =>
            await transport.SendAsync(unitId: 1, pdu, TestContext.Current.CancellationToken));
    }

    // --- STAB-3: per-op timeout + framing violations must be classified connection-fatal so the
    //     socket is torn down instead of left desynchronized. With auto-reconnect on, the transport
    //     transparently reconnects + resends once and returns the CORRECT response (no stale read);
    //     the reconnect is observable as a second accepted connection on the server. Before the fix
    //     the timeout OCE / InvalidDataException are not socket-level, so no teardown happens and the
    //     SendAsync throws (and the socket stays desynchronized forever). ---

    private static readonly byte[] Fc03Qty1 = { 0x03, 0x00, 0x00, 0x00, 0x01 };

    /// <summary>A per-op response timeout tears down the socket and (auto-reconnect) resends cleanly.</summary>
    [Fact]
    public async Task SendAsync_AfterResponseTimeout_TearsDownAndReconnects()
    {
        await using var server = new FlakeyModbusServer { StallFirstNResponses = 1 };
        await using var transport = new ModbusTcpTransport(
            "127.0.0.1", server.Port, TimeSpan.FromMilliseconds(300), autoReconnect: true);
        await transport.ConnectAsync(TestContext.Current.CancellationToken);

        var resp = await transport.SendAsync(unitId: 1, Fc03Qty1, TestContext.Current.CancellationToken);

        resp[0].ShouldBe((byte)0x03);          // correct FC, not stale bytes
        resp.Length.ShouldBe(2 + 1 * 2);       // FC + byteCount + qty*2 data
        server.AcceptedConnectionCount.ShouldBe(2); // torn down + reconnected
    }

    /// <summary>A TxId-mismatch framing violation tears down the socket and (auto-reconnect) resends cleanly.</summary>
    [Fact]
    public async Task SendAsync_AfterTxIdMismatch_TearsDownSocket()
    {
        await using var server = new FlakeyModbusServer { WrongTxIdFirstNResponses = 1 };
        await using var transport = new ModbusTcpTransport(
            "127.0.0.1", server.Port, TimeSpan.FromSeconds(2), autoReconnect: true);
        await transport.ConnectAsync(TestContext.Current.CancellationToken);

        var resp = await transport.SendAsync(unitId: 1, Fc03Qty1, TestContext.Current.CancellationToken);

        resp[0].ShouldBe((byte)0x03);
        server.AcceptedConnectionCount.ShouldBe(2);
    }

    /// <summary>A truncated-length framing violation tears down the socket and (auto-reconnect) resends cleanly.</summary>
    [Fact]
    public async Task SendAsync_AfterTruncatedHeader_TearsDownSocket()
    {
        await using var server = new FlakeyModbusServer { InvalidLengthFirstNResponses = 1 };
        await using var transport = new ModbusTcpTransport(
            "127.0.0.1", server.Port, TimeSpan.FromSeconds(2), autoReconnect: true);
        await transport.ConnectAsync(TestContext.Current.CancellationToken);

        var resp = await transport.SendAsync(unitId: 1, Fc03Qty1, TestContext.Current.CancellationToken);

        resp[0].ShouldBe((byte)0x03);
        server.AcceptedConnectionCount.ShouldBe(2);
    }

    /// <summary>Caller cancellation propagates as OCE and must NOT tear down the socket (a legitimate shutdown is not a desync).</summary>
    [Fact]
    public async Task SendAsync_CallerCancellation_DoesNotTearDown()
    {
        await using var server = new FlakeyModbusServer { StallFirstNResponses = 1 };
        await using var transport = new ModbusTcpTransport(
            "127.0.0.1", server.Port, TimeSpan.FromSeconds(5), autoReconnect: true);
        await transport.ConnectAsync(TestContext.Current.CancellationToken);

        using var cts = new CancellationTokenSource();
        var pending = transport.SendAsync(unitId: 1, Fc03Qty1, cts.Token);
        await Task.Delay(150, TestContext.Current.CancellationToken); // let the request land + block on the response read
        await cts.CancelAsync();

        await Should.ThrowAsync<OperationCanceledException>(async () => await pending);

        // The socket must still be usable — a subsequent transaction succeeds on the SAME connection
        // (no reconnect), proving caller-cancel did not trip the desync teardown.
        var resp = await transport.SendAsync(unitId: 1, Fc03Qty1, TestContext.Current.CancellationToken);
        resp[0].ShouldBe((byte)0x03);
        server.AcceptedConnectionCount.ShouldBe(1);
    }
}