natsdotnet/src/NATS.Server/Mqtt/MqttPacketWriter.cs

using System.Buffers.Binary;
using System.Text;

namespace NATS.Server.Mqtt;

public static class MqttPacketWriter
{
    public static byte[] WriteString(string value)
        => WriteBytes(Encoding.UTF8.GetBytes(value));

    public static byte[] WriteBytes(ReadOnlySpan<byte> bytes)
    {
        if (bytes.Length > ushort.MaxValue)
            throw new ArgumentOutOfRangeException(nameof(bytes), "MQTT length-prefixed field cannot exceed 65535 bytes.");

        var buffer = new byte[2 + bytes.Length];
        BinaryPrimitives.WriteUInt16BigEndian(buffer.AsSpan(0, 2), (ushort)bytes.Length);
        bytes.CopyTo(buffer.AsSpan(2));
        return buffer;
    }

    public static byte[] Write(MqttControlPacketType type, ReadOnlySpan<byte> payload, byte flags = 0)
    {
        if (type == MqttControlPacketType.Reserved)
            throw new ArgumentOutOfRangeException(nameof(type), "MQTT control packet type must be non-zero.");

        var remainingLength = payload.Length;
        var encodedRemainingLength = EncodeRemainingLength(remainingLength);
        var buffer = new byte[1 + encodedRemainingLength.Length + remainingLength];
        buffer[0] = (byte)(((byte)type << 4) | (flags & 0x0F));
        encodedRemainingLength.CopyTo(buffer.AsSpan(1));
        payload.CopyTo(buffer.AsSpan(1 + encodedRemainingLength.Length));
        return buffer;
    }

    /// <summary>
    /// Writes a CONNACK packet. Go reference: mqtt.go mqttConnAck.
    /// </summary>
    /// <param name="sessionPresent">0x01 if resuming existing session, 0x00 otherwise.</param>
    /// <param name="returnCode">CONNACK return code (0x00 = accepted).</param>
    public static byte[] WriteConnAck(byte sessionPresent, byte returnCode)
    {
        ReadOnlySpan<byte> payload = [sessionPresent, returnCode];
        return Write(MqttControlPacketType.ConnAck, payload);
    }

    /// <summary>
    /// Writes a PUBACK packet (QoS 1 acknowledgment).
    /// </summary>
    public static byte[] WritePubAck(ushort packetId)
    {
        Span<byte> payload = stackalloc byte[2];
        BinaryPrimitives.WriteUInt16BigEndian(payload, packetId);
        return Write(MqttControlPacketType.PubAck, payload);
    }

    /// <summary>
    /// Writes a SUBACK packet with granted QoS values per subscription filter.
    /// </summary>
    public static byte[] WriteSubAck(ushort packetId, ReadOnlySpan<byte> grantedQoS)
    {
        var payload = new byte[2 + grantedQoS.Length];
        BinaryPrimitives.WriteUInt16BigEndian(payload.AsSpan(0, 2), packetId);
        grantedQoS.CopyTo(payload.AsSpan(2));
        return Write(MqttControlPacketType.SubAck, payload);
    }

    /// <summary>
    /// Writes an UNSUBACK packet.
    /// </summary>
    public static byte[] WriteUnsubAck(ushort packetId)
    {
        Span<byte> payload = stackalloc byte[2];
        BinaryPrimitives.WriteUInt16BigEndian(payload, packetId);
        return Write(MqttControlPacketType.UnsubAck, payload);
    }

    /// <summary>
    /// Writes a PINGRESP packet (no payload).
    /// </summary>
    public static byte[] WritePingResp()
        => Write(MqttControlPacketType.PingResp, []);

    /// <summary>
    /// Writes a PUBREC packet (QoS 2 step 1 response).
    /// </summary>
    public static byte[] WritePubRec(ushort packetId)
    {
        Span<byte> payload = stackalloc byte[2];
        BinaryPrimitives.WriteUInt16BigEndian(payload, packetId);
        return Write(MqttControlPacketType.PubRec, payload);
    }

    /// <summary>
    /// Writes a PUBREL packet (QoS 2 step 2). Fixed-header flags must be 0x02 per MQTT spec.
    /// </summary>
    public static byte[] WritePubRel(ushort packetId)
    {
        Span<byte> payload = stackalloc byte[2];
        BinaryPrimitives.WriteUInt16BigEndian(payload, packetId);
        return Write(MqttControlPacketType.PubRel, payload, flags: 0x02);
    }

    /// <summary>
    /// Writes a PUBCOMP packet (QoS 2 step 3 response).
    /// </summary>
    public static byte[] WritePubComp(ushort packetId)
    {
        Span<byte> payload = stackalloc byte[2];
        BinaryPrimitives.WriteUInt16BigEndian(payload, packetId);
        return Write(MqttControlPacketType.PubComp, payload);
    }

