feat: add WebSocket frame writer with masking and close status mapping

2026-02-23 04:40:44 -05:00
parent 6981a38b72
commit 8ded10d49b
2 changed files with 312 additions and 0 deletions
--- a/tests/NATS.Server.Tests/WebSocket/WsFrameWriterTests.cs
+++ b/tests/NATS.Server.Tests/WebSocket/WsFrameWriterTests.cs
@@ -0,0 +1,152 @@
+using System.Buffers.Binary;
+using NATS.Server.WebSocket;
+using Shouldly;
+
+namespace NATS.Server.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);
+    }
+}