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

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

namespace NATS.Server.Transport.Tests.WebSocket;

public class WsConnectionTests
{
    [Fact]
    public async Task ReadAsync_DecodesFrameAndReturnsPayload()
    {
        var payload = "SUB test 1\r\n"u8.ToArray();
        var frame = BuildUnmaskedFrame(payload);
        var inner = new MemoryStream(frame);
        var ws = new WsConnection(inner, compress: false, maskRead: false, maskWrite: false, browser: false, noCompFrag: false);

        var buf = new byte[256];
        int n = await ws.ReadAsync(buf);

        n.ShouldBe(payload.Length);
        buf[..n].ShouldBe(payload);
    }

    [Fact]
    public async Task WriteAsync_FramesPayload()
    {
        var inner = new MemoryStream();
        var ws = new WsConnection(inner, compress: false, maskRead: false, maskWrite: false, browser: false, noCompFrag: false);

        var payload = "MSG test 1 5\r\nHello\r\n"u8.ToArray();
        await ws.WriteAsync(payload);
        await ws.FlushAsync();

        inner.Position = 0;
        var written = inner.ToArray();
        // First 2 bytes should be WS frame header
        (written[0] & WsConstants.FinalBit).ShouldNotBe(0);
        (written[0] & 0x0F).ShouldBe(WsConstants.BinaryMessage);
        int len = written[1] & 0x7F;
        len.ShouldBe(payload.Length);
        written[2..].ShouldBe(payload);
    }

    [Fact]
    public async Task WriteAsync_WithCompression_CompressesLargePayload()
    {
        var inner = new MemoryStream();
        var ws = new WsConnection(inner, compress: true, maskRead: false, maskWrite: false, browser: false, noCompFrag: false);

        var payload = new byte[200];
        Array.Fill<byte>(payload, 0x41); // 'A' repeated - very compressible
        await ws.WriteAsync(payload);
        await ws.FlushAsync();

        inner.Position = 0;
        var written = inner.ToArray();
        // RSV1 bit should be set for compressed frame
        (written[0] & WsConstants.Rsv1Bit).ShouldNotBe(0);
        // Compressed size should be less than original
        written.Length.ShouldBeLessThan(payload.Length + 10);
    }

    [Fact]
    public async Task WriteAsync_SmallPayload_NotCompressedEvenWhenEnabled()
    {
        var inner = new MemoryStream();
        var ws = new WsConnection(inner, compress: true, maskRead: false, maskWrite: false, browser: false, noCompFrag: false);

        var payload = "Hi"u8.ToArray(); // Below CompressThreshold
        await ws.WriteAsync(payload);
        await ws.FlushAsync();

        inner.Position = 0;
        var written = inner.ToArray();
        // RSV1 bit should NOT be set for small payloads
        (written[0] & WsConstants.Rsv1Bit).ShouldBe(0);
    }

    [Fact]
    public async Task ReadAsync_DecodesMaskedFrame()
    {
        var payload = "CONNECT {}\r\n"u8.ToArray();
        var (header, _) = WsFrameWriter.CreateFrameHeader(
            useMasking: true, compressed: false,
            opcode: WsConstants.BinaryMessage, payloadLength: payload.Length);
        var maskKey = header[^4..];
        WsFrameWriter.MaskBuf(maskKey, payload);

        var frame = new byte[header.Length + payload.Length];
        header.CopyTo(frame, 0);
        payload.CopyTo(frame, header.Length);

        var inner = new MemoryStream(frame);
        var ws = new WsConnection(inner, compress: false, maskRead: true, maskWrite: false, browser: false, noCompFrag: false);

        var buf = new byte[256];
        int n = await ws.ReadAsync(buf);

        n.ShouldBe("CONNECT {}\r\n".Length);
        System.Text.Encoding.ASCII.GetString(buf, 0, n).ShouldBe("CONNECT {}\r\n");
    }

    [Fact]
    public async Task ReadAsync_ReturnsZero_OnEndOfStream()
    {
        // Empty stream should return 0 (true end of stream)
        var inner = new MemoryStream([]);
        var ws = new WsConnection(inner, compress: false, maskRead: false, maskWrite: false, browser: false, noCompFrag: false);

        var buf = new byte[256];
        int n = await ws.ReadAsync(buf);
        n.ShouldBe(0);
    }

    private static byte[] BuildUnmaskedFrame(byte[] payload)
    {
        var header = new byte[2];
        header[0] = (byte)(WsConstants.FinalBit | WsConstants.BinaryMessage);
        header[1] = (byte)payload.Length;
        var frame = new byte[2 + payload.Length];
        header.CopyTo(frame, 0);
        payload.CopyTo(frame, 2);
        return frame;
    }
}