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

using System.Net;
using System.Net.Sockets;

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

/// <summary>
///     Loopback Modbus-TCP stub whose reply to each request is scripted by the test, and whose listener can
///     be stopped (connection refused) and restarted on the SAME port. Used by the transport self-heal tests
///     that need a device which disappears and comes back.
/// </summary>
internal sealed class ScriptedModbusTcpServer : IAsyncDisposable
{
    /// <summary>What the server does with one request.</summary>
    public enum Action
    {
        /// <summary>Answer with an FC03-shaped response of zeroed registers.</summary>
        Answer,
        /// <summary>Read the request and never answer it (the client's per-op timeout fires).</summary>
        Stall,
        /// <summary>Close the socket without answering.</summary>
        Drop,
    }

    private TcpListener? _listener;
    private CancellationTokenSource _stop = new();
    private int _requests;
    private int _accepted;

    /// <summary>Optional scripted reply. Receives the 0-based request index and the request PDU; returns
    /// the raw bytes to write AFTER the MBAP TxId (i.e. the full ADU is built by the test), or null to fall
    /// through to <see cref="Script"/>.</summary>
    public Func<int, byte[], byte[], byte[]?>? RawReply { get; set; }

    /// <summary>Per-request action; defaults to <see cref="Action.Answer"/>.</summary>
    public Func<int, Action> Script { get; set; } = _ => Action.Answer;

    /// <summary>Gets the port the server listens on (stable across <see cref="Stop"/>/<see cref="Start"/>).</summary>
    public int Port { get; }

    /// <summary>Gets the number of accepted TCP connections.</summary>
    public int AcceptedConnections => Volatile.Read(ref _accepted);

    /// <summary>Gets the number of requests received.</summary>
    public int Requests => Volatile.Read(ref _requests);

    public ScriptedModbusTcpServer()
    {
        var probe = new TcpListener(IPAddress.Loopback, 0);
        probe.Start();
        Port = ((IPEndPoint)probe.LocalEndpoint).Port;
        probe.Stop();
        Start();
    }

    /// <summary>(Re)starts listening on <see cref="Port"/>.</summary>
    public void Start()
    {
        _stop = new CancellationTokenSource();
        _listener = new TcpListener(IPAddress.Loopback, Port);
        _listener.Server.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.ReuseAddress, true);
        _listener.Start();
        var listener = _listener;
        var token = _stop.Token;
        _ = Task.Run(() => AcceptLoopAsync(listener, token));
    }

    /// <summary>Stops listening and closes every live connection — new connects are refused.</summary>
    public void Stop()
    {
        try { _stop.Cancel(); } catch { }
        try { _listener?.Stop(); } catch { }
        _listener = null;
    }

    private async Task AcceptLoopAsync(TcpListener listener, CancellationToken token)
    {
        while (!token.IsCancellationRequested)
        {
            TcpClient client;
            try { client = await listener.AcceptTcpClientAsync(token); }
            catch { return; }
            Interlocked.Increment(ref _accepted);
            _ = Task.Run(() => ServeAsync(client, token));
        }
    }

    private async Task ServeAsync(TcpClient client, CancellationToken token)
    {
        using var _ = client;
        using var reg = token.Register(() => { try { client.Close(); } catch { } });
        try
        {
            var stream = client.GetStream();
            while (!token.IsCancellationRequested)
            {
                var header = new byte[7];
                if (!await ReadExactly(stream, header, token)) return;
                var len = (ushort)((header[4] << 8) | header[5]);
                var pdu = new byte[len - 1];
                if (!await ReadExactly(stream, pdu, token)) return;
                var idx = Interlocked.Increment(ref _requests) - 1;

                var raw = RawReply?.Invoke(idx, header, pdu);
                if (raw is not null)
                {
                    await stream.WriteAsync(raw, token);
                    await stream.FlushAsync(token);
                    continue;
                }

                switch (Script(idx))
                {
                    case Action.Stall:
                        continue;
                    case Action.Drop:
                        client.Client.Shutdown(SocketShutdown.Both);
                        return;
                }

                var qty = (ushort)((pdu[3] << 8) | pdu[4]);
                var respPdu = new byte[2 + qty * 2];
                respPdu[0] = pdu[0];
                respPdu[1] = (byte)(qty * 2);
                await stream.WriteAsync(Adu(header, respPdu), token);
                await stream.FlushAsync(token);
            }
        }
        catch { /* connection torn down */ }
    }

    /// <summary>Builds a reply ADU echoing the request's TxId/unit around <paramref name="respPdu"/>.</summary>
    public static byte[] Adu(byte[] requestHeader, byte[] respPdu)
    {
        var respLen = (ushort)(1 + respPdu.Length);
        var adu = new byte[7 + respPdu.Length];
        adu[0] = requestHeader[0]; adu[1] = requestHeader[1];
        adu[4] = (byte)(respLen >> 8); adu[5] = (byte)(respLen & 0xFF);
        adu[6] = requestHeader[6];
        Buffer.BlockCopy(respPdu, 0, adu, 7, respPdu.Length);
        return adu;
    }

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

    public ValueTask DisposeAsync()
    {
        Stop();
        return ValueTask.CompletedTask;
    }
}