lmxopcua/tests/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Modbus.IntegrationTests/TcpRelay.cs

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

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

/// <summary>
///     Loopback TCP relay in front of the real Modbus fixture, so a test can make the device "disappear"
///     (<see cref="Cut"/>: every relayed connection is closed and new connects are refused) and "come back"
///     (<see cref="Restore"/>: listen again on the SAME port) without touching the shared container.
/// </summary>
internal sealed class TcpRelay : IAsyncDisposable
{
    private readonly string _upstreamHost;
    private readonly int _upstreamPort;
    private readonly object _lock = new();
    private readonly List<TcpClient> _live = [];
    private TcpListener? _listener;
    private CancellationTokenSource _stop = new();
    private int _accepted;

    public TcpRelay(string upstreamHost, int upstreamPort)
    {
        _upstreamHost = upstreamHost;
        _upstreamPort = upstreamPort;
        var probe = new TcpListener(IPAddress.Loopback, 0);
        probe.Start();
        Port = ((IPEndPoint)probe.LocalEndpoint).Port;
        probe.Stop();
        Restore();
    }

    /// <summary>Gets the loopback port clients dial.</summary>
    public int Port { get; }

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

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

    /// <summary>Closes every relayed connection and refuses new ones until <see cref="Restore"/>.</summary>
    public void Cut()
    {
        try { _stop.Cancel(); } catch { }
        try { _listener?.Stop(); } catch { }
        _listener = null;
        lock (_lock)
        {
            foreach (var c in _live) { try { c.Close(); } catch { } }
            _live.Clear();
        }
    }

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

    private async Task RelayAsync(TcpClient downstream, CancellationToken token)
    {
        var upstream = new TcpClient(AddressFamily.InterNetwork);
        lock (_lock) { _live.Add(downstream); _live.Add(upstream); }
        try
        {
            var addr = (await Dns.GetHostAddressesAsync(_upstreamHost, token))
                .First(a => a.AddressFamily == AddressFamily.InterNetwork);
            await upstream.ConnectAsync(addr, _upstreamPort, token);
            var a = downstream.GetStream();
            var b = upstream.GetStream();
            await Task.WhenAny(a.CopyToAsync(b, token), b.CopyToAsync(a, token));
        }
        catch { /* cut or upstream gone */ }
        finally
        {
            try { downstream.Close(); } catch { }
            try { upstream.Close(); } catch { }
        }
    }

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