natsnet/dotnet/tests/ZB.MOM.NatsNet.Server.Tests/Mqtt/MqttParserTests.cs

// Copyright 2020-2026 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0

using Shouldly;
using ZB.MOM.NatsNet.Server;
using ZB.MOM.NatsNet.Server.Internal;
using ZB.MOM.NatsNet.Server.Mqtt;

namespace ZB.MOM.NatsNet.Server.Tests.Mqtt;

/// <summary>
/// Unit tests for <see cref="MqttParser"/> — validates packet type extraction,
/// remaining length decoding, CONNECT-first enforcement, partial packet handling,
/// and PINGREQ dispatch.
/// </summary>
public sealed class MqttParserTests
{
    /// <summary>
    /// Creates a minimal ClientConnection with MQTT handler for testing.
    /// </summary>
    private static ClientConnection CreateMqttClient()
    {
        var c = new ClientConnection(ClientKind.Client);
        c.InitMqtt(new MqttHandler());
        return c;
    }

    // =========================================================================
    // CONNECT-first enforcement
    // =========================================================================

    [Fact]
    public void Parse_NonConnectFirst_ShouldReturnError()
    {
        var c = CreateMqttClient();
        // PINGREQ before CONNECT → error
        var buf = new byte[] { MqttPacket.Ping, 0x00 };
        var err = MqttParser.Parse(c, buf, buf.Length);
        err.ShouldNotBeNull();
        err.Message.ShouldContain("first packet should be a CONNECT");
    }

    [Fact]
    public void Parse_PublishBeforeConnect_ShouldReturnError()
    {
        var c = CreateMqttClient();
        // PUBLISH QoS 0 before CONNECT
        var buf = new byte[] { MqttPacket.Pub, 0x05, 0x00, 0x01, (byte)'t', 0x68, 0x69 };
        var err = MqttParser.Parse(c, buf, buf.Length);
        err.ShouldNotBeNull();
        err.Message.ShouldContain("first packet should be a CONNECT");
    }

    [Fact]
    public void Parse_ConnectFirst_ShouldAcceptConnect()
    {
        var c = CreateMqttClient();
        // Use a MemoryStream so CONNACK can be written.
        typeof(ClientConnection)
            .GetField("_nc", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance)!
            .SetValue(c, new MemoryStream());

        // Build a valid CONNECT packet.
        var payload = new List<byte>();
        payload.AddRange(new byte[] { 0x00, 0x04 }); // protocol name length
        payload.AddRange(System.Text.Encoding.UTF8.GetBytes("MQTT"));
        payload.Add(0x04); // level
        payload.Add(0x02); // flags: clean session
        payload.AddRange(new byte[] { 0x00, 0x3C }); // keep alive = 60
        payload.AddRange(new byte[] { 0x00, 0x04 }); // client id length
        payload.AddRange(System.Text.Encoding.UTF8.GetBytes("test"));

        var buf = new List<byte> { MqttPacket.Connect };
        // Remaining length
        var remLen = payload.Count;
        do
        {
            var b = (byte)(remLen & 0x7F);
            remLen >>= 7;
            if (remLen > 0) b |= 0x80;
            buf.Add(b);
        } while (remLen > 0);
        buf.AddRange(payload);

        var err = MqttParser.Parse(c, buf.ToArray(), buf.Count);
        err.ShouldBeNull("CONNECT should be accepted, not rejected as non-CONNECT");
        (c.Flags & ClientFlags.ConnectReceived).ShouldNotBe((ClientFlags)0);
    }

    [Fact]
    public void Parse_SecondConnect_ShouldReturnError()
    {
        var c = CreateMqttClient();
        // Simulate that CONNECT was already received.
        c.Flags |= ClientFlags.ConnectReceived;
        var buf = new byte[] { MqttPacket.Connect, 0x00 };
        var err = MqttParser.Parse(c, buf, buf.Length);
        err.ShouldNotBeNull();
        err.Message.ShouldContain("second CONNECT");
    }

    // =========================================================================
    // PINGREQ dispatch
    // =========================================================================

    [Fact]
    public void Parse_PingReq_ShouldEnqueuePingResp()
    {
        var c = CreateMqttClient();
        // Mark as connected so ping is accepted.
        c.Flags |= ClientFlags.ConnectReceived;

        // Use a memory stream to capture EnqueueProto output.
        var ms = new MemoryStream();
        // Set up the connection's network stream.
        typeof(ClientConnection)
            .GetField("_nc", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance)!
            .SetValue(c, ms);

        var buf = new byte[] { MqttPacket.Ping, 0x00 };
        var err = MqttParser.Parse(c, buf, buf.Length);
        err.ShouldBeNull();

        // Verify PINGRESP was written.
        var written = ms.ToArray();
        written.Length.ShouldBe(2);
        written[0].ShouldBe(MqttPacket.PingResp);
        written[1].ShouldBe((byte)0x00);
    }

    [Fact]
    public void Parse_MultiplePings_ShouldSucceed()
    {
        var c = CreateMqttClient();
        c.Flags |= ClientFlags.ConnectReceived;

        var ms = new MemoryStream();
        typeof(ClientConnection)
            .GetField("_nc", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance)!
            .SetValue(c, ms);

        // Two PINGREQ packets in the same buffer.
        var buf = new byte[] { MqttPacket.Ping, 0x00, MqttPacket.Ping, 0x00 };
        var err = MqttParser.Parse(c, buf, buf.Length);
        err.ShouldBeNull();

        // Two PINGRESP packets should have been written.
        var written = ms.ToArray();
        written.Length.ShouldBe(4);
        written[0].ShouldBe(MqttPacket.PingResp);
        written[1].ShouldBe((byte)0x00);
        written[2].ShouldBe(MqttPacket.PingResp);
        written[3].ShouldBe((byte)0x00);
    }