    /// <summary>
    /// Writes an MQTT PUBLISH packet for delivery to a client.
    /// </summary>
    public static byte[] WritePublish(string topic, ReadOnlySpan<byte> payload, byte qos = 0,
        bool retain = false, bool dup = false, ushort packetId = 0)
    {
        var topicBytes = Encoding.UTF8.GetBytes(topic);
        var variableHeaderLen = 2 + topicBytes.Length + (qos > 0 ? 2 : 0);
        var totalPayload = new byte[variableHeaderLen + payload.Length];
        var pos = 0;

        // Topic name (length-prefixed)
        BinaryPrimitives.WriteUInt16BigEndian(totalPayload.AsSpan(pos, 2), (ushort)topicBytes.Length);
        pos += 2;
        topicBytes.CopyTo(totalPayload.AsSpan(pos));
        pos += topicBytes.Length;

        // Packet ID (only for QoS > 0)
        if (qos > 0)
        {
            BinaryPrimitives.WriteUInt16BigEndian(totalPayload.AsSpan(pos, 2), packetId);
            pos += 2;
        }

        // Application payload
        payload.CopyTo(totalPayload.AsSpan(pos));

        byte flags = 0;
        if (dup) flags |= 0x08;
        flags |= (byte)((qos & 0x03) << 1);
        if (retain) flags |= 0x01;

        return Write(MqttControlPacketType.Publish, totalPayload, flags);
    }

    /// <summary>
    /// Writes a complete MQTT PUBLISH packet directly into a destination span.
    /// Returns the number of bytes written. Zero-allocation hot path for message delivery.
    /// </summary>
    public static int WritePublishTo(Span<byte> dest, ReadOnlySpan<byte> topicUtf8,
        ReadOnlySpan<byte> payload, byte qos = 0, bool retain = false, bool dup = false, ushort packetId = 0)
    {
        // Calculate remaining length: 2 (topic len) + topic + optional 2 (packet id) + payload
        var remainingLength = 2 + topicUtf8.Length + (qos > 0 ? 2 : 0) + payload.Length;

        // Encode remaining length into scratch
        Span<byte> rlScratch = stackalloc byte[4];
        var rlLen = EncodeRemainingLengthTo(rlScratch, remainingLength);

        var totalLen = 1 + rlLen + remainingLength;

        // Fixed header byte
        byte flags = 0;
        if (dup) flags |= 0x08;
        flags |= (byte)((qos & 0x03) << 1);
        if (retain) flags |= 0x01;
        dest[0] = (byte)(((byte)MqttControlPacketType.Publish << 4) | flags);

        var pos = 1;

        // Remaining length
        rlScratch[..rlLen].CopyTo(dest[pos..]);
        pos += rlLen;

        // Topic name (length-prefixed)
        BinaryPrimitives.WriteUInt16BigEndian(dest[pos..], (ushort)topicUtf8.Length);
        pos += 2;
        topicUtf8.CopyTo(dest[pos..]);
        pos += topicUtf8.Length;

        // Packet ID (only for QoS > 0)
        if (qos > 0)
        {
            BinaryPrimitives.WriteUInt16BigEndian(dest[pos..], packetId);
            pos += 2;
        }

        // Application payload
        payload.CopyTo(dest[pos..]);
        pos += payload.Length;

        return totalLen;
    }

    /// <summary>
    /// Calculates the total wire size of a PUBLISH packet without writing it.
    /// </summary>
    public static int MeasurePublish(int topicLen, int payloadLen, byte qos)
    {
        var remainingLength = 2 + topicLen + (qos > 0 ? 2 : 0) + payloadLen;
        var rlLen = MeasureRemainingLength(remainingLength);
        return 1 + rlLen + remainingLength;
    }

    internal static int EncodeRemainingLengthTo(Span<byte> dest, int value)
    {
        var index = 0;
        do
        {
            var digit = (byte)(value % 128);
            value /= 128;
            if (value > 0)
                digit |= 0x80;
            dest[index++] = digit;
        } while (value > 0);

        return index;
    }

    internal static int MeasureRemainingLength(int value)
    {
        var count = 0;
        do
        {
            value /= 128;
            count++;
        } while (value > 0);

        return count;
    }

    internal static byte[] EncodeRemainingLength(int value)
    {
        if (value < 0 || value > MqttProtocolConstants.MaxPayloadSize)
            throw new ArgumentOutOfRangeException(
                nameof(value),
                $"MQTT remaining length must be between 0 and {MqttProtocolConstants.MaxPayloadSize}.");

        Span<byte> scratch = stackalloc byte[4];
        var index = 0;

        do
        {
            var digit = (byte)(value % 128);
            value /= 128;
            if (value > 0)
                digit |= 0x80;
            scratch[index++] = digit;
        } while (value > 0);

        return scratch[..index].ToArray();
    }
}