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); } }