    // =========================================================================
    // Remaining length decoding
    // =========================================================================

    [Fact]
    public void Parse_SingleByteRemainingLength_ShouldWork()
    {
        // SUB packet with remaining length = 5 (single byte < 128)
        // After CONNECT is received, a SUB packet should parse the remaining length correctly.
        var c = CreateMqttClient();
        c.Flags |= ClientFlags.ConnectReceived;

        // SUBSCRIBE: type=0x82, remaining len=5, then 5 bytes of payload
        var buf = new byte[] { 0x82, 0x05, 0x00, 0x01, 0x00, 0x01, 0x74 };
        var err = MqttParser.Parse(c, buf, buf.Length);
        // Will hit NotImplementedException for SUBSCRIBE — that's fine, it proves parsing worked.
        err.ShouldNotBeNull();
        err.ShouldBeOfType<NotImplementedException>();
    }

    [Fact]
    public void Parse_TwoByteRemainingLength_ShouldWork()
    {
        var c = CreateMqttClient();
        c.Flags |= ClientFlags.ConnectReceived;

        // Remaining length = 200 → encoded as [0xC8, 0x01]
        // (200 & 0x7F) | 0x80 = 0xC8, 200 >> 7 = 1 → 0x01
        // type(1) + remlen(2) + payload(200) = 203 bytes total.
        var buf = new byte[203];
        buf[0] = MqttPacket.Pub;
        buf[1] = 0xC8;
        buf[2] = 0x01;
        // Remaining 200 bytes are zero (payload).

        var err = MqttParser.Parse(c, buf, buf.Length);
        err.ShouldNotBeNull();
        err.ShouldBeOfType<NotImplementedException>(); // PUBLISH not yet implemented
    }

    // =========================================================================
    // Partial packet handling
    // =========================================================================

    [Fact]
    public void Parse_PartialPacket_ShouldSaveState()
    {
        var c = CreateMqttClient();
        c.Flags |= ClientFlags.ConnectReceived;

        // Send only the first byte of a PING packet (missing the 0x00 remaining length).
        var buf = new byte[] { MqttPacket.Ping };
        var err = MqttParser.Parse(c, buf, buf.Length);
        // Should succeed (partial packet saved) — no error because it just stops.
        // Actually, the ReadByte for packet type succeeds, then ReadPacketLen has no data,
        // so it returns (0, false) — incomplete, state saved.
        err.ShouldBeNull();

        // Now send the remaining length byte.
        var buf2 = new byte[] { 0x00 };
        // The reader's pending buffer should have the first byte saved.
        var ms = new MemoryStream();
        typeof(ClientConnection)
            .GetField("_nc", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance)!
            .SetValue(c, ms);

        err = MqttParser.Parse(c, buf2, buf2.Length);
        err.ShouldBeNull();

        // PINGRESP should have been written.
        var written = ms.ToArray();
        written.Length.ShouldBe(2);
        written[0].ShouldBe(MqttPacket.PingResp);
    }

    [Fact]
    public void Parse_PartialPayload_ShouldSaveState()
    {
        var c = CreateMqttClient();
        c.Flags |= ClientFlags.ConnectReceived;

        // SUBSCRIBE packet: type=0x82, remaining length=10, but only send 5 payload bytes.
        var buf = new byte[] { 0x82, 0x0A, 0x01, 0x02, 0x03, 0x04, 0x05 };
        // remaining length = 10, but only 5 bytes of payload present → partial.
        var err = MqttParser.Parse(c, buf, buf.Length);
        // Should save state and return null (needs more data).
        err.ShouldBeNull();
    }

    // =========================================================================
    // Unknown packet type
    // =========================================================================

    [Fact]
    public void Parse_UnknownPacketType_ShouldReturnError()
    {
        var c = CreateMqttClient();
        c.Flags |= ClientFlags.ConnectReceived;

        // Packet type 0x00 is reserved/invalid.
        var buf = new byte[] { 0x00, 0x00 };
        var err = MqttParser.Parse(c, buf, buf.Length);
        err.ShouldNotBeNull();
        err.Message.ShouldContain("unknown MQTT packet type");
    }

    // =========================================================================
    // Empty buffer
    // =========================================================================

    [Fact]
    public void Parse_EmptyBuffer_ShouldSucceed()
    {
        var c = CreateMqttClient();
        var buf = Array.Empty<byte>();
        var err = MqttParser.Parse(c, buf, 0);
        err.ShouldBeNull(); // Nothing to parse.
    }

    // =========================================================================
    // Buffer length parameter
    // =========================================================================

    [Fact]
    public void Parse_LenSmallerThanBuffer_ShouldOnlyParseLenBytes()
    {
        var c = CreateMqttClient();
        c.Flags |= ClientFlags.ConnectReceived;

        var ms = new MemoryStream();
        typeof(ClientConnection)
            .GetField("_nc", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance)!
            .SetValue(c, ms);

        // Buffer has two PING packets, but len says only the first one.
        var buf = new byte[] { MqttPacket.Ping, 0x00, MqttPacket.Ping, 0x00 };
        var err = MqttParser.Parse(c, buf, 2); // Only first 2 bytes.
        err.ShouldBeNull();

        var written = ms.ToArray();
        written.Length.ShouldBe(2); // Only one PINGRESP.
    }
}