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, SUBSCRIBE with a SUBACK (unless told not to), and record
///     PUBLISHes (PUBACKing QoS 1 only when told to). It never forwards a PUBLISH to a subscriber. 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>
    ///     When true, a QoS 1 PUBLISH is answered with a PUBACK. False (the default) is a broker that takes
    ///     the PUBLISH and never acknowledges it — the stalled-broker shape a write deadline exists for.
    /// </summary>
    public bool AckPublish { get; set; }

    /// <summary>Every PUBLISH received, in arrival order, as <c>(topic, payload)</c>.</summary>
    public System.Collections.Concurrent.ConcurrentQueue<(string Topic, byte[] Payload)> Published { get; } = new();

    /// <summary>Handles one complete packet; returns false when the connection must close.</summary>
    private async Task<bool> HandlePacketAsync(NetworkStream stream, byte[] packet)
    {
        var packetType = packet[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.
            return ConnAckReturnCode == ReturnCodeAccepted;
        }

        if (packetType == 0xC0)
        {
            await stream.WriteAsync(PingResp, _cts.Token);
        }
        else if (packetType == 0x80 && AnswerSubscribe && TryBuildSubAck(packet, packet.Length) is { } subAck)
        {
            await stream.WriteAsync(subAck, _cts.Token);
        }
        else if (packetType == 0x30)
        {
            var qos = (packet[0] >> 1) & 0x03;
            var at = HeaderLength(packet);
            var topicLength = (packet[at] << 8) | packet[at + 1];
            var topic = System.Text.Encoding.UTF8.GetString(packet, at + 2, topicLength);
            at += 2 + topicLength;
            byte idHigh = 0, idLow = 0;
            if (qos > 0)
            {
                idHigh = packet[at];
                idLow = packet[at + 1];
                at += 2;
            }

            Published.Enqueue((topic, packet[at..]));
            if (qos == 1 && AckPublish)
            {
                await stream.WriteAsync(new byte[] { 0x40, 0x02, idHigh, idLow }, _cts.Token);
            }
        }

        return true;
    }

    /// <summary>Length of a packet's fixed header (type byte + variable-length remaining length).</summary>
    private static int HeaderLength(byte[] packet)
    {
        var at = 1;
        while ((packet[at] & 0x80) != 0) at++;
        return at + 1;
    }

    /// <summary>Removes one complete packet from the front of <paramref name="pending"/>, if one is there.</summary>
    private static bool TryTakePacket(List<byte> pending, out byte[] packet)
    {
        packet = [];
        if (pending.Count < 2) return false;

        var at = 1;
        var remaining = 0;
        var multiplier = 1;
        byte digit;
        do
        {
            if (at >= pending.Count) return false;
            digit = pending[at++];
            remaining += (digit & 0x7F) * multiplier;
            multiplier *= 128;
        }
        while ((digit & 0x80) != 0);

        var total = at + remaining;
        if (pending.Count < total) return false;

        packet = pending.GetRange(0, total).ToArray();
        pending.RemoveRange(0, total);
        return true;
    }

    /// <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[4096];
        var pending = new List<byte>();
        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;
                }

                // Frame properly: several packets can share one TCP read (concurrent publishes do), and
                // one packet can span two.
                pending.AddRange(buffer.AsSpan(0, read).ToArray());
                var close = false;
                while (TryTakePacket(pending, out var packet))
                {
                    if (!await HandlePacketAsync(stream, packet))
                    {
                        close = true;
                        break;
                    }
                }

                if (close)
                {
                    break;
                }
            }
        }
        catch (Exception)
        {
            // Socket torn down by either side.
        }
        finally
        {
            lock (_live)
            {
                _live.Remove(client);
            }

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