lmxopcua/src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport/Eip/EncapsulationClient.cs

using System.Net.Sockets;
namespace ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Eip;

/// <summary>
///     One EtherNet/IP TCP session: connects, registers, and sends SendRRData (unconnected) and
///     SendUnitData (connected) requests with <b>at most one request in flight</b>.
/// </summary>
/// <remarks>
///     <para><b>Single flight.</b> Every exchange (write one frame, read until its reply) runs under a
///     <see cref="SemaphoreSlim"/>(1,1). The connected sequence count is advanced inside that gate, so
///     concurrent callers can never put sequences on the wire out of order.</para>
///     <para><b>Caller cancellation never faults the connection, and never misattributes a reply</b>
///     (rulings R5-1/R5-2). The caller gets <see cref="OperationCanceledException"/> at once, but the
///     abandoned exchange keeps reading <i>its own</i> reply in the background while still holding the
///     gate, bounded by the original timeout. The reply is discarded and the connection stays healthy;
///     if the timeout elapses first the connection faults exactly as a normal timeout does. Because the
///     next request is not written until the previous reply has been consumed, a late reply can never be
///     taken for the next request's — which matters because ab_server zeroes the reply sender context,
///     so context matching alone cannot tell two unconnected replies apart.</para>
///     <para><b>Faults.</b> A timeout, an I/O error or an encapsulation-level protocol error
///     (<see cref="AbStatus.BadReply"/>) closes the socket; <see cref="IsConnected"/> turns false and
///     every later call throws <see cref="AbTransportException"/>(<see cref="AbStatus.BadConnection"/>).
///     The owner reconnects with a fresh client.</para>
///     <para><b>No half-written frames.</b> A frame write never observes a cancellation token. It is
///     nevertheless bounded: the exchange timeout also covers the write, and if it elapses first the
///     connection is <i>faulted</i> (socket closed, which fails the pending write) rather than the write
///     being cancelled — so a peer that stops reading cannot hold the single-flight gate forever, and a
///     torn frame can only ever sit on a socket that is already dead.</para>
///     <para><b>Timeouts are finite.</b> Every entry point rejects a non-positive or infinite timeout
///     (<see cref="ArgumentOutOfRangeException"/>): the abandoned-exchange drain is bounded by it, so an
///     infinite one would let a cancelled caller's drain hold the gate forever. The gate wait itself is
///     bounded by the same timeout (<see cref="AbStatus.Timeout"/>, nothing written, connection intact).</para>
///     <para><b>RegisterSession.</b> It goes out with a zero sender context (ab_server refuses a non-zero
///     one) and therefore relies on the unconnected "reply context == 0 is tolerated" rule in
///     <c>TryMatch</c> to match its reply; both ends of that dependency are pinned by tests.</para>
/// </remarks>
public sealed class EncapsulationClient : IAsyncDisposable
{
    private static readonly TimeSpan UnregisterBudget = TimeSpan.FromMilliseconds(500);

    private readonly TcpClient _tcp;
    private readonly Stream _stream;

    // Never disposed: an abandoned exchange may release it after DisposeAsync returns, and
    // SemaphoreSlim.Release on a disposed instance throws. It holds no unmanaged resource unless
    // AvailableWaitHandle is touched, which this class never does.
    private readonly SemaphoreSlim _gate = new(1, 1);
    private long _ctx;
    private ushort _seq;             // connected sequence count; only touched under _gate
    private volatile bool _faulted;
    private int _disposed;

    /// <summary>The session handle the target assigned in its RegisterSession reply.</summary>
    public uint SessionHandle { get; private set; }

    /// <summary>False once the connection has faulted or been disposed.</summary>
    public bool IsConnected => !_faulted && Volatile.Read(ref _disposed) == 0 && _tcp.Connected;

    private EncapsulationClient(TcpClient tcp, Stream? stream = null)
    {
        _tcp = tcp;
        _stream = stream ?? tcp.GetStream();
    }

    /// <summary>Test seam: wraps an already-connected <paramref name="tcp"/> but routes all frame I/O through
    /// <paramref name="stream"/> (e.g. one whose write never completes), skipping RegisterSession. The TCP
    /// client is still what <see cref="Fault"/> closes and <see cref="IsConnected"/> consults.</summary>
    internal static EncapsulationClient ForStream(TcpClient tcp, Stream stream) => new(tcp, stream);

    /// <summary>Open the TCP connection and register an encapsulation session. <paramref name="timeout"/>
    /// bounds the TCP connect and the RegisterSession exchange separately.</summary>
    public static async Task<EncapsulationClient> ConnectAsync(string host, int port, TimeSpan timeout, CancellationToken ct)
    {
        ValidateTimeout(timeout);
        var tcp = new TcpClient { NoDelay = true };
        try
        {
            using var cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
            cts.CancelAfter(timeout);
            await tcp.ConnectAsync(host, port, cts.Token).ConfigureAwait(false);
        }
        catch (OperationCanceledException) when (!ct.IsCancellationRequested)
        {
            tcp.Dispose();
            throw new AbTransportException(AbStatus.Timeout, $"TCP connect to {host}:{port} timed out.");
        }
        catch (SocketException ex)
        {
            tcp.Dispose();
            throw new AbTransportException(AbStatus.BadConnection, $"TCP connect to {host}:{port} failed: {ex.SocketErrorCode}", ex);
        }
        catch
        {
            tcp.Dispose();
            throw;
        }

        var client = new EncapsulationClient(tcp);
        try
        {
            var (header, _) = await client.ExchangeAsync(connected: false,
                (_, _) => EncapsulationCodec.RegisterSession(), timeout, ct).ConfigureAwait(false);
            client.SessionHandle = header.SessionHandle;
            return client;
        }
        catch
        {
            client.Abort();
            throw;
        }
    }

