lmxopcua/tests/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Tests/Fakes/FakeEipServer.cs

using System.Buffers.Binary;
using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using System.Threading.Channels;
using ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Eip;
namespace ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Tests.Fakes;

/// <summary>
///     Scriptable loopback EtherNet/IP endpoint. RegisterSession is answered automatically (session
///     <see cref="Session"/>) and UnRegisterSession is recorded but never answered. Every other frame is
///     answered by the next queued <see cref="Script"/> step, else by <see cref="HandlerAsync"/>; a
///     non-null result is written back verbatim (it may hold several concatenated frames).
/// </summary>
/// <remarks>
///     Each connection runs a reader loop and a responder loop, so frames are recorded the moment they
///     arrive — even while a reply is still pending. <see cref="MaxOutstanding"/> therefore really does
///     see a client that pipelines a second request before the first reply.
/// </remarks>
internal sealed class FakeEipServer : IAsyncDisposable
{
    public const uint Session = 0x11223344;

    private readonly TcpListener _listener = new(IPAddress.Loopback, 0);
    private readonly CancellationTokenSource _cts = new();
    private readonly ConcurrentQueue<Func<byte[], Task<byte[]?>>> _script = new();
    private readonly List<byte[]> _received = [];
    private readonly List<(Predicate<byte[]> Match, TaskCompletionSource<byte[]> Tcs)> _waiters = [];
    private readonly ConcurrentBag<TcpClient> _clients = [];
    private readonly ConcurrentBag<Exception> _handlerErrors = [];
    private int _connections;
    private int _open;
    private int _outstanding;
    private int _maxOutstanding;

    public FakeEipServer()
    {
        _listener.Start();
        _ = AcceptLoop();
    }

    public int Port => ((IPEndPoint)_listener.LocalEndpoint).Port;

    /// <summary>Every frame received, in arrival order (snapshot).</summary>
    public IReadOnlyList<byte[]> Received { get { lock (_received) return [.. _received]; } }

    public int Connections => Volatile.Read(ref _connections);

    /// <summary>TCP connections accepted and not yet closed by the client (Gitea #579 leak checks).</summary>
    public int OpenConnections => Volatile.Read(ref _open);

    /// <summary>Largest number of non-session requests ever received but not yet answered.</summary>
    public int MaxOutstanding => Volatile.Read(ref _maxOutstanding);

    /// <summary>Exceptions thrown by handlers / script steps (a test assertion inside a handler lands here).</summary>
    public IReadOnlyCollection<Exception> HandlerErrors => _handlerErrors;

    /// <summary>Default responder, used once the script is exhausted.</summary>
    public Func<byte[], Task<byte[]?>> HandlerAsync { get; set; } = _ => Task.FromResult<byte[]?>(null);

    /// <summary>Synchronous convenience setter for <see cref="HandlerAsync"/>.</summary>
    public Func<byte[], byte[]?> Handler { set => HandlerAsync = req => Task.FromResult(value(req)); }

    /// <summary>Awaited before a RegisterSession is answered (the request is already recorded).</summary>
    public Func<Task>? RegisterSessionGate { get; set; }

    /// <summary>Never answer non-session requests.</summary>
    public bool Stall { get; set; }

    /// <summary>Zero the sender context (header bytes 12-19) of every reply frame, as ab_server does.</summary>
    public bool ZeroReplyContext { get; set; }

    /// <summary>Close the connection instead of answering the next non-session request.</summary>
    public bool DropOnNextRequest { get; set; }

    /// <summary>Queue one-shot responders, consumed one per non-session request.</summary>
    public void Script(params Func<byte[], byte[]?>[] steps)
    {
        foreach (var s in steps) _script.Enqueue(req => Task.FromResult(s(req)));
    }

    /// <summary>Queue one-shot asynchronous responders (e.g. one that waits for the test to release it).</summary>
    public void ScriptAsync(params Func<byte[], Task<byte[]?>>[] steps)
    {
        foreach (var s in steps) _script.Enqueue(s);
    }

    /// <summary>Completes with the first received frame (past or future) matching <paramref name="match"/>.</summary>
    public Task<byte[]> WaitForFrameAsync(Predicate<byte[]> match, TimeSpan timeout)
    {
        var tcs = new TaskCompletionSource<byte[]>(TaskCreationOptions.RunContinuationsAsynchronously);
        lock (_received)
        {
            var hit = _received.FirstOrDefault(f => match(f));
            if (hit is not null) return Task.FromResult(hit);
            _waiters.Add((match, tcs));
        }
        return tcs.Task.WaitAsync(timeout);
    }

    /// <summary>Completes once <paramref name="count"/> frames with encapsulation command <paramref name="command"/> have arrived.</summary>
    public async Task WaitForCommandCountAsync(ushort command, int count, TimeSpan timeout)
    {
        var seen = 0;
        await WaitForFrameAsync(f => Command(f) == command && ++seen >= count, timeout);
    }

    public static ushort Command(byte[] frame) => BinaryPrimitives.ReadUInt16LittleEndian(frame);

    private async Task AcceptLoop()
    {
        while (!_cts.IsCancellationRequested)
        {
            TcpClient c;
            try { c = await _listener.AcceptTcpClientAsync(_cts.Token); }
            catch { return; }
            Interlocked.Increment(ref _connections);
            Interlocked.Increment(ref _open);
            _clients.Add(c);
            _ = Serve(c);
        }
    }

    private async Task Serve(TcpClient c)
    {
        try { await ServeCore(c); }
        finally { Interlocked.Decrement(ref _open); }
    }

    private async Task ServeCore(TcpClient c)
    {
        using var _ = c;
        var s = c.GetStream();
        var inbox = Channel.CreateUnbounded<byte[]>();
        var responder = Respond(c, s, inbox.Reader);
        var head = new byte[EncapsulationHeader.Size];
        while (!_cts.IsCancellationRequested)
        {
            if (!await ReadExact(s, head)) break;
            var len = head[2] | (head[3] << 8);
            var frame = new byte[EncapsulationHeader.Size + len];
            head.CopyTo(frame, 0);
            if (!await ReadExact(s, frame.AsMemory(EncapsulationHeader.Size))) break;
            if (Command(frame) is not (EncapCommand.RegisterSession or EncapCommand.UnRegisterSession))
            {
                var now = Interlocked.Increment(ref _outstanding);
                int max;
                while (now > (max = Volatile.Read(ref _maxOutstanding)) &&
                       Interlocked.CompareExchange(ref _maxOutstanding, now, max) != max) { }
            }
            Record(frame);
            inbox.Writer.TryWrite(frame);
        }
        inbox.Writer.TryComplete();
        try { await responder; } catch { /* connection gone */ }
    }

    private async Task Respond(TcpClient c, NetworkStream s, ChannelReader<byte[]> inbox)
    {
        await foreach (var frame in inbox.ReadAllAsync(_cts.Token))
        {
            byte[]? reply;
            switch (Command(frame))
            {
                case EncapCommand.RegisterSession:
                    if (RegisterSessionGate is { } gate) await gate();
                    reply = (byte[])frame.Clone();
                    BinaryPrimitives.WriteUInt32LittleEndian(reply.AsSpan(4), Session);
                    break;
                case EncapCommand.UnRegisterSession:
                    reply = null;
                    break;
                default:
                    if (DropOnNextRequest)
                    {
                        DropOnNextRequest = false;
                        c.Close();
                        return;
                    }
                    if (Stall) continue;          // stays outstanding forever
                    var step = _script.TryDequeue(out var scripted) ? scripted : HandlerAsync;
                    try { reply = await step(frame); }
                    catch (Exception ex) { _handlerErrors.Add(ex); reply = null; }
                    Interlocked.Decrement(ref _outstanding);   // before the write: the client cannot send
                    break;                                     // its next request until it has this reply
            }
            if (reply is null) continue;
            if (ZeroReplyContext) ZeroContexts(reply);
            await s.WriteAsync(reply, _cts.Token);
        }
    }

    private void Record(byte[] frame)
    {
        List<TaskCompletionSource<byte[]>> fire = [];
        lock (_received)
        {
            _received.Add(frame);
            for (var i = _waiters.Count - 1; i >= 0; i--)
            {
                if (!_waiters[i].Match(frame)) continue;
                fire.Add(_waiters[i].Tcs);
                _waiters.RemoveAt(i);
            }
        }
        foreach (var t in fire) t.TrySetResult(frame);
    }

    private static void ZeroContexts(byte[] buffer)
    {
        for (var off = 0; off + EncapsulationHeader.Size <= buffer.Length;
             off += EncapsulationHeader.Size + BinaryPrimitives.ReadUInt16LittleEndian(buffer.AsSpan(off + 2)))
            buffer.AsSpan(off + 12, 8).Clear();
    }

    private async Task<bool> ReadExact(NetworkStream s, Memory<byte> buf)
    {
        var got = 0;
        while (got < buf.Length)
        {
            int n;
            try { n = await s.ReadAsync(buf[got..], _cts.Token); }
            catch { return false; }
            if (n == 0) return false;
            got += n;
        }
        return true;
    }

    // ---- reply builders -------------------------------------------------------------------------

    /// <summary>The CIP payload of a SendRRData / SendUnitData request.</summary>
    public static byte[] RequestCip(byte[] request) =>
        EncapsulationCodec.ParseReply(request, expectConnected: Command(request) == EncapCommand.SendUnitData).Cip.ToArray();

    /// <summary>The connected sequence count of a SendUnitData request.</summary>
    public static ushort RequestSequence(byte[] request) =>
        EncapsulationCodec.ParseReply(request, expectConnected: true).Sequence;

    /// <summary>A SendRRData reply echoing the request's session + context, carrying <paramref name="cipReply"/>.</summary>
    public static byte[] RRReply(byte[] request, byte[] cipReply) =>
        RRReplyWithContext(request, BinaryPrimitives.ReadUInt64LittleEndian(request.AsSpan(12)), cipReply);

    /// <summary>A SendRRData reply with an explicit sender context.</summary>
    public static byte[] RRReplyWithContext(byte[] request, ulong ctx, byte[] cipReply) =>
        EncapsulationCodec.SendRRData(BinaryPrimitives.ReadUInt32LittleEndian(request.AsSpan(4)), ctx, cipReply, 0);

    /// <summary>A SendUnitData reply echoing the request's sequence count, carrying <paramref name="cipReply"/>.</summary>
    public static byte[] UnitReply(byte[] request, byte[] cipReply) =>
        UnitReplyWithSequence(request, RequestSequence(request), cipReply);

    /// <summary>A SendUnitData reply with an explicit sequence count (T→O connection id 0x5A5A5A5A).</summary>
    public static byte[] UnitReplyWithSequence(byte[] request, ushort sequence, byte[] cipReply) =>
        EncapsulationCodec.SendUnitData(BinaryPrimitives.ReadUInt32LittleEndian(request.AsSpan(4)),
            BinaryPrimitives.ReadUInt64LittleEndian(request.AsSpan(12)), 0x5A5A5A5A, sequence, cipReply);

    public async ValueTask DisposeAsync()
    {
        await _cts.CancelAsync();
        _listener.Stop();
        foreach (var c in _clients)
        {
            try { c.Dispose(); } catch { /* already gone */ }
        }
    }
}