natsdotnet/src/NATS.Server/WebSocket/WsFrameWriter.cs

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)
    {
        var bodyBytes = Encoding.UTF8.GetBytes(body);
        int maxBody = WsConstants.MaxControlPayloadSize - WsConstants.CloseStatusSize;

        if (bodyBytes.Length > maxBody)
        {
            var suffix = "..."u8;
            int truncLen = maxBody - suffix.Length;
            // Find a valid UTF-8 boundary by walking back from truncation point
            while (truncLen > 0 && (bodyBytes[truncLen] & 0xC0) == 0x80)
                truncLen--;
            var buf = new byte[WsConstants.CloseStatusSize + truncLen + suffix.Length];
            BinaryPrimitives.WriteUInt16BigEndian(buf, (ushort)status);
            bodyBytes.AsSpan(0, truncLen).CopyTo(buf.AsSpan(WsConstants.CloseStatusSize));
            suffix.CopyTo(buf.AsSpan(WsConstants.CloseStatusSize + truncLen));
            return buf;
        }

        var result = new byte[WsConstants.CloseStatusSize + bodyBytes.Length];
        BinaryPrimitives.WriteUInt16BigEndian(result, (ushort)status);
        bodyBytes.CopyTo(result.AsSpan(WsConstants.CloseStatusSize));
        return result;
    }

    /// <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,
    };
}