    /// <summary>Send one unconnected (UCMM) CIP request and return the CIP reply bytes.</summary>
    public async Task<ReadOnlyMemory<byte>> SendRRDataAsync(ReadOnlyMemory<byte> cip, TimeSpan timeout, CancellationToken ct)
    {
        ValidateTimeout(timeout);
        var secs = (ushort)Math.Clamp((int)Math.Ceiling(timeout.TotalSeconds), 1, ushort.MaxValue);
        var (_, reply) = await ExchangeAsync(connected: false,
            (ctx, _) => EncapsulationCodec.SendRRData(SessionHandle, ctx, cip.Span, secs), timeout, ct).ConfigureAwait(false);
        return reply;
    }

    /// <summary>Send one connected (class 3) CIP request over <paramref name="otConnectionId"/> and return
    /// the CIP reply bytes. The client owns the connected sequence count (review B1).</summary>
    public async Task<ReadOnlyMemory<byte>> SendUnitDataAsync(uint otConnectionId, ReadOnlyMemory<byte> cip,
        TimeSpan timeout, CancellationToken ct)
    {
        ValidateTimeout(timeout);
        var (_, reply) = await ExchangeAsync(connected: true,
            (ctx, seq) => EncapsulationCodec.SendUnitData(SessionHandle, ctx, otConnectionId, seq, cip.Span), timeout, ct)
            .ConfigureAwait(false);
        return reply;
    }

    private static void ValidateTimeout(TimeSpan timeout)
    {
        if (timeout <= TimeSpan.Zero || timeout == Timeout.InfiniteTimeSpan)
            throw new ArgumentOutOfRangeException(nameof(timeout), timeout, "Timeout must be finite and positive.");
    }

    /// <summary>The connected sequence count after <paramref name="current"/>: increments, wrapping
    /// 65535 → 1 (0 is never put on the wire).</summary>
    internal static ushort NextSequence(ushort current) => current == ushort.MaxValue ? (ushort)1 : (ushort)(current + 1);

    /// <summary>Acquire the gate, build the request (sequence assigned under the gate), then hand the
    /// gate to <see cref="RunExchangeAsync"/>, which owns its release from that point on.</summary>
    private async Task<(EncapsulationHeader Header, ReadOnlyMemory<byte> Cip)> ExchangeAsync(
        bool connected, Func<ulong, ushort, byte[]> build, TimeSpan timeout, CancellationToken ct)
    {
        ThrowIfUnusable();
        if (!await _gate.WaitAsync(timeout, ct).ConfigureAwait(false))
            throw new AbTransportException(AbStatus.Timeout,
                $"Timed out after {timeout.TotalMilliseconds:F0} ms waiting for the session's single-flight gate.");

        Task<(EncapsulationHeader, ReadOnlyMemory<byte>)> exchange;
        try
        {
            ThrowIfUnusable();                     // a drain we queued behind may have faulted
            ct.ThrowIfCancellationRequested();     // nothing written yet: cancel cleanly
            var ctx = (ulong)Interlocked.Increment(ref _ctx);
            ushort seq = 0;
            if (connected) seq = _seq = NextSequence(_seq);
            var request = build(ctx, seq);
            exchange = RunExchangeAsync(request, connected, ctx, seq, timeout);
        }
        catch
        {
            _gate.Release();                       // RunExchangeAsync never started: release here
            throw;
        }

        try
        {
            return await exchange.WaitAsync(ct).ConfigureAwait(false);
        }
        catch (OperationCanceledException) when (ct.IsCancellationRequested)
        {
            // Caller abandoned the wait (R5-1). The exchange keeps draining its own reply under the
            // gate and faults on its own timeout; observe its outcome so a fault is never unobserved.
            _ = exchange.ContinueWith(static t => _ = t.Exception, CancellationToken.None,
                TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
            throw;
        }
    }

    /// <summary>Write one frame, then read frames until the matching reply arrives. Releases the gate
    /// exactly once, in its <c>finally</c>. Never throws <see cref="OperationCanceledException"/>: every
    /// failure surfaces as <see cref="AbTransportException"/> and faults the connection first.</summary>
    private async Task<(EncapsulationHeader, ReadOnlyMemory<byte>)> RunExchangeAsync(
        byte[] request, bool connected, ulong ctx, ushort seq, TimeSpan timeout)
    {
        var cmd = (ushort)(request[0] | (request[1] << 8));
        try
        {
            using var timeoutCts = new CancellationTokenSource(timeout);
            try
            {
                // The write ignores cancellation (a half-written frame would desync the stream) but is
                // bounded: on expiry the catch below faults the connection, closing the socket, which
                // is what fails the still-pending write.
                var write = _stream.WriteAsync(request, CancellationToken.None).AsTask();
                _ = write.ContinueWith(static t => _ = t.Exception, CancellationToken.None,
                    TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
                await write.WaitAsync(timeoutCts.Token).ConfigureAwait(false);
                while (true)
                {
                    var frame = await ReadFrameAsync(timeoutCts.Token).ConfigureAwait(false);
                    if (TryMatch(frame, cmd, connected, ctx, seq, out var reply)) return reply;
                }
            }
            catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested)
            {
                Fault();
                throw new AbTransportException(AbStatus.Timeout,
                    $"No reply to command 0x{cmd:X4} within {timeout.TotalMilliseconds:F0} ms.");
            }
            catch (AbTransportException)
            {
                Fault();
                throw;
            }
            catch (Exception ex) when (ex is IOException or SocketException or ObjectDisposedException
                                           or EndOfStreamException or InvalidOperationException)
            {
                Fault();
                throw new AbTransportException(AbStatus.BadConnection, $"Connection lost: {ex.Message}", ex);
            }
        }
        finally
        {
            _gate.Release();
        }
    }

    /// <summary>Does <paramref name="frame"/> answer the request in flight? Unconnected replies match on
    /// command and sender context, tolerating a zeroed context (ab_server, R5-2); connected replies match
    /// on the echoed sequence count. Anything else is skipped. A matching frame is fully validated, so an
    /// encapsulation error throws BadReply.</summary>
    private static bool TryMatch(byte[] frame, ushort cmd, bool connected, ulong ctx, ushort seq,
        out (EncapsulationHeader, ReadOnlyMemory<byte>) reply)
    {
        reply = default;
        var h = EncapsulationCodec.ReadHeader(frame);
        if (h.Command != cmd) return false;

        if (!connected)
        {
            // "== 0" is load-bearing, not just ab_server tolerance: RegisterSession is sent with a ZERO
            // sender context (ab_server refuses non-zero), so its reply echoes 0 and matches ONLY here.
            if (h.SenderContext != ctx && h.SenderContext != 0) return false;
            var p = EncapsulationCodec.ParseReply(frame, expectConnected: false);
            reply = (p.Header, p.Cip);
            return true;
        }

        // A connected error reply (status != 0) carries no readable sequence; with one request in flight
        // and abandoned exchanges drained, it can only be ours — ParseReply throws BadReply for it.
        var c = EncapsulationCodec.ParseReply(frame, expectConnected: true);
        if (c.Sequence != seq) return false;
        reply = (c.Header, c.Cip);
        return true;
    }

    private async Task<byte[]> ReadFrameAsync(CancellationToken ct)
    {
        var head = new byte[EncapsulationHeader.Size];
        await _stream.ReadExactlyAsync(head, ct).ConfigureAwait(false);
        var len = head[2] | (head[3] << 8);
        var frame = new byte[EncapsulationHeader.Size + len];
        head.CopyTo(frame, 0);
        await _stream.ReadExactlyAsync(frame.AsMemory(EncapsulationHeader.Size), ct).ConfigureAwait(false);
        return frame;
    }

    private void ThrowIfUnusable()
    {
        if (Volatile.Read(ref _disposed) != 0)
            throw new AbTransportException(AbStatus.BadConnection, "Session is disposed.");
        if (_faulted)
            throw new AbTransportException(AbStatus.BadConnection, "Session is faulted; reconnect.");
    }

    private void Fault()
    {
        _faulted = true;
        try { _tcp.Close(); } catch { /* already closed */ }
    }

    /// <summary>Close without UnRegisterSession (registration never completed).</summary>
    private void Abort()
    {
        Interlocked.Exchange(ref _disposed, 1);
        Fault();
    }

    /// <summary>Best-effort UnRegisterSession, then close. Waits briefly for an in-flight exchange so the
    /// UnRegisterSession frame is never interleaved with another frame; if the gate is not free in time
    /// the socket is closed without it.</summary>
    public async ValueTask DisposeAsync()
    {
        if (Interlocked.Exchange(ref _disposed, 1) != 0) return;
        var acquired = await _gate.WaitAsync(UnregisterBudget).ConfigureAwait(false);
        try
        {
            if (acquired && !_faulted && _tcp.Connected)
            {
                try
                {
                    var bytes = EncapsulationCodec.UnRegisterSession(SessionHandle, (ulong)Interlocked.Increment(ref _ctx));
                    using var cts = new CancellationTokenSource(UnregisterBudget);
                    await _stream.WriteAsync(bytes, cts.Token).ConfigureAwait(false);
                    await _stream.FlushAsync(cts.Token).ConfigureAwait(false);
                }
                catch { /* best effort: the socket is closed next either way */ }
            }
        }
        finally
        {
            Fault();
            _tcp.Dispose();
            if (acquired) _gate.Release();
        }
    }
}