natsdotnet/tests/NATS.Server.Transport.Tests/WebSocket/WsFrameWriterTests.cs

using System.Buffers.Binary;
using NATS.Server.WebSocket;
using Shouldly;

namespace NATS.Server.Transport.Tests.WebSocket;

public class WsFrameWriterTests
{
    [Fact]
    public void CreateFrameHeader_SmallPayload_7BitLength()
    {
        var (header, _) = WsFrameWriter.CreateFrameHeader(
            useMasking: false, compressed: false,
            opcode: WsConstants.BinaryMessage, payloadLength: 100);
        header.Length.ShouldBe(2);
        (header[0] & WsConstants.FinalBit).ShouldNotBe(0); // FIN set
        (header[0] & 0x0F).ShouldBe(WsConstants.BinaryMessage);
        (header[1] & 0x7F).ShouldBe(100);
    }

    [Fact]
    public void CreateFrameHeader_MediumPayload_16BitLength()
    {
        var (header, _) = WsFrameWriter.CreateFrameHeader(
            useMasking: false, compressed: false,
            opcode: WsConstants.BinaryMessage, payloadLength: 1000);
        header.Length.ShouldBe(4);
        (header[1] & 0x7F).ShouldBe(126);
        BinaryPrimitives.ReadUInt16BigEndian(header.AsSpan(2)).ShouldBe((ushort)1000);
    }

    [Fact]
    public void CreateFrameHeader_LargePayload_64BitLength()
    {
        var (header, _) = WsFrameWriter.CreateFrameHeader(
            useMasking: false, compressed: false,
            opcode: WsConstants.BinaryMessage, payloadLength: 70000);
        header.Length.ShouldBe(10);
        (header[1] & 0x7F).ShouldBe(127);
        BinaryPrimitives.ReadUInt64BigEndian(header.AsSpan(2)).ShouldBe(70000UL);
    }

    [Fact]
    public void CreateFrameHeader_WithMasking_Adds4ByteKey()
    {
        var (header, key) = WsFrameWriter.CreateFrameHeader(
            useMasking: true, compressed: false,
            opcode: WsConstants.BinaryMessage, payloadLength: 10);
        header.Length.ShouldBe(6); // 2 header + 4 mask key
        (header[1] & WsConstants.MaskBit).ShouldNotBe(0);
        key.ShouldNotBeNull();
        key.Length.ShouldBe(4);
    }

    [Fact]
    public void CreateFrameHeader_Compressed_SetsRsv1Bit()
    {
        var (header, _) = WsFrameWriter.CreateFrameHeader(
            useMasking: false, compressed: true,
            opcode: WsConstants.BinaryMessage, payloadLength: 10);
        (header[0] & WsConstants.Rsv1Bit).ShouldNotBe(0);
    }

    [Fact]
    public void MaskBuf_XorsCorrectly()
    {
        byte[] key = [0xAA, 0xBB, 0xCC, 0xDD];
        byte[] data = [0x01, 0x02, 0x03, 0x04, 0x05, 0x06];
        byte[] expected = new byte[data.Length];
        for (int i = 0; i < data.Length; i++)
            expected[i] = (byte)(data[i] ^ key[i & 3]);

        WsFrameWriter.MaskBuf(key, data);
        data.ShouldBe(expected);
    }

    [Fact]
    public void MaskBuf_RoundTrip()
    {
        byte[] key = [0x12, 0x34, 0x56, 0x78];
        byte[] original = "Hello, WebSocket!"u8.ToArray();
        var data = original.ToArray();

        WsFrameWriter.MaskBuf(key, data);
        data.ShouldNotBe(original);
        WsFrameWriter.MaskBuf(key, data);
        data.ShouldBe(original);
    }

    [Fact]
    public void CreateCloseMessage_WithStatusAndBody()
    {
        var msg = WsFrameWriter.CreateCloseMessage(1000, "normal closure");
        msg.Length.ShouldBe(2 + "normal closure".Length);
        BinaryPrimitives.ReadUInt16BigEndian(msg).ShouldBe((ushort)1000);
    }

    [Fact]
    public void CreateCloseMessage_LongBody_Truncated()
    {
        var longBody = new string('x', 200);
        var msg = WsFrameWriter.CreateCloseMessage(1000, longBody);
        msg.Length.ShouldBeLessThanOrEqualTo(WsConstants.MaxControlPayloadSize);
    }

    [Fact]
    public void MapCloseStatus_ClientClosed_NormalClosure()
    {
        WsFrameWriter.MapCloseStatus(ClientClosedReason.ClientClosed)
            .ShouldBe(WsConstants.CloseStatusNormalClosure);
    }

    [Fact]
    public void MapCloseStatus_AuthTimeout_PolicyViolation()
    {
        WsFrameWriter.MapCloseStatus(ClientClosedReason.AuthenticationTimeout)
            .ShouldBe(WsConstants.CloseStatusPolicyViolation);
    }

    [Fact]
    public void MapCloseStatus_ParseError_ProtocolError()
    {
        WsFrameWriter.MapCloseStatus(ClientClosedReason.ParseError)
            .ShouldBe(WsConstants.CloseStatusProtocolError);
    }

    [Fact]
    public void MapCloseStatus_MaxPayload_MessageTooBig()
    {
        WsFrameWriter.MapCloseStatus(ClientClosedReason.MaxPayloadExceeded)
            .ShouldBe(WsConstants.CloseStatusMessageTooBig);
    }

    [Fact]
    public void BuildControlFrame_PingNomask()
    {
        var frame = WsFrameWriter.BuildControlFrame(WsConstants.PingMessage, [], useMasking: false);
        frame.Length.ShouldBe(2);
        (frame[0] & WsConstants.FinalBit).ShouldNotBe(0);
        (frame[0] & 0x0F).ShouldBe(WsConstants.PingMessage);
        (frame[1] & 0x7F).ShouldBe(0);
    }

    [Fact]
    public void BuildControlFrame_PongWithPayload()
    {
        byte[] payload = [1, 2, 3, 4];
        var frame = WsFrameWriter.BuildControlFrame(WsConstants.PongMessage, payload, useMasking: false);
        frame.Length.ShouldBe(2 + 4);
        frame[2..].ShouldBe(payload);
    }
}