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

src/NATS.Server/WebSocket/WsFrameWriter.cs

@@ -0,0 +1,160 @@
+using System.Buffers.Binary;
+using System.Security.Cryptography;
+using System.Text;
+
+namespace NATS.Server.WebSocket;
+
+/// <summary>
+/// WebSocket frame construction, masking, and control message creation.
+/// Ported from golang/nats-server/server/websocket.go lines 543-726.
+/// </summary>
+public static class WsFrameWriter
+{
+    /// <summary>
+    /// Creates a complete frame header for a single-frame message (first=true, final=true).
+    /// Returns (header bytes, mask key or null).
+    /// </summary>
+    public static (byte[] header, byte[]? key) CreateFrameHeader(
+        bool useMasking, bool compressed, int opcode, int payloadLength)
+    {
+        var fh = new byte[WsConstants.MaxFrameHeaderSize];
+        var (n, key) = FillFrameHeader(fh, useMasking,
+            first: true, final: true, compressed: compressed, opcode: opcode, payloadLength: payloadLength);
+        return (fh[..n], key);
+    }
+
+    /// <summary>
+    /// Fills a pre-allocated frame header buffer.
+    /// Returns (bytes written, mask key or null).
+    /// </summary>
+    public static (int written, byte[]? key) FillFrameHeader(
+        Span<byte> fh, bool useMasking, bool first, bool final, bool compressed, int opcode, int payloadLength)
+    {
+        byte b0 = first ? (byte)opcode : (byte)0;
+        if (final) b0 |= WsConstants.FinalBit;
+        if (compressed) b0 |= WsConstants.Rsv1Bit;
+
+        byte b1 = 0;
+        if (useMasking) b1 |= WsConstants.MaskBit;
+
+        int n;
+        switch (payloadLength)
+        {
+            case <= 125:
+                n = 2;
+                fh[0] = b0;
+                fh[1] = (byte)(b1 | (byte)payloadLength);
+                break;
+            case < 65536:
+                n = 4;
+                fh[0] = b0;
+                fh[1] = (byte)(b1 | 126);
+                BinaryPrimitives.WriteUInt16BigEndian(fh[2..], (ushort)payloadLength);
+                break;
+            default:
+                n = 10;
+                fh[0] = b0;
+                fh[1] = (byte)(b1 | 127);
+                BinaryPrimitives.WriteUInt64BigEndian(fh[2..], (ulong)payloadLength);
+                break;
+        }
+
+        byte[]? key = null;
+        if (useMasking)
+        {
+            key = new byte[4];
+            RandomNumberGenerator.Fill(key);
+            key.CopyTo(fh[n..]);
+            n += 4;
+        }
+
+        return (n, key);
+    }
+
+    /// <summary>
+    /// XOR masks a buffer with a 4-byte key. Applies in-place.
+    /// </summary>
+    public static void MaskBuf(ReadOnlySpan<byte> key, Span<byte> buf)
+    {
+        for (int i = 0; i < buf.Length; i++)
+            buf[i] ^= key[i & 3];
+    }
+
+    /// <summary>
+    /// XOR masks multiple contiguous buffers as if they were one.
+    /// </summary>
+    public static void MaskBufs(ReadOnlySpan<byte> key, List<byte[]> bufs)
+    {
+        int pos = 0;
+        foreach (var buf in bufs)
+        {
+            for (int j = 0; j < buf.Length; j++)
+            {
+                buf[j] ^= key[pos & 3];
+                pos++;
+            }
+        }
+    }
+
+    /// <summary>
+    /// Creates a close message payload: 2-byte status code + optional UTF-8 body.
+    /// Body truncated to fit MaxControlPayloadSize with "..." suffix.
+    /// </summary>
+    public static byte[] CreateCloseMessage(int status, string body)
+    {
+        if (body.Length > WsConstants.MaxControlPayloadSize - WsConstants.CloseStatusSize)
+        {
+            body = body[..(WsConstants.MaxControlPayloadSize - WsConstants.CloseStatusSize - 3)] + "...";
+        }
+
+        var bodyBytes = Encoding.UTF8.GetBytes(body);
+        var buf = new byte[WsConstants.CloseStatusSize + bodyBytes.Length];
+        BinaryPrimitives.WriteUInt16BigEndian(buf, (ushort)status);
+        bodyBytes.CopyTo(buf.AsSpan(WsConstants.CloseStatusSize));
+        return buf;
+    }
+
+    /// <summary>
+    /// Builds a complete control frame (header + payload, optional masking).
+    /// </summary>
+    public static byte[] BuildControlFrame(int opcode, ReadOnlySpan<byte> payload, bool useMasking)
+    {
+        int headerSize = 2 + (useMasking ? 4 : 0);
+        var frame = new byte[headerSize + payload.Length];
+        var span = frame.AsSpan();
+        var (n, key) = FillFrameHeader(span, useMasking,
+            first: true, final: true, compressed: false, opcode: opcode, payloadLength: payload.Length);
+        if (payload.Length > 0)
+        {
+            payload.CopyTo(span[n..]);
+            if (useMasking && key != null)
+                MaskBuf(key, span[n..]);
+        }
+
+        return frame;
+    }
+
+    /// <summary>
+    /// Maps a ClientClosedReason to a WebSocket close status code.
+    /// Matches Go wsEnqueueCloseMessage in websocket.go lines 668-694.
+    /// </summary>
+    public static int MapCloseStatus(ClientClosedReason reason) => reason switch
+    {
+        ClientClosedReason.ClientClosed => WsConstants.CloseStatusNormalClosure,
+        ClientClosedReason.AuthenticationTimeout or
+        ClientClosedReason.AuthenticationViolation or
+        ClientClosedReason.SlowConsumerPendingBytes or
+        ClientClosedReason.SlowConsumerWriteDeadline or
+        ClientClosedReason.MaxSubscriptionsExceeded or
+        ClientClosedReason.AuthenticationExpired => WsConstants.CloseStatusPolicyViolation,
+        ClientClosedReason.TlsHandshakeError => WsConstants.CloseStatusTlsHandshake,
+        ClientClosedReason.ParseError or
+        ClientClosedReason.ProtocolViolation => WsConstants.CloseStatusProtocolError,
+        ClientClosedReason.MaxPayloadExceeded => WsConstants.CloseStatusMessageTooBig,
+        ClientClosedReason.WriteError or
+        ClientClosedReason.ReadError or
+        ClientClosedReason.StaleConnection or
+        ClientClosedReason.ServerShutdown => WsConstants.CloseStatusGoingAway,
+        _ => WsConstants.CloseStatusInternalSrvError,
+    };
+}