lmxopcua/tests/Server/ZB.MOM.WW.OtOpcUa.OpcUaServer.IntegrationTests/TcpStallProxy.cs

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

namespace ZB.MOM.WW.OtOpcUa.OpcUaServer.IntegrationTests;

/// <summary>
///     Gitea #676 — an in-test TCP forwarder that can STALL a connection without closing it: the shape of a
///     <c>docker pause</c>, a hung switch port or a congested WAN link, where the socket stays open and simply stops
///     carrying bytes. Each accepted connection gets a fresh upstream connection and two copy pumps; every pump awaits a
///     shared gate before each write, so <see cref="Pause"/> holds all traffic in both directions (bytes already read
///     are held, the rest back up in the kernel buffers) and <see cref="Resume"/> lets it all through in order. Nothing
///     is ever closed by a pause — the outage is invisible at the TCP level, exactly as in production.
/// </summary>
internal sealed class TcpStallProxy : IAsyncDisposable
{
    private readonly TcpListener _listener = new(IPAddress.Loopback, 0);
    private readonly int _targetPort;
    private readonly CancellationTokenSource _cts = new();
    private readonly List<TcpClient> _sockets = [];
    private readonly Lock _lock = new();
    private TaskCompletionSource _gate = Open();
    private volatile bool _refusing;
    private Task _acceptLoop = Task.CompletedTask;

    private TcpStallProxy(int targetPort) => _targetPort = targetPort;

    /// <summary>The loopback port clients connect to.</summary>
    public int Port => ((IPEndPoint)_listener.LocalEndpoint).Port;

    /// <summary>Number of client connections accepted so far — a new one means the client re-connected.</summary>
    public int AcceptedConnections { get; private set; }

    /// <summary>Starts forwarding <c>127.0.0.1:{Port}</c> to <c>127.0.0.1:{targetPort}</c>.</summary>
    public static TcpStallProxy Start(int targetPort)
    {
        var proxy = new TcpStallProxy(targetPort);
        proxy._listener.Start();
        proxy._acceptLoop = proxy.AcceptLoopAsync();
        return proxy;
    }

    /// <summary>Rewrites the port of an <c>opc.tcp</c> URL to this proxy's.</summary>
    public string Through(string endpointUrl) => new UriBuilder(endpointUrl) { Port = Port }.Uri.ToString();

    /// <summary>Stops all traffic in both directions; sockets stay open.</summary>
    public void Pause()
    {
        lock (_lock)
            if (_gate.Task.IsCompleted) _gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
    }

    /// <summary>
    ///     The shape of a killed / stopped server process (#676 live gate): every open connection is RESET (linger 0 ⇒
    ///     RST, both sides) and new connections are refused (accepted and reset at once) until <see cref="Reopen"/>.
    /// </summary>
    public void HardClose()
    {
        List<TcpClient> victims;
        lock (_lock)
        {
            _refusing = true;
            victims = [.. _sockets];
            _sockets.Clear();
        }
        foreach (var s in victims) Reset(s);
    }

    /// <summary>Accept and forward new connections again (the server is back); the reset ones stay dead.</summary>
    public void Reopen() => _refusing = false;

    /// <summary>Connections refused while <see cref="HardClose"/> was in force — a reconnect ATTEMPT by the client.</summary>
    public int RefusedConnections { get; private set; }

    private static void Reset(TcpClient s)
    {
        try { s.Client.LingerState = new LingerOption(true, 0); } catch { }
        try { s.Dispose(); } catch { }
    }

    /// <summary>Releases everything held since <see cref="Pause"/>, in order.</summary>
    public void Resume()
    {
        lock (_lock) _gate.TrySetResult();
    }

    private static TaskCompletionSource Open()
    {
        var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        tcs.SetResult();
        return tcs;
    }

    private Task CurrentGate()
    {
        lock (_lock) return _gate.Task;
    }

    private async Task AcceptLoopAsync()
    {
        while (!_cts.IsCancellationRequested)
        {
            TcpClient client;
            try { client = await _listener.AcceptTcpClientAsync(_cts.Token); }
            catch { return; }
            if (_refusing)
            {
                lock (_lock) RefusedConnections++;
                Reset(client);
                continue;
            }
            var upstream = new TcpClient();
            try { await upstream.ConnectAsync(IPAddress.Loopback, _targetPort, _cts.Token); }
            catch
            {
                client.Dispose();
                upstream.Dispose();
                continue;
            }
            lock (_lock)
            {
                AcceptedConnections++;
                _sockets.Add(client);
                _sockets.Add(upstream);
            }
            _ = PumpAsync(client.GetStream(), upstream.GetStream());
            _ = PumpAsync(upstream.GetStream(), client.GetStream());
        }
    }

    private async Task PumpAsync(NetworkStream from, NetworkStream to)
    {
        var buffer = new byte[64 * 1024];
        try
        {
            while (true)
            {
                var n = await from.ReadAsync(buffer, _cts.Token);
                if (n == 0) break;
                await CurrentGate().WaitAsync(_cts.Token);
                await to.WriteAsync(buffer.AsMemory(0, n), _cts.Token);
            }
        }
        catch { /* either side gone or the proxy disposed */ }
        // Half-close propagation is not needed: a real close (server stopped) tears the pair down on the next read.
        try { to.Socket.Shutdown(SocketShutdown.Send); } catch { }
    }

    public async ValueTask DisposeAsync()
    {
        Resume();
        await _cts.CancelAsync();
        _listener.Stop();
        try { await _acceptLoop; } catch { }
        lock (_lock)
        {
            foreach (var s in _sockets) s.Dispose();
            _sockets.Clear();
        }
        _cts.Dispose();
    }
}