natsdotnet/tests/NATS.Server.Mqtt.Tests/Mqtt/MqttSessionParityTests.cs

// Ports session management behavior from Go reference:
// golang/nats-server/server/mqtt_test.go — TestMQTTCleanSession, TestMQTTPersistedSession,
// TestMQTTDuplicateClientID, TestMQTTRecoverSessionAndAddNewSub

using System.Net;
using System.Net.Sockets;
using System.Text;
using NATS.Server.Mqtt;

namespace NATS.Server.Mqtt.Tests.Mqtt;

public class MqttSessionParityTests
{
    // Go ref: TestMQTTCleanSession — connecting with clean=true discards any previous session state.
    // A clean-session client never receives redeliveries from prior disconnected sessions.
    [Fact]
    public async Task Clean_session_true_discards_previous_session_state()
    {
        await using var listener = new MqttListener("127.0.0.1", 0);
        listener.UseBinaryProtocol = false;
        using var cts = new CancellationTokenSource();
        await listener.StartAsync(cts.Token);

        // First connection: send a QoS 1 publish that goes unacked (session-client, persistent)
        using (var first = new TcpClient())
        {
            await first.ConnectAsync(IPAddress.Loopback, listener.Port);
            var s = first.GetStream();
            await MqttSessionWire.WriteLineAsync(s, "CONNECT clean-test-client clean=false");
            (await MqttSessionWire.ReadLineAsync(s, 1000)).ShouldBe("CONNACK");

            // Publish QoS 1 — server records pending, client disconnects without ACKing
            await MqttSessionWire.WriteLineAsync(s, "PUBQ1 5 device.status online");
            (await MqttSessionWire.ReadLineAsync(s, 1000)).ShouldBe("PUBACK 5");
        }

        // Second connection with clean=true — session state must be purged, no REDLIVER
        using var second = new TcpClient();
        await second.ConnectAsync(IPAddress.Loopback, listener.Port);
        var stream = second.GetStream();
        await MqttSessionWire.WriteLineAsync(stream, "CONNECT clean-test-client clean=true");
        (await MqttSessionWire.ReadLineAsync(stream, 1000)).ShouldBe("CONNACK");

        // No redelivery expected because clean session wiped state
        (await MqttSessionWire.ReadLineAsync(stream, 300)).ShouldBeNull();
    }

    // Go ref: TestMQTTPersistedSession — clean=false preserves unacked QoS 1 publishes across reconnect.
    [Fact]
    public async Task Clean_session_false_preserves_unacked_publishes_across_reconnect()
    {
        await using var listener = new MqttListener("127.0.0.1", 0);
        listener.UseBinaryProtocol = false;
        using var cts = new CancellationTokenSource();
        await listener.StartAsync(cts.Token);

        // First connection: publish QoS 1 without sending ACK, then drop
        using (var first = new TcpClient())
        {
            await first.ConnectAsync(IPAddress.Loopback, listener.Port);
            var s = first.GetStream();
            await MqttSessionWire.WriteLineAsync(s, "CONNECT persist-client clean=false");
            (await MqttSessionWire.ReadLineAsync(s, 1000)).ShouldBe("CONNACK");

            await MqttSessionWire.WriteLineAsync(s, "PUBQ1 12 alarm.fire detected");
            (await MqttSessionWire.ReadLineAsync(s, 1000)).ShouldBe("PUBACK 12");
            // Disconnect without sending ACK 12
        }

        // Second connection with same clientId, clean=false — server must redeliver
        using var second = new TcpClient();
        await second.ConnectAsync(IPAddress.Loopback, listener.Port);
        var stream = second.GetStream();
        await MqttSessionWire.WriteLineAsync(stream, "CONNECT persist-client clean=false");
        (await MqttSessionWire.ReadLineAsync(stream, 1000)).ShouldBe("CONNACK");
        (await MqttSessionWire.ReadLineAsync(stream, 1000)).ShouldBe("REDLIVER 12 alarm.fire detected");
    }

    // Go ref: TestMQTTCleanSession — after clean disconnect the session entry is removed;
    // a subsequent persistent reconnect starts fresh with no pending messages.
    [Fact]
    public async Task Session_disconnect_cleans_up_client_tracking_on_clean_session()
    {
        await using var listener = new MqttListener("127.0.0.1", 0);
        listener.UseBinaryProtocol = false;
        using var cts = new CancellationTokenSource();
        await listener.StartAsync(cts.Token);

        // Connect and immediately disconnect without publishing anything (clean=true)
        using (var first = new TcpClient())
        {
            await first.ConnectAsync(IPAddress.Loopback, listener.Port);
            var s = first.GetStream();
            await MqttSessionWire.WriteLineAsync(s, "CONNECT transient-client clean=true");
            (await MqttSessionWire.ReadLineAsync(s, 1000)).ShouldBe("CONNACK");
        }

        // Reconnect with clean=false — no session was saved, so no redeliveries
        using var second = new TcpClient();
        await second.ConnectAsync(IPAddress.Loopback, listener.Port);
        var stream = second.GetStream();
        await MqttSessionWire.WriteLineAsync(stream, "CONNECT transient-client clean=false");
        (await MqttSessionWire.ReadLineAsync(stream, 1000)).ShouldBe("CONNACK");

        // Nothing pending from the previous clean-session connection
        (await MqttSessionWire.ReadLineAsync(stream, 300)).ShouldBeNull();
    }

    // Go ref: TestMQTTDuplicateClientID — multiple concurrent sessions on distinct client IDs
    // operate independently with no cross-contamination of messages or session state.
    [Fact]
    public async Task Multiple_concurrent_sessions_on_different_client_ids_work_independently()
    {
        await using var listener = new MqttListener("127.0.0.1", 0);
        listener.UseBinaryProtocol = false;
        using var cts = new CancellationTokenSource();
        await listener.StartAsync(cts.Token);

        // Client A — persistent session, QoS 1 publish unacked
        using var clientA = new TcpClient();
        await clientA.ConnectAsync(IPAddress.Loopback, listener.Port);
        var streamA = clientA.GetStream();
        await MqttSessionWire.WriteLineAsync(streamA, "CONNECT client-alpha clean=false");
        (await MqttSessionWire.ReadLineAsync(streamA, 1000)).ShouldBe("CONNACK");
        await MqttSessionWire.WriteLineAsync(streamA, "PUBQ1 7 alpha.topic alpha-payload");
        (await MqttSessionWire.ReadLineAsync(streamA, 1000)).ShouldBe("PUBACK 7");

        // Client B — independent persistent session, different topic and packet ID
        using var clientB = new TcpClient();
        await clientB.ConnectAsync(IPAddress.Loopback, listener.Port);
        var streamB = clientB.GetStream();
        await MqttSessionWire.WriteLineAsync(streamB, "CONNECT client-beta clean=false");
        (await MqttSessionWire.ReadLineAsync(streamB, 1000)).ShouldBe("CONNACK");
        await MqttSessionWire.WriteLineAsync(streamB, "PUBQ1 8 beta.topic beta-payload");
        (await MqttSessionWire.ReadLineAsync(streamB, 1000)).ShouldBe("PUBACK 8");

        // Disconnect both without ACKing
        clientA.Dispose();
        clientB.Dispose();

        // Reconnect alpha — must only redeliver alpha's pending publish
        using var reconnectA = new TcpClient();
        await reconnectA.ConnectAsync(IPAddress.Loopback, listener.Port);
        var rsA = reconnectA.GetStream();
        await MqttSessionWire.WriteLineAsync(rsA, "CONNECT client-alpha clean=false");
        (await MqttSessionWire.ReadLineAsync(rsA, 1000)).ShouldBe("CONNACK");
        (await MqttSessionWire.ReadLineAsync(rsA, 1000)).ShouldBe("REDLIVER 7 alpha.topic alpha-payload");

        // Reconnect beta — must only redeliver beta's pending publish
        using var reconnectB = new TcpClient();
        await reconnectB.ConnectAsync(IPAddress.Loopback, listener.Port);
        var rsB = reconnectB.GetStream();
        await MqttSessionWire.WriteLineAsync(rsB, "CONNECT client-beta clean=false");
        (await MqttSessionWire.ReadLineAsync(rsB, 1000)).ShouldBe("CONNACK");
        (await MqttSessionWire.ReadLineAsync(rsB, 1000)).ShouldBe("REDLIVER 8 beta.topic beta-payload");

        // Alpha should not see beta's message and vice-versa (no cross-contamination)
        (await MqttSessionWire.ReadLineAsync(rsA, 200)).ShouldBeNull();
        (await MqttSessionWire.ReadLineAsync(rsB, 200)).ShouldBeNull();
    }
}

// Duplicated per-file as required — each test file is self-contained.
internal static class MqttSessionWire
{
    public static async Task WriteLineAsync(NetworkStream stream, string line)
    {
        var bytes = Encoding.UTF8.GetBytes(line + "\n");
        await stream.WriteAsync(bytes);
        await stream.FlushAsync();
    }

    public static async Task<string?> ReadLineAsync(NetworkStream stream, int timeoutMs)
    {
        using var timeout = new CancellationTokenSource(timeoutMs);
        var bytes = new List<byte>();
        var one = new byte[1];
        try
        {
            while (true)
            {
                var read = await stream.ReadAsync(one.AsMemory(0, 1), timeout.Token);
                if (read == 0)
                    return null;
                if (one[0] == (byte)'\n')
                    break;
                if (one[0] != (byte)'\r')
                    bytes.Add(one[0]);
            }
        }
        catch (OperationCanceledException)
        {
            return null;
        }

        return Encoding.UTF8.GetString([.. bytes]);
    }

    public static async Task<string?> ReadRawAsync(NetworkStream stream, int timeoutMs)
    {
        using var timeout = new CancellationTokenSource(timeoutMs);
        var one = new byte[1];
        try
        {
            var read = await stream.ReadAsync(one.AsMemory(0, 1), timeout.Token);
            if (read == 0)
                return null;

            return Encoding.UTF8.GetString(one, 0, read);
        }
        catch (OperationCanceledException)
        {
            return "__timeout__";
        }
    }
}