feat: add WsConnection Stream wrapper for transparent framing

2026-02-23 04:58:56 -05:00
parent fe304dfe01
commit 6d0a4d259e
2 changed files with 280 additions and 0 deletions
--- a/tests/NATS.Server.Tests/WebSocket/WsConnectionTests.cs
+++ b/tests/NATS.Server.Tests/WebSocket/WsConnectionTests.cs
@@ -0,0 +1,88 @@
+using System.Buffers.Binary;
+using NATS.Server.WebSocket;
+
+namespace NATS.Server.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);
+    }
+
+    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;
+    }
+}