lmxopcua/tests/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Mqtt.Tests/MiniBroker.cs

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

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

/// <summary>
///     The smallest broker that can complete an MQTT 3.1.1 handshake: accept TCP, answer CONNECT with
///     a CONNACK, answer PINGREQ with PINGRESP, and nothing else. That is enough for MQTTnet to report
///     a real connection — which is what makes a real drop (<see cref="DropAll"/>) drive a real
///     <c>DisconnectedAsync</c> and a real reconnect, rather than a test seam standing in for one.
/// </summary>
/// <remarks>
///     Shared across suites rather than nested in one, because both the connection-level and the
///     driver-level rejection tests need it: the driver-level one is the pin for the actual
///     operator-visible symptom (a rejected CONNECT must not report <c>Healthy</c>).
/// </remarks>
internal sealed class MiniBroker : IDisposable
{
    /// <summary>CONNACK "Connection Accepted" (MQTT 3.1.1 return code 0).</summary>
    internal const byte ReturnCodeAccepted = 0x00;

    /// <summary>CONNACK "Identifier rejected" — MQTTnet maps it to <c>ClientIdentifierNotValid</c>.</summary>
    internal const byte ReturnCodeIdentifierRejected = 0x02;

    /// <summary>CONNACK "Server unavailable" — MQTTnet maps it to <c>ServerUnavailable</c> (transient).</summary>
    internal const byte ReturnCodeServerUnavailable = 0x03;

    /// <summary>CONNACK "Bad user name or password" — MQTTnet maps it to <c>BadUserNameOrPassword</c>.</summary>
    internal const byte ReturnCodeBadUserNameOrPassword = 0x04;

    /// <summary>CONNACK "Not authorized" — MQTTnet maps it to <c>NotAuthorized</c> (unrecoverable).</summary>
    internal const byte ReturnCodeNotAuthorized = 0x05;

    private static readonly byte[] PingResp = [0xD0, 0x00];

    private readonly CancellationTokenSource _cts = new();
    private readonly HashSet<TcpClient> _live = [];
    private readonly TcpListener _listener;
    private int _acceptedCount;

    public MiniBroker()
    {
        _listener = new TcpListener(IPAddress.Loopback, 0);
        _listener.Start();
        Port = ((IPEndPoint)_listener.LocalEndpoint).Port;
        _ = Task.Run(AcceptLoopAsync);
    }

    public int Port { get; }

    /// <summary>When set, TCP is still accepted but CONNECT is never answered (frozen peer).</summary>
    public bool Blackhole { get; set; }

    /// <summary>
    ///     The MQTT 3.1.1 return code answered to every CONNECT. Anything but
    ///     <see cref="ReturnCodeAccepted"/> is a broker that completes the handshake and then
    ///     <i>refuses</i> — the shape MQTTnet reports as a result rather than an exception, and the one
    ///     a real Mosquitto with the wrong password produces.
    /// </summary>
    public byte ConnAckReturnCode { get; set; } = ReturnCodeAccepted;

    /// <summary>
    ///     When true (the default) a SUBSCRIBE is answered with a SUBACK granting each filter's requested
    ///     QoS. False is a broker that accepts SUBSCRIBE and never answers — the subscribe-deadline shape.
    /// </summary>
    public bool AnswerSubscribe { get; set; } = true;

    /// <summary>
    ///     Builds the SUBACK for one SUBSCRIBE packet read whole into <paramref name="buffer"/>, or null when
    ///     the packet is not parseable from one read (the fake handles one packet per read, like CONNECT).
    /// </summary>
    private static byte[]? TryBuildSubAck(byte[] buffer, int read)
    {
        // Fixed header: type byte, then the variable-length "remaining length".
        var at = 1;
        var remaining = 0;
        var multiplier = 1;
        byte digit;
        do
        {
            if (at >= read) return null;
            digit = buffer[at++];
            remaining += (digit & 0x7F) * multiplier;
            multiplier *= 128;
        }
        while ((digit & 0x80) != 0);

        var end = at + remaining;
        if (end > read || remaining < 2) return null;

        var packetIdHigh = buffer[at];
        var packetIdLow = buffer[at + 1];
        at += 2;

        var grants = new List<byte>();
        while (at + 2 <= end)
        {
            var length = (buffer[at] << 8) | buffer[at + 1];
            at += 2 + length;
            if (at >= end) return null;
            grants.Add((byte)(buffer[at] & 0x03));
            at++;
        }

        var ack = new List<byte> { 0x90, (byte)(2 + grants.Count), packetIdHigh, packetIdLow };
        ack.AddRange(grants);
        return [.. ack];
    }

    /// <summary>Total TCP connections accepted since construction — counts reconnect attempts.</summary>
    public int AcceptedCount => Volatile.Read(ref _acceptedCount);

    /// <summary>
    ///     Sockets the broker still has open. Drops to zero only once the peer actually closes, so a
    ///     leaked-but-live client keeps this above zero.
    /// </summary>
    public int LiveConnections
    {
        get
        {
            lock (_live)
            {
                return _live.Count;
            }
        }
    }

    /// <summary>Kills every established session — the "broker went away" event.</summary>
    public void DropAll()
    {
        lock (_live)
        {
            foreach (var client in _live)
            {
                try
                {
                    client.Close();
                }
                catch (Exception)
                {
                    // Already gone.
                }
            }

            _live.Clear();
        }
    }

    public void Dispose()
    {
        _cts.Cancel();
        _listener.Stop();
        DropAll();
        _cts.Dispose();
    }

    private async Task AcceptLoopAsync()
    {
        try
        {
            while (!_cts.IsCancellationRequested)
            {
                var client = await _listener.AcceptTcpClientAsync(_cts.Token);
                Interlocked.Increment(ref _acceptedCount);
                lock (_live)
                {
                    _live.Add(client);
                }

                _ = Task.Run(() => ServeAsync(client));
            }
        }
        catch (Exception)
        {
            // Listener stopped / token cancelled — the fixture is going away.
        }
    }

    private async Task ServeAsync(TcpClient client)
    {
        var buffer = new byte[1024];
        try
        {
            var stream = client.GetStream();
            while (!_cts.IsCancellationRequested)
            {
                var read = await stream.ReadAsync(buffer, _cts.Token);
                if (read == 0)
                {
                    break; // peer closed — this is how a disposed client stops counting as live
                }

                if (Blackhole)
                {
                    continue;
                }

                var packetType = buffer[0] & 0xF0;
                if (packetType == 0x10)
                {
                    await stream.WriteAsync(new byte[] { 0x20, 0x02, 0x00, ConnAckReturnCode }, _cts.Token);

                    // A real broker closes the socket straight after refusing, which is what drives
                    // MQTTnet's DisconnectedAsync. Mirroring it keeps the fake honest.
                    if (ConnAckReturnCode != ReturnCodeAccepted)
                    {
                        break;
                    }
                }
                else if (packetType == 0xC0)
                {
                    await stream.WriteAsync(PingResp, _cts.Token);
                }
                else if (packetType == 0x80 && AnswerSubscribe && TryBuildSubAck(buffer, read) is { } subAck)
                {
                    await stream.WriteAsync(subAck, _cts.Token);
                }
            }
        }
        catch (Exception)
        {
            // Socket torn down by either side.
        }
        finally
        {
            lock (_live)
            {
                _live.Remove(client);
            }

            try
            {
                client.Dispose();
            }
            catch (Exception)
            {
                // Already gone.
            }
        }
    }
}