Merge branch 'feature/websocket'
# Conflicts: # differences.md
This commit is contained in:
@@ -11,6 +11,7 @@ using NATS.Server.Auth;
|
||||
using NATS.Server.Protocol;
|
||||
using NATS.Server.Subscriptions;
|
||||
using NATS.Server.Tls;
|
||||
using NATS.Server.WebSocket;
|
||||
|
||||
namespace NATS.Server;
|
||||
|
||||
@@ -93,6 +94,9 @@ public sealed class NatsClient : IDisposable
|
||||
private long _rtt;
|
||||
public TimeSpan Rtt => new(Interlocked.Read(ref _rtt));
|
||||
|
||||
public bool IsWebSocket { get; set; }
|
||||
public WsUpgradeResult? WsInfo { get; set; }
|
||||
|
||||
public TlsConnectionState? TlsState { get; set; }
|
||||
public bool InfoAlreadySent { get; set; }
|
||||
|
||||
|
||||
@@ -116,4 +116,32 @@ public sealed class NatsOptions
|
||||
public Dictionary<string, string>? SubjectMappings { get; set; }
|
||||
|
||||
public bool HasTls => TlsCert != null && TlsKey != null;
|
||||
|
||||
// WebSocket
|
||||
public WebSocketOptions WebSocket { get; set; } = new();
|
||||
}
|
||||
|
||||
public sealed class WebSocketOptions
|
||||
{
|
||||
public string Host { get; set; } = "0.0.0.0";
|
||||
public int Port { get; set; } = -1;
|
||||
public string? Advertise { get; set; }
|
||||
public string? NoAuthUser { get; set; }
|
||||
public string? JwtCookie { get; set; }
|
||||
public string? UsernameCookie { get; set; }
|
||||
public string? PasswordCookie { get; set; }
|
||||
public string? TokenCookie { get; set; }
|
||||
public string? Username { get; set; }
|
||||
public string? Password { get; set; }
|
||||
public string? Token { get; set; }
|
||||
public TimeSpan AuthTimeout { get; set; } = TimeSpan.FromSeconds(2);
|
||||
public bool NoTls { get; set; }
|
||||
public string? TlsCert { get; set; }
|
||||
public string? TlsKey { get; set; }
|
||||
public bool SameOrigin { get; set; }
|
||||
public List<string>? AllowedOrigins { get; set; }
|
||||
public bool Compression { get; set; }
|
||||
public TimeSpan HandshakeTimeout { get; set; } = TimeSpan.FromSeconds(2);
|
||||
public TimeSpan? PingInterval { get; set; }
|
||||
public Dictionary<string, string>? Headers { get; set; }
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ using NATS.Server.Monitoring;
|
||||
using NATS.Server.Protocol;
|
||||
using NATS.Server.Subscriptions;
|
||||
using NATS.Server.Tls;
|
||||
using NATS.Server.WebSocket;
|
||||
|
||||
namespace NATS.Server;
|
||||
|
||||
@@ -39,6 +40,8 @@ public sealed class NatsServer : IMessageRouter, ISubListAccess, IDisposable
|
||||
private readonly TlsRateLimiter? _tlsRateLimiter;
|
||||
private readonly SubjectTransform[] _subjectTransforms;
|
||||
private Socket? _listener;
|
||||
private Socket? _wsListener;
|
||||
private readonly TaskCompletionSource _wsAcceptLoopExited = new(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
private MonitorServer? _monitorServer;
|
||||
private ulong _nextClientId;
|
||||
private long _startTimeTicks;
|
||||
@@ -93,11 +96,13 @@ public sealed class NatsServer : IMessageRouter, ISubListAccess, IDisposable
|
||||
// Signal all internal loops to stop
|
||||
await _quitCts.CancelAsync();
|
||||
|
||||
// Close listener to stop accept loop
|
||||
// Close listeners to stop accept loops
|
||||
_listener?.Close();
|
||||
_wsListener?.Close();
|
||||
|
||||
// Wait for accept loop to exit
|
||||
// Wait for accept loops to exit
|
||||
await _acceptLoopExited.Task.WaitAsync(TimeSpan.FromSeconds(5)).ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
|
||||
await _wsAcceptLoopExited.Task.WaitAsync(TimeSpan.FromSeconds(5)).ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
|
||||
|
||||
// Close all client connections — flush first, then mark closed
|
||||
var flushTasks = new List<Task>();
|
||||
@@ -138,11 +143,13 @@ public sealed class NatsServer : IMessageRouter, ISubListAccess, IDisposable
|
||||
|
||||
_logger.LogInformation("Entering lame duck mode, stop accepting new clients");
|
||||
|
||||
// Close listener to stop accepting new connections
|
||||
// Close listeners to stop accepting new connections
|
||||
_listener?.Close();
|
||||
_wsListener?.Close();
|
||||
|
||||
// Wait for accept loop to exit
|
||||
// Wait for accept loops to exit
|
||||
await _acceptLoopExited.Task.WaitAsync(TimeSpan.FromSeconds(5)).ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
|
||||
await _wsAcceptLoopExited.Task.WaitAsync(TimeSpan.FromSeconds(5)).ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
|
||||
|
||||
var gracePeriod = _options.LameDuckGracePeriod;
|
||||
if (gracePeriod < TimeSpan.Zero) gracePeriod = -gracePeriod;
|
||||
@@ -369,8 +376,6 @@ public sealed class NatsServer : IMessageRouter, ISubListAccess, IDisposable
|
||||
BuildCachedInfo();
|
||||
}
|
||||
|
||||
_listeningStarted.TrySetResult();
|
||||
|
||||
_logger.LogInformation("Listening for client connections on {Host}:{Port}", _options.Host, _options.Port);
|
||||
|
||||
// Warn about stub features
|
||||
@@ -386,6 +391,31 @@ public sealed class NatsServer : IMessageRouter, ISubListAccess, IDisposable
|
||||
WritePidFile();
|
||||
WritePortsFile();
|
||||
|
||||
if (_options.WebSocket.Port >= 0)
|
||||
{
|
||||
_wsListener = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
|
||||
_wsListener.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.ReuseAddress, true);
|
||||
_wsListener.Bind(new IPEndPoint(
|
||||
_options.WebSocket.Host == "0.0.0.0" ? IPAddress.Any : IPAddress.Parse(_options.WebSocket.Host),
|
||||
_options.WebSocket.Port));
|
||||
_wsListener.Listen(128);
|
||||
|
||||
if (_options.WebSocket.Port == 0)
|
||||
{
|
||||
_options.WebSocket.Port = ((IPEndPoint)_wsListener.LocalEndPoint!).Port;
|
||||
}
|
||||
|
||||
_logger.LogInformation("Listening for WebSocket clients on {Host}:{Port}",
|
||||
_options.WebSocket.Host, _options.WebSocket.Port);
|
||||
|
||||
if (_options.WebSocket.NoTls)
|
||||
_logger.LogWarning("WebSocket not configured with TLS. DO NOT USE IN PRODUCTION!");
|
||||
|
||||
_ = RunWebSocketAcceptLoopAsync(linked.Token);
|
||||
}
|
||||
|
||||
_listeningStarted.TrySetResult();
|
||||
|
||||
var tmpDelay = AcceptMinSleep;
|
||||
|
||||
try
|
||||
@@ -531,6 +561,102 @@ public sealed class NatsServer : IMessageRouter, ISubListAccess, IDisposable
|
||||
}
|
||||
}
|
||||
|
||||
private async Task RunWebSocketAcceptLoopAsync(CancellationToken ct)
|
||||
{
|
||||
var tmpDelay = AcceptMinSleep;
|
||||
try
|
||||
{
|
||||
while (!ct.IsCancellationRequested)
|
||||
{
|
||||
Socket socket;
|
||||
try
|
||||
{
|
||||
socket = await _wsListener!.AcceptAsync(ct);
|
||||
tmpDelay = AcceptMinSleep;
|
||||
}
|
||||
catch (OperationCanceledException) { break; }
|
||||
catch (ObjectDisposedException) { break; }
|
||||
catch (SocketException ex)
|
||||
{
|
||||
if (IsShuttingDown || IsLameDuckMode) break;
|
||||
_logger.LogError(ex, "Temporary WebSocket accept error, sleeping {Delay}ms", tmpDelay.TotalMilliseconds);
|
||||
try { await Task.Delay(tmpDelay, ct); } catch (OperationCanceledException) { break; }
|
||||
tmpDelay = TimeSpan.FromTicks(Math.Min(tmpDelay.Ticks * 2, AcceptMaxSleep.Ticks));
|
||||
continue;
|
||||
}
|
||||
|
||||
if (_options.MaxConnections > 0 && _clients.Count >= _options.MaxConnections)
|
||||
{
|
||||
socket.Dispose();
|
||||
continue;
|
||||
}
|
||||
|
||||
var clientId = Interlocked.Increment(ref _nextClientId);
|
||||
Interlocked.Increment(ref _stats.TotalConnections);
|
||||
Interlocked.Increment(ref _activeClientCount);
|
||||
|
||||
_ = AcceptWebSocketClientAsync(socket, clientId, ct);
|
||||
}
|
||||
}
|
||||
finally
|
||||
{
|
||||
_wsAcceptLoopExited.TrySetResult();
|
||||
}
|
||||
}
|
||||
|
||||
private async Task AcceptWebSocketClientAsync(Socket socket, ulong clientId, CancellationToken ct)
|
||||
{
|
||||
try
|
||||
{
|
||||
var networkStream = new NetworkStream(socket, ownsSocket: false);
|
||||
Stream stream = networkStream;
|
||||
|
||||
// TLS negotiation if configured
|
||||
if (_sslOptions != null && !_options.WebSocket.NoTls)
|
||||
{
|
||||
var (tlsStream, _) = await TlsConnectionWrapper.NegotiateAsync(
|
||||
socket, networkStream, _options, _sslOptions, _serverInfo,
|
||||
_loggerFactory.CreateLogger("NATS.Server.Tls"), ct);
|
||||
stream = tlsStream;
|
||||
}
|
||||
|
||||
// HTTP upgrade handshake
|
||||
var upgradeResult = await WsUpgrade.TryUpgradeAsync(stream, stream, _options.WebSocket, ct);
|
||||
if (!upgradeResult.Success)
|
||||
{
|
||||
_logger.LogDebug("WebSocket upgrade failed for client {ClientId}", clientId);
|
||||
socket.Dispose();
|
||||
Interlocked.Decrement(ref _activeClientCount);
|
||||
return;
|
||||
}
|
||||
|
||||
// Create WsConnection wrapper
|
||||
var wsConn = new WsConnection(stream,
|
||||
compress: upgradeResult.Compress,
|
||||
maskRead: upgradeResult.MaskRead,
|
||||
maskWrite: upgradeResult.MaskWrite,
|
||||
browser: upgradeResult.Browser,
|
||||
noCompFrag: upgradeResult.NoCompFrag);
|
||||
|
||||
var clientLogger = _loggerFactory.CreateLogger($"NATS.Server.NatsClient[{clientId}]");
|
||||
var client = new NatsClient(clientId, wsConn, socket, _options, _serverInfo,
|
||||
_authService, null, clientLogger, _stats);
|
||||
client.Router = this;
|
||||
client.IsWebSocket = true;
|
||||
client.WsInfo = upgradeResult;
|
||||
_clients[clientId] = client;
|
||||
|
||||
await RunClientAsync(client, ct);
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
_logger.LogDebug(ex, "Failed to accept WebSocket client {ClientId}", clientId);
|
||||
try { socket.Shutdown(SocketShutdown.Both); } catch { }
|
||||
socket.Dispose();
|
||||
Interlocked.Decrement(ref _activeClientCount);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task RunClientAsync(NatsClient client, CancellationToken ct)
|
||||
{
|
||||
try
|
||||
@@ -942,6 +1068,7 @@ public sealed class NatsServer : IMessageRouter, ISubListAccess, IDisposable
|
||||
_quitCts.Dispose();
|
||||
_tlsRateLimiter?.Dispose();
|
||||
_listener?.Dispose();
|
||||
_wsListener?.Dispose();
|
||||
foreach (var client in _clients.Values)
|
||||
client.Dispose();
|
||||
foreach (var account in _accounts.Values)
|
||||
|
||||
94
src/NATS.Server/WebSocket/WsCompression.cs
Normal file
94
src/NATS.Server/WebSocket/WsCompression.cs
Normal file
@@ -0,0 +1,94 @@
|
||||
using System.IO.Compression;
|
||||
|
||||
namespace NATS.Server.WebSocket;
|
||||
|
||||
/// <summary>
|
||||
/// permessage-deflate compression/decompression for WebSocket frames (RFC 7692).
|
||||
/// Ported from golang/nats-server/server/websocket.go lines 403-440 and 1391-1466.
|
||||
/// </summary>
|
||||
public static class WsCompression
|
||||
{
|
||||
/// <summary>
|
||||
/// Compresses data using deflate. Removes trailing 4 bytes (sync marker)
|
||||
/// per RFC 7692 Section 7.2.1.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// We call Flush() but intentionally do not Dispose() the DeflateStream before
|
||||
/// reading output, because Dispose writes a final deflate block (0x03 0x00) that
|
||||
/// would be corrupted by the 4-byte tail strip. Flush() alone writes a sync flush
|
||||
/// ending with 0x00 0x00 0xff 0xff, matching Go's flate.Writer.Flush() behavior.
|
||||
/// </remarks>
|
||||
public static byte[] Compress(ReadOnlySpan<byte> data)
|
||||
{
|
||||
var output = new MemoryStream();
|
||||
var deflate = new DeflateStream(output, CompressionLevel.Fastest, leaveOpen: true);
|
||||
try
|
||||
{
|
||||
deflate.Write(data);
|
||||
deflate.Flush();
|
||||
|
||||
var compressed = output.ToArray();
|
||||
|
||||
// Remove trailing 4-byte sync marker (0x00 0x00 0xff 0xff) per RFC 7692
|
||||
if (compressed.Length >= 4)
|
||||
return compressed[..^4];
|
||||
|
||||
return compressed;
|
||||
}
|
||||
finally
|
||||
{
|
||||
deflate.Dispose();
|
||||
output.Dispose();
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Decompresses collected compressed buffers.
|
||||
/// Appends trailer bytes before decompressing per RFC 7692 Section 7.2.2.
|
||||
/// Ported from golang/nats-server/server/websocket.go lines 403-440.
|
||||
/// The Go code appends compressLastBlock (9 bytes) which includes the sync
|
||||
/// marker plus a final empty stored block to signal end-of-stream to the
|
||||
/// flate reader.
|
||||
/// </summary>
|
||||
public static byte[] Decompress(List<byte[]> compressedBuffers, int maxPayload)
|
||||
{
|
||||
if (maxPayload <= 0)
|
||||
maxPayload = 1024 * 1024; // Default 1MB
|
||||
|
||||
// Concatenate all compressed buffers + trailer.
|
||||
// Per RFC 7692 Section 7.2.2, append the sync flush marker (0x00 0x00 0xff 0xff)
|
||||
// that was stripped during compression. The Go reference appends compressLastBlock
|
||||
// (9 bytes) for Go's flate reader; .NET's DeflateStream only needs the 4-byte trailer.
|
||||
int totalLen = 0;
|
||||
foreach (var buf in compressedBuffers)
|
||||
totalLen += buf.Length;
|
||||
totalLen += WsConstants.DecompressTrailer.Length;
|
||||
|
||||
var combined = new byte[totalLen];
|
||||
int offset = 0;
|
||||
foreach (var buf in compressedBuffers)
|
||||
{
|
||||
buf.CopyTo(combined, offset);
|
||||
offset += buf.Length;
|
||||
}
|
||||
|
||||
WsConstants.DecompressTrailer.CopyTo(combined, offset);
|
||||
|
||||
using var input = new MemoryStream(combined);
|
||||
using var deflate = new DeflateStream(input, CompressionMode.Decompress);
|
||||
using var output = new MemoryStream();
|
||||
|
||||
var readBuf = new byte[4096];
|
||||
int totalRead = 0;
|
||||
int n;
|
||||
while ((n = deflate.Read(readBuf, 0, readBuf.Length)) > 0)
|
||||
{
|
||||
totalRead += n;
|
||||
if (totalRead > maxPayload)
|
||||
throw new InvalidOperationException("decompressed data exceeds maximum payload size");
|
||||
output.Write(readBuf, 0, n);
|
||||
}
|
||||
|
||||
return output.ToArray();
|
||||
}
|
||||
}
|
||||
202
src/NATS.Server/WebSocket/WsConnection.cs
Normal file
202
src/NATS.Server/WebSocket/WsConnection.cs
Normal file
@@ -0,0 +1,202 @@
|
||||
namespace NATS.Server.WebSocket;
|
||||
|
||||
/// <summary>
|
||||
/// Stream wrapper that transparently frames/deframes WebSocket around raw TCP I/O.
|
||||
/// NatsClient uses this as its _stream -- FillPipeAsync and RunWriteLoopAsync work unchanged.
|
||||
/// Ported from golang/nats-server/server/websocket.go wsUpgrade/wrapWebsocket pattern.
|
||||
/// </summary>
|
||||
public sealed class WsConnection : Stream
|
||||
{
|
||||
private readonly Stream _inner;
|
||||
private readonly bool _compress;
|
||||
private readonly bool _maskRead;
|
||||
private readonly bool _maskWrite;
|
||||
private readonly bool _browser;
|
||||
private readonly bool _noCompFrag;
|
||||
private WsReadInfo _readInfo;
|
||||
// Read-side state: accessed only from the single FillPipeAsync reader task (no synchronization needed)
|
||||
private readonly Queue<byte[]> _readQueue = new();
|
||||
private int _readOffset;
|
||||
private readonly object _writeLock = new();
|
||||
private readonly List<ControlFrameAction> _pendingControlWrites = [];
|
||||
|
||||
public bool CloseReceived => _readInfo.CloseReceived;
|
||||
public int CloseStatus => _readInfo.CloseStatus;
|
||||
|
||||
public WsConnection(Stream inner, bool compress, bool maskRead, bool maskWrite, bool browser, bool noCompFrag)
|
||||
{
|
||||
_inner = inner;
|
||||
_compress = compress;
|
||||
_maskRead = maskRead;
|
||||
_maskWrite = maskWrite;
|
||||
_browser = browser;
|
||||
_noCompFrag = noCompFrag;
|
||||
_readInfo = new WsReadInfo(expectMask: maskRead);
|
||||
}
|
||||
|
||||
public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken ct = default)
|
||||
{
|
||||
// Drain any buffered decoded payloads first
|
||||
if (_readQueue.Count > 0)
|
||||
return DrainReadQueue(buffer.Span);
|
||||
|
||||
while (true)
|
||||
{
|
||||
// Read raw bytes from inner stream
|
||||
var rawBuf = new byte[Math.Max(buffer.Length, 4096)];
|
||||
int bytesRead = await _inner.ReadAsync(rawBuf.AsMemory(), ct);
|
||||
if (bytesRead == 0) return 0;
|
||||
|
||||
// Decode frames
|
||||
var payloads = WsReadInfo.ReadFrames(_readInfo, new MemoryStream(rawBuf, 0, bytesRead), bytesRead, maxPayload: 1024 * 1024);
|
||||
|
||||
// Collect control frame responses
|
||||
if (_readInfo.PendingControlFrames.Count > 0)
|
||||
{
|
||||
lock (_writeLock)
|
||||
_pendingControlWrites.AddRange(_readInfo.PendingControlFrames);
|
||||
_readInfo.PendingControlFrames.Clear();
|
||||
// Write pending control frames
|
||||
await FlushControlFramesAsync(ct);
|
||||
}
|
||||
|
||||
if (_readInfo.CloseReceived)
|
||||
return 0;
|
||||
|
||||
foreach (var payload in payloads)
|
||||
_readQueue.Enqueue(payload);
|
||||
|
||||
// If no payloads were decoded (e.g. only frame headers were read),
|
||||
// continue reading instead of returning 0 which signals end-of-stream
|
||||
if (_readQueue.Count > 0)
|
||||
return DrainReadQueue(buffer.Span);
|
||||
}
|
||||
}
|
||||
|
||||
public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken ct = default)
|
||||
{
|
||||
var data = buffer.Span;
|
||||
|
||||
if (_compress && data.Length > WsConstants.CompressThreshold)
|
||||
{
|
||||
var compressed = WsCompression.Compress(data);
|
||||
await WriteFramedAsync(compressed, compressed: true, ct);
|
||||
}
|
||||
else
|
||||
{
|
||||
await WriteFramedAsync(data.ToArray(), compressed: false, ct);
|
||||
}
|
||||
}
|
||||
|
||||
private async ValueTask WriteFramedAsync(byte[] payload, bool compressed, CancellationToken ct)
|
||||
{
|
||||
if (_browser && payload.Length > WsConstants.FrameSizeForBrowsers && !(_noCompFrag && compressed))
|
||||
{
|
||||
// Fragment for browsers
|
||||
int offset = 0;
|
||||
bool first = true;
|
||||
while (offset < payload.Length)
|
||||
{
|
||||
int chunkLen = Math.Min(WsConstants.FrameSizeForBrowsers, payload.Length - offset);
|
||||
bool final = offset + chunkLen >= payload.Length;
|
||||
var fh = new byte[WsConstants.MaxFrameHeaderSize];
|
||||
var (n, key) = WsFrameWriter.FillFrameHeader(fh, _maskWrite,
|
||||
first: first, final: final, compressed: first && compressed,
|
||||
opcode: WsConstants.BinaryMessage, payloadLength: chunkLen);
|
||||
|
||||
var chunk = payload.AsSpan(offset, chunkLen).ToArray();
|
||||
if (_maskWrite && key != null)
|
||||
WsFrameWriter.MaskBuf(key, chunk);
|
||||
|
||||
await _inner.WriteAsync(fh.AsMemory(0, n), ct);
|
||||
await _inner.WriteAsync(chunk.AsMemory(), ct);
|
||||
offset += chunkLen;
|
||||
first = false;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
var (header, key) = WsFrameWriter.CreateFrameHeader(_maskWrite, compressed, WsConstants.BinaryMessage, payload.Length);
|
||||
if (_maskWrite && key != null)
|
||||
WsFrameWriter.MaskBuf(key, payload);
|
||||
await _inner.WriteAsync(header.AsMemory(), ct);
|
||||
await _inner.WriteAsync(payload.AsMemory(), ct);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task FlushControlFramesAsync(CancellationToken ct)
|
||||
{
|
||||
List<ControlFrameAction> toWrite;
|
||||
lock (_writeLock)
|
||||
{
|
||||
if (_pendingControlWrites.Count == 0) return;
|
||||
toWrite = [.. _pendingControlWrites];
|
||||
_pendingControlWrites.Clear();
|
||||
}
|
||||
|
||||
foreach (var action in toWrite)
|
||||
{
|
||||
var frame = WsFrameWriter.BuildControlFrame(action.Opcode, action.Payload, _maskWrite);
|
||||
await _inner.WriteAsync(frame, ct);
|
||||
}
|
||||
await _inner.FlushAsync(ct);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Sends a WebSocket close frame.
|
||||
/// </summary>
|
||||
public async Task SendCloseAsync(ClientClosedReason reason, CancellationToken ct = default)
|
||||
{
|
||||
var status = WsFrameWriter.MapCloseStatus(reason);
|
||||
var closePayload = WsFrameWriter.CreateCloseMessage(status, reason.ToReasonString());
|
||||
var frame = WsFrameWriter.BuildControlFrame(WsConstants.CloseMessage, closePayload, _maskWrite);
|
||||
await _inner.WriteAsync(frame, ct);
|
||||
await _inner.FlushAsync(ct);
|
||||
}
|
||||
|
||||
private int DrainReadQueue(Span<byte> buffer)
|
||||
{
|
||||
int written = 0;
|
||||
while (_readQueue.Count > 0 && written < buffer.Length)
|
||||
{
|
||||
var current = _readQueue.Peek();
|
||||
int available = current.Length - _readOffset;
|
||||
int toCopy = Math.Min(available, buffer.Length - written);
|
||||
current.AsSpan(_readOffset, toCopy).CopyTo(buffer[written..]);
|
||||
written += toCopy;
|
||||
_readOffset += toCopy;
|
||||
if (_readOffset >= current.Length)
|
||||
{
|
||||
_readQueue.Dequeue();
|
||||
_readOffset = 0;
|
||||
}
|
||||
}
|
||||
return written;
|
||||
}
|
||||
|
||||
// Stream abstract members
|
||||
public override bool CanRead => true;
|
||||
public override bool CanWrite => true;
|
||||
public override bool CanSeek => false;
|
||||
public override long Length => throw new NotSupportedException();
|
||||
public override long Position { get => throw new NotSupportedException(); set => throw new NotSupportedException(); }
|
||||
public override void Flush() => _inner.Flush();
|
||||
public override Task FlushAsync(CancellationToken ct) => _inner.FlushAsync(ct);
|
||||
public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException("Use ReadAsync");
|
||||
public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException("Use WriteAsync");
|
||||
public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
|
||||
public override void SetLength(long value) => throw new NotSupportedException();
|
||||
|
||||
protected override void Dispose(bool disposing)
|
||||
{
|
||||
if (disposing)
|
||||
_inner.Dispose();
|
||||
base.Dispose(disposing);
|
||||
}
|
||||
|
||||
public override async ValueTask DisposeAsync()
|
||||
{
|
||||
await _inner.DisposeAsync();
|
||||
GC.SuppressFinalize(this);
|
||||
}
|
||||
}
|
||||
72
src/NATS.Server/WebSocket/WsConstants.cs
Normal file
72
src/NATS.Server/WebSocket/WsConstants.cs
Normal file
@@ -0,0 +1,72 @@
|
||||
namespace NATS.Server.WebSocket;
|
||||
|
||||
/// <summary>
|
||||
/// WebSocket protocol constants (RFC 6455).
|
||||
/// Ported from golang/nats-server/server/websocket.go lines 41-106.
|
||||
/// </summary>
|
||||
public static class WsConstants
|
||||
{
|
||||
// Opcodes (RFC 6455 Section 5.2)
|
||||
public const int TextMessage = 1;
|
||||
public const int BinaryMessage = 2;
|
||||
public const int CloseMessage = 8;
|
||||
public const int PingMessage = 9;
|
||||
public const int PongMessage = 10;
|
||||
public const int ContinuationFrame = 0;
|
||||
|
||||
// Frame header bits
|
||||
public const byte FinalBit = 0x80; // 1 << 7
|
||||
public const byte Rsv1Bit = 0x40; // 1 << 6 (compression, RFC 7692)
|
||||
public const byte Rsv2Bit = 0x20; // 1 << 5
|
||||
public const byte Rsv3Bit = 0x10; // 1 << 4
|
||||
public const byte MaskBit = 0x80; // 1 << 7 (in second byte)
|
||||
|
||||
// Frame size limits
|
||||
public const int MaxFrameHeaderSize = 14;
|
||||
public const int MaxControlPayloadSize = 125;
|
||||
public const int FrameSizeForBrowsers = 4096;
|
||||
public const int CompressThreshold = 64;
|
||||
public const int CloseStatusSize = 2;
|
||||
|
||||
// Close status codes (RFC 6455 Section 11.7)
|
||||
public const int CloseStatusNormalClosure = 1000;
|
||||
public const int CloseStatusGoingAway = 1001;
|
||||
public const int CloseStatusProtocolError = 1002;
|
||||
public const int CloseStatusUnsupportedData = 1003;
|
||||
public const int CloseStatusNoStatusReceived = 1005;
|
||||
public const int CloseStatusInvalidPayloadData = 1007;
|
||||
public const int CloseStatusPolicyViolation = 1008;
|
||||
public const int CloseStatusMessageTooBig = 1009;
|
||||
public const int CloseStatusInternalSrvError = 1011;
|
||||
public const int CloseStatusTlsHandshake = 1015;
|
||||
|
||||
// Compression constants (RFC 7692)
|
||||
public const string PmcExtension = "permessage-deflate";
|
||||
public const string PmcSrvNoCtx = "server_no_context_takeover";
|
||||
public const string PmcCliNoCtx = "client_no_context_takeover";
|
||||
public static readonly string PmcReqHeaderValue = $"{PmcExtension}; {PmcSrvNoCtx}; {PmcCliNoCtx}";
|
||||
public static readonly string PmcFullResponse = $"Sec-WebSocket-Extensions: {PmcExtension}; {PmcSrvNoCtx}; {PmcCliNoCtx}\r\n";
|
||||
|
||||
// Header names
|
||||
public const string NoMaskingHeader = "Nats-No-Masking";
|
||||
public const string NoMaskingValue = "true";
|
||||
public static readonly string NoMaskingFullResponse = $"{NoMaskingHeader}: {NoMaskingValue}\r\n";
|
||||
public const string XForwardedForHeader = "X-Forwarded-For";
|
||||
|
||||
// Path routing
|
||||
public const string ClientPath = "/";
|
||||
public const string LeafNodePath = "/leafnode";
|
||||
public const string MqttPath = "/mqtt";
|
||||
|
||||
// Decompression trailer appended before decompressing (RFC 7692 Section 7.2.2)
|
||||
public static readonly byte[] DecompressTrailer = [0x00, 0x00, 0xff, 0xff];
|
||||
|
||||
public static bool IsControlFrame(int opcode) => opcode >= CloseMessage;
|
||||
}
|
||||
|
||||
public enum WsClientKind
|
||||
{
|
||||
Client,
|
||||
Leaf,
|
||||
Mqtt,
|
||||
}
|
||||
171
src/NATS.Server/WebSocket/WsFrameWriter.cs
Normal file
171
src/NATS.Server/WebSocket/WsFrameWriter.cs
Normal file
@@ -0,0 +1,171 @@
|
||||
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,
|
||||
};
|
||||
}
|
||||
81
src/NATS.Server/WebSocket/WsOriginChecker.cs
Normal file
81
src/NATS.Server/WebSocket/WsOriginChecker.cs
Normal file
@@ -0,0 +1,81 @@
|
||||
namespace NATS.Server.WebSocket;
|
||||
|
||||
/// <summary>
|
||||
/// Validates WebSocket Origin headers per RFC 6455 Section 10.2.
|
||||
/// Ported from golang/nats-server/server/websocket.go lines 933-1000.
|
||||
/// </summary>
|
||||
public sealed class WsOriginChecker
|
||||
{
|
||||
private readonly bool _sameOrigin;
|
||||
private readonly Dictionary<string, AllowedOrigin>? _allowedOrigins;
|
||||
|
||||
public WsOriginChecker(bool sameOrigin, List<string>? allowedOrigins)
|
||||
{
|
||||
_sameOrigin = sameOrigin;
|
||||
if (allowedOrigins is { Count: > 0 })
|
||||
{
|
||||
_allowedOrigins = new Dictionary<string, AllowedOrigin>(StringComparer.OrdinalIgnoreCase);
|
||||
foreach (var ao in allowedOrigins)
|
||||
{
|
||||
if (Uri.TryCreate(ao, UriKind.Absolute, out var uri))
|
||||
{
|
||||
var (host, port) = GetHostAndPort(uri.Scheme == "https", uri.Host, uri.Port);
|
||||
_allowedOrigins[host] = new AllowedOrigin(uri.Scheme, port);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Returns null if origin is allowed, or an error message if rejected.
|
||||
/// </summary>
|
||||
public string? CheckOrigin(string? origin, string requestHost, bool isTls)
|
||||
{
|
||||
if (!_sameOrigin && _allowedOrigins == null)
|
||||
return null;
|
||||
|
||||
if (string.IsNullOrEmpty(origin))
|
||||
return null;
|
||||
|
||||
if (!Uri.TryCreate(origin, UriKind.Absolute, out var originUri))
|
||||
return $"invalid origin: {origin}";
|
||||
|
||||
var (oh, op) = GetHostAndPort(originUri.Scheme == "https", originUri.Host, originUri.Port);
|
||||
|
||||
if (_sameOrigin)
|
||||
{
|
||||
var (rh, rp) = ParseHostPort(requestHost, isTls);
|
||||
if (!string.Equals(oh, rh, StringComparison.OrdinalIgnoreCase) || op != rp)
|
||||
return "not same origin";
|
||||
}
|
||||
|
||||
if (_allowedOrigins != null)
|
||||
{
|
||||
if (!_allowedOrigins.TryGetValue(oh, out var allowed) ||
|
||||
!string.Equals(originUri.Scheme, allowed.Scheme, StringComparison.OrdinalIgnoreCase) ||
|
||||
op != allowed.Port)
|
||||
{
|
||||
return "not in the allowed list";
|
||||
}
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
private static (string host, int port) GetHostAndPort(bool tls, string host, int port)
|
||||
{
|
||||
if (port <= 0)
|
||||
port = tls ? 443 : 80;
|
||||
return (host.ToLowerInvariant(), port);
|
||||
}
|
||||
|
||||
private static (string host, int port) ParseHostPort(string hostPort, bool isTls)
|
||||
{
|
||||
var colonIdx = hostPort.LastIndexOf(':');
|
||||
if (colonIdx > 0 && int.TryParse(hostPort.AsSpan(colonIdx + 1), out var port))
|
||||
return (hostPort[..colonIdx].ToLowerInvariant(), port);
|
||||
return (hostPort.ToLowerInvariant(), isTls ? 443 : 80);
|
||||
}
|
||||
|
||||
private readonly record struct AllowedOrigin(string Scheme, int Port);
|
||||
}
|
||||
322
src/NATS.Server/WebSocket/WsReadInfo.cs
Normal file
322
src/NATS.Server/WebSocket/WsReadInfo.cs
Normal file
@@ -0,0 +1,322 @@
|
||||
using System.Buffers.Binary;
|
||||
using System.Text;
|
||||
|
||||
namespace NATS.Server.WebSocket;
|
||||
|
||||
/// <summary>
|
||||
/// Per-connection WebSocket frame reading state machine.
|
||||
/// Ported from golang/nats-server/server/websocket.go lines 156-506.
|
||||
/// </summary>
|
||||
public class WsReadInfo
|
||||
{
|
||||
public int Remaining;
|
||||
public bool FrameStart;
|
||||
public bool FirstFrame;
|
||||
public bool FrameCompressed;
|
||||
public bool ExpectMask;
|
||||
public byte MaskKeyPos;
|
||||
public byte[] MaskKey;
|
||||
public List<byte[]>? CompressedBuffers;
|
||||
public int CompressedOffset;
|
||||
|
||||
// Control frame outputs
|
||||
public List<ControlFrameAction> PendingControlFrames;
|
||||
public bool CloseReceived;
|
||||
public int CloseStatus;
|
||||
public string? CloseBody;
|
||||
|
||||
public WsReadInfo(bool expectMask)
|
||||
{
|
||||
Remaining = 0;
|
||||
FrameStart = true;
|
||||
FirstFrame = true;
|
||||
FrameCompressed = false;
|
||||
ExpectMask = expectMask;
|
||||
MaskKeyPos = 0;
|
||||
MaskKey = new byte[4];
|
||||
CompressedBuffers = null;
|
||||
CompressedOffset = 0;
|
||||
PendingControlFrames = [];
|
||||
CloseReceived = false;
|
||||
CloseStatus = 0;
|
||||
CloseBody = null;
|
||||
}
|
||||
|
||||
public void SetMaskKey(ReadOnlySpan<byte> key)
|
||||
{
|
||||
key[..4].CopyTo(MaskKey);
|
||||
MaskKeyPos = 0;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Unmask buffer in-place using current mask key and position.
|
||||
/// Optimized for 8-byte chunks when buffer is large enough.
|
||||
/// Ported from websocket.go lines 509-536.
|
||||
/// </summary>
|
||||
public void Unmask(Span<byte> buf)
|
||||
{
|
||||
int p = MaskKeyPos;
|
||||
if (buf.Length < 16)
|
||||
{
|
||||
for (int i = 0; i < buf.Length; i++)
|
||||
{
|
||||
buf[i] ^= MaskKey[p & 3];
|
||||
p++;
|
||||
}
|
||||
MaskKeyPos = (byte)(p & 3);
|
||||
return;
|
||||
}
|
||||
|
||||
// Build 8-byte key for bulk XOR
|
||||
Span<byte> k = stackalloc byte[8];
|
||||
for (int i = 0; i < 8; i++)
|
||||
k[i] = MaskKey[(p + i) & 3];
|
||||
ulong km = BinaryPrimitives.ReadUInt64BigEndian(k);
|
||||
|
||||
int n = (buf.Length / 8) * 8;
|
||||
for (int i = 0; i < n; i += 8)
|
||||
{
|
||||
ulong tmp = BinaryPrimitives.ReadUInt64BigEndian(buf[i..]);
|
||||
tmp ^= km;
|
||||
BinaryPrimitives.WriteUInt64BigEndian(buf[i..], tmp);
|
||||
}
|
||||
|
||||
// Handle remaining bytes
|
||||
p += n;
|
||||
var tail = buf[n..];
|
||||
for (int i = 0; i < tail.Length; i++)
|
||||
{
|
||||
tail[i] ^= MaskKey[p & 3];
|
||||
p++;
|
||||
}
|
||||
MaskKeyPos = (byte)(p & 3);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Read and decode WebSocket frames from a buffer.
|
||||
/// Returns list of decoded payload byte arrays.
|
||||
/// Ported from websocket.go lines 208-351.
|
||||
/// </summary>
|
||||
public static List<byte[]> ReadFrames(WsReadInfo r, Stream stream, int available, int maxPayload)
|
||||
{
|
||||
var bufs = new List<byte[]>();
|
||||
var buf = new byte[available];
|
||||
int bytesRead = 0;
|
||||
|
||||
// Fill the buffer from the stream
|
||||
while (bytesRead < available)
|
||||
{
|
||||
int n = stream.Read(buf, bytesRead, available - bytesRead);
|
||||
if (n == 0) break;
|
||||
bytesRead += n;
|
||||
}
|
||||
|
||||
int pos = 0;
|
||||
int max = bytesRead;
|
||||
|
||||
while (pos < max)
|
||||
{
|
||||
if (r.FrameStart)
|
||||
{
|
||||
if (pos >= max) break;
|
||||
byte b0 = buf[pos];
|
||||
int frameType = b0 & 0x0F;
|
||||
bool final = (b0 & WsConstants.FinalBit) != 0;
|
||||
bool compressed = (b0 & WsConstants.Rsv1Bit) != 0;
|
||||
pos++;
|
||||
|
||||
// Read second byte
|
||||
var (b1Buf, newPos) = WsGet(stream, buf, pos, max, 1);
|
||||
pos = newPos;
|
||||
byte b1 = b1Buf[0];
|
||||
|
||||
// Check mask bit
|
||||
if (r.ExpectMask && (b1 & WsConstants.MaskBit) == 0)
|
||||
throw new InvalidOperationException("mask bit missing");
|
||||
|
||||
r.Remaining = b1 & 0x7F;
|
||||
|
||||
// Validate frame types
|
||||
if (WsConstants.IsControlFrame(frameType))
|
||||
{
|
||||
if (r.Remaining > WsConstants.MaxControlPayloadSize)
|
||||
throw new InvalidOperationException("control frame length too large");
|
||||
if (!final)
|
||||
throw new InvalidOperationException("control frame does not have final bit set");
|
||||
}
|
||||
else if (frameType == WsConstants.TextMessage || frameType == WsConstants.BinaryMessage)
|
||||
{
|
||||
if (!r.FirstFrame)
|
||||
throw new InvalidOperationException("new message before previous finished");
|
||||
r.FirstFrame = final;
|
||||
r.FrameCompressed = compressed;
|
||||
}
|
||||
else if (frameType == WsConstants.ContinuationFrame)
|
||||
{
|
||||
if (r.FirstFrame || compressed)
|
||||
throw new InvalidOperationException("invalid continuation frame");
|
||||
r.FirstFrame = final;
|
||||
}
|
||||
else
|
||||
{
|
||||
throw new InvalidOperationException($"unknown opcode {frameType}");
|
||||
}
|
||||
|
||||
// Extended payload length
|
||||
switch (r.Remaining)
|
||||
{
|
||||
case 126:
|
||||
{
|
||||
var (lenBuf, p2) = WsGet(stream, buf, pos, max, 2);
|
||||
pos = p2;
|
||||
r.Remaining = BinaryPrimitives.ReadUInt16BigEndian(lenBuf);
|
||||
break;
|
||||
}
|
||||
case 127:
|
||||
{
|
||||
var (lenBuf, p2) = WsGet(stream, buf, pos, max, 8);
|
||||
pos = p2;
|
||||
var len64 = BinaryPrimitives.ReadUInt64BigEndian(lenBuf);
|
||||
if (len64 > (ulong)maxPayload)
|
||||
throw new InvalidOperationException($"frame payload length {len64} exceeds max payload {maxPayload}");
|
||||
r.Remaining = (int)len64;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Read mask key (mask bit already validated at line 134)
|
||||
if (r.ExpectMask)
|
||||
{
|
||||
var (keyBuf, p2) = WsGet(stream, buf, pos, max, 4);
|
||||
pos = p2;
|
||||
keyBuf.AsSpan(0, 4).CopyTo(r.MaskKey);
|
||||
r.MaskKeyPos = 0;
|
||||
}
|
||||
|
||||
// Handle control frames
|
||||
if (WsConstants.IsControlFrame(frameType))
|
||||
{
|
||||
pos = HandleControlFrame(r, frameType, stream, buf, pos, max);
|
||||
continue;
|
||||
}
|
||||
|
||||
r.FrameStart = false;
|
||||
}
|
||||
|
||||
if (pos < max)
|
||||
{
|
||||
int n = r.Remaining;
|
||||
if (pos + n > max) n = max - pos;
|
||||
|
||||
var payloadSlice = buf.AsSpan(pos, n).ToArray();
|
||||
pos += n;
|
||||
r.Remaining -= n;
|
||||
|
||||
if (r.ExpectMask)
|
||||
r.Unmask(payloadSlice);
|
||||
|
||||
bool addToBufs = true;
|
||||
if (r.FrameCompressed)
|
||||
{
|
||||
addToBufs = false;
|
||||
r.CompressedBuffers ??= [];
|
||||
r.CompressedBuffers.Add(payloadSlice);
|
||||
|
||||
if (r.FirstFrame && r.Remaining == 0)
|
||||
{
|
||||
var decompressed = WsCompression.Decompress(r.CompressedBuffers, maxPayload);
|
||||
r.CompressedBuffers = null;
|
||||
r.FrameCompressed = false;
|
||||
addToBufs = true;
|
||||
payloadSlice = decompressed;
|
||||
}
|
||||
}
|
||||
|
||||
if (addToBufs && payloadSlice.Length > 0)
|
||||
bufs.Add(payloadSlice);
|
||||
|
||||
if (r.Remaining == 0)
|
||||
r.FrameStart = true;
|
||||
}
|
||||
}
|
||||
|
||||
return bufs;
|
||||
}
|
||||
|
||||
private static int HandleControlFrame(WsReadInfo r, int frameType, Stream stream, byte[] buf, int pos, int max)
|
||||
{
|
||||
byte[]? payload = null;
|
||||
if (r.Remaining > 0)
|
||||
{
|
||||
var (payloadBuf, newPos) = WsGet(stream, buf, pos, max, r.Remaining);
|
||||
pos = newPos;
|
||||
payload = payloadBuf;
|
||||
if (r.ExpectMask)
|
||||
r.Unmask(payload);
|
||||
r.Remaining = 0;
|
||||
}
|
||||
|
||||
switch (frameType)
|
||||
{
|
||||
case WsConstants.CloseMessage:
|
||||
r.CloseReceived = true;
|
||||
r.CloseStatus = WsConstants.CloseStatusNoStatusReceived;
|
||||
if (payload != null && payload.Length >= WsConstants.CloseStatusSize)
|
||||
{
|
||||
r.CloseStatus = BinaryPrimitives.ReadUInt16BigEndian(payload);
|
||||
if (payload.Length > WsConstants.CloseStatusSize)
|
||||
r.CloseBody = Encoding.UTF8.GetString(payload.AsSpan(WsConstants.CloseStatusSize));
|
||||
}
|
||||
// Per RFC 6455 Section 5.5.1, always send a close response
|
||||
if (r.CloseStatus != WsConstants.CloseStatusNoStatusReceived)
|
||||
{
|
||||
var closeMsg = WsFrameWriter.CreateCloseMessage(r.CloseStatus, r.CloseBody ?? "");
|
||||
r.PendingControlFrames.Add(new ControlFrameAction(WsConstants.CloseMessage, closeMsg));
|
||||
}
|
||||
else
|
||||
{
|
||||
// Empty close frame — respond with empty close
|
||||
r.PendingControlFrames.Add(new ControlFrameAction(WsConstants.CloseMessage, []));
|
||||
}
|
||||
break;
|
||||
|
||||
case WsConstants.PingMessage:
|
||||
r.PendingControlFrames.Add(new ControlFrameAction(WsConstants.PongMessage, payload ?? []));
|
||||
break;
|
||||
|
||||
case WsConstants.PongMessage:
|
||||
// Nothing to do
|
||||
break;
|
||||
}
|
||||
|
||||
return pos;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Gets needed bytes from buffer or reads from stream.
|
||||
/// Ported from websocket.go lines 178-193.
|
||||
/// </summary>
|
||||
private static (byte[] data, int newPos) WsGet(Stream stream, byte[] buf, int pos, int max, int needed)
|
||||
{
|
||||
int avail = max - pos;
|
||||
if (avail >= needed)
|
||||
return (buf[pos..(pos + needed)], pos + needed);
|
||||
|
||||
var b = new byte[needed];
|
||||
int start = 0;
|
||||
if (avail > 0)
|
||||
{
|
||||
Buffer.BlockCopy(buf, pos, b, 0, avail);
|
||||
start = avail;
|
||||
}
|
||||
while (start < needed)
|
||||
{
|
||||
int n = stream.Read(b, start, needed - start);
|
||||
if (n == 0) throw new IOException("unexpected end of stream");
|
||||
start += n;
|
||||
}
|
||||
return (b, pos + avail);
|
||||
}
|
||||
}
|
||||
|
||||
public readonly record struct ControlFrameAction(int Opcode, byte[] Payload);
|
||||
268
src/NATS.Server/WebSocket/WsUpgrade.cs
Normal file
268
src/NATS.Server/WebSocket/WsUpgrade.cs
Normal file
@@ -0,0 +1,268 @@
|
||||
using System.Net;
|
||||
using System.Security.Cryptography;
|
||||
using System.Text;
|
||||
|
||||
namespace NATS.Server.WebSocket;
|
||||
|
||||
/// <summary>
|
||||
/// WebSocket HTTP upgrade handshake handler.
|
||||
/// Ported from golang/nats-server/server/websocket.go lines 731-917.
|
||||
/// </summary>
|
||||
public static class WsUpgrade
|
||||
{
|
||||
public static async Task<WsUpgradeResult> TryUpgradeAsync(
|
||||
Stream inputStream, Stream outputStream, WebSocketOptions options,
|
||||
CancellationToken ct = default)
|
||||
{
|
||||
try
|
||||
{
|
||||
using var cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
|
||||
cts.CancelAfter(options.HandshakeTimeout);
|
||||
var (method, path, headers) = await ReadHttpRequestAsync(inputStream, cts.Token);
|
||||
|
||||
if (!string.Equals(method, "GET", StringComparison.OrdinalIgnoreCase))
|
||||
return await FailAsync(outputStream, 405, "request method must be GET");
|
||||
|
||||
if (!headers.ContainsKey("Host"))
|
||||
return await FailAsync(outputStream, 400, "'Host' missing in request");
|
||||
|
||||
if (!HeaderContains(headers, "Upgrade", "websocket"))
|
||||
return await FailAsync(outputStream, 400, "invalid value for header 'Upgrade'");
|
||||
|
||||
if (!HeaderContains(headers, "Connection", "Upgrade"))
|
||||
return await FailAsync(outputStream, 400, "invalid value for header 'Connection'");
|
||||
|
||||
if (!headers.TryGetValue("Sec-WebSocket-Key", out var key) || string.IsNullOrEmpty(key))
|
||||
return await FailAsync(outputStream, 400, "key missing");
|
||||
|
||||
if (!HeaderContains(headers, "Sec-WebSocket-Version", "13"))
|
||||
return await FailAsync(outputStream, 400, "invalid version");
|
||||
|
||||
var kind = path switch
|
||||
{
|
||||
_ when path.EndsWith("/leafnode") => WsClientKind.Leaf,
|
||||
_ when path.EndsWith("/mqtt") => WsClientKind.Mqtt,
|
||||
_ => WsClientKind.Client,
|
||||
};
|
||||
|
||||
// Origin checking
|
||||
if (options.SameOrigin || options.AllowedOrigins is { Count: > 0 })
|
||||
{
|
||||
var checker = new WsOriginChecker(options.SameOrigin, options.AllowedOrigins);
|
||||
headers.TryGetValue("Origin", out var origin);
|
||||
if (string.IsNullOrEmpty(origin))
|
||||
headers.TryGetValue("Sec-WebSocket-Origin", out origin);
|
||||
var originErr = checker.CheckOrigin(origin, headers.GetValueOrDefault("Host", ""), isTls: false);
|
||||
if (originErr != null)
|
||||
return await FailAsync(outputStream, 403, $"origin not allowed: {originErr}");
|
||||
}
|
||||
|
||||
// Compression negotiation
|
||||
bool compress = options.Compression;
|
||||
if (compress)
|
||||
{
|
||||
compress = headers.TryGetValue("Sec-WebSocket-Extensions", out var ext) &&
|
||||
ext.Contains(WsConstants.PmcExtension, StringComparison.OrdinalIgnoreCase);
|
||||
}
|
||||
|
||||
// No-masking support (leaf nodes only — browser clients must always mask)
|
||||
bool noMasking = kind == WsClientKind.Leaf &&
|
||||
headers.TryGetValue(WsConstants.NoMaskingHeader, out var nmVal) &&
|
||||
string.Equals(nmVal.Trim(), WsConstants.NoMaskingValue, StringComparison.OrdinalIgnoreCase);
|
||||
|
||||
// Browser detection
|
||||
bool browser = false;
|
||||
bool noCompFrag = false;
|
||||
if (kind is WsClientKind.Client or WsClientKind.Mqtt &&
|
||||
headers.TryGetValue("User-Agent", out var ua) && ua.StartsWith("Mozilla/"))
|
||||
{
|
||||
browser = true;
|
||||
// Disable fragmentation of compressed frames for Safari browsers.
|
||||
// Safari has both "Version/" and "Safari/" in the user agent string,
|
||||
// while Chrome on macOS has "Safari/" but not "Version/".
|
||||
noCompFrag = compress && ua.Contains("Version/") && ua.Contains("Safari/");
|
||||
}
|
||||
|
||||
// Cookie extraction
|
||||
string? cookieJwt = null, cookieUsername = null, cookiePassword = null, cookieToken = null;
|
||||
if ((kind is WsClientKind.Client or WsClientKind.Mqtt) &&
|
||||
headers.TryGetValue("Cookie", out var cookieHeader))
|
||||
{
|
||||
var cookies = ParseCookies(cookieHeader);
|
||||
if (options.JwtCookie != null) cookies.TryGetValue(options.JwtCookie, out cookieJwt);
|
||||
if (options.UsernameCookie != null) cookies.TryGetValue(options.UsernameCookie, out cookieUsername);
|
||||
if (options.PasswordCookie != null) cookies.TryGetValue(options.PasswordCookie, out cookiePassword);
|
||||
if (options.TokenCookie != null) cookies.TryGetValue(options.TokenCookie, out cookieToken);
|
||||
}
|
||||
|
||||
// X-Forwarded-For client IP extraction
|
||||
string? clientIp = null;
|
||||
if (headers.TryGetValue(WsConstants.XForwardedForHeader, out var xff))
|
||||
{
|
||||
var ip = xff.Split(',')[0].Trim();
|
||||
if (IPAddress.TryParse(ip, out _))
|
||||
clientIp = ip;
|
||||
}
|
||||
|
||||
// Build the 101 Switching Protocols response
|
||||
var response = new StringBuilder();
|
||||
response.Append("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: ");
|
||||
response.Append(ComputeAcceptKey(key));
|
||||
response.Append("\r\n");
|
||||
if (compress)
|
||||
response.Append(WsConstants.PmcFullResponse);
|
||||
if (noMasking)
|
||||
response.Append(WsConstants.NoMaskingFullResponse);
|
||||
if (options.Headers != null)
|
||||
{
|
||||
foreach (var (k, v) in options.Headers)
|
||||
{
|
||||
response.Append(k);
|
||||
response.Append(": ");
|
||||
response.Append(v);
|
||||
response.Append("\r\n");
|
||||
}
|
||||
}
|
||||
|
||||
response.Append("\r\n");
|
||||
|
||||
var responseBytes = Encoding.ASCII.GetBytes(response.ToString());
|
||||
await outputStream.WriteAsync(responseBytes);
|
||||
await outputStream.FlushAsync();
|
||||
|
||||
return new WsUpgradeResult(
|
||||
Success: true, Compress: compress, Browser: browser, NoCompFrag: noCompFrag,
|
||||
MaskRead: !noMasking, MaskWrite: false,
|
||||
CookieJwt: cookieJwt, CookieUsername: cookieUsername,
|
||||
CookiePassword: cookiePassword, CookieToken: cookieToken,
|
||||
ClientIp: clientIp, Kind: kind);
|
||||
}
|
||||
catch (Exception)
|
||||
{
|
||||
return WsUpgradeResult.Failed;
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Computes the Sec-WebSocket-Accept value per RFC 6455 Section 4.2.2.
|
||||
/// </summary>
|
||||
public static string ComputeAcceptKey(string clientKey)
|
||||
{
|
||||
var combined = Encoding.ASCII.GetBytes(clientKey + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
|
||||
var hash = SHA1.HashData(combined);
|
||||
return Convert.ToBase64String(hash);
|
||||
}
|
||||
|
||||
private static async Task<WsUpgradeResult> FailAsync(Stream output, int statusCode, string reason)
|
||||
{
|
||||
var statusText = statusCode switch
|
||||
{
|
||||
400 => "Bad Request",
|
||||
403 => "Forbidden",
|
||||
405 => "Method Not Allowed",
|
||||
_ => "Internal Server Error",
|
||||
};
|
||||
var response = $"HTTP/1.1 {statusCode} {statusText}\r\nSec-WebSocket-Version: 13\r\nContent-Type: text/plain\r\nContent-Length: {reason.Length}\r\n\r\n{reason}";
|
||||
await output.WriteAsync(Encoding.ASCII.GetBytes(response));
|
||||
await output.FlushAsync();
|
||||
return WsUpgradeResult.Failed;
|
||||
}
|
||||
|
||||
private static async Task<(string method, string path, Dictionary<string, string> headers)> ReadHttpRequestAsync(
|
||||
Stream stream, CancellationToken ct)
|
||||
{
|
||||
var headerBytes = new List<byte>(4096);
|
||||
var buf = new byte[512];
|
||||
while (true)
|
||||
{
|
||||
int n = await stream.ReadAsync(buf, ct);
|
||||
if (n == 0) throw new IOException("connection closed during handshake");
|
||||
for (int i = 0; i < n; i++)
|
||||
{
|
||||
headerBytes.Add(buf[i]);
|
||||
if (headerBytes.Count >= 4 &&
|
||||
headerBytes[^4] == '\r' && headerBytes[^3] == '\n' &&
|
||||
headerBytes[^2] == '\r' && headerBytes[^1] == '\n')
|
||||
goto done;
|
||||
if (headerBytes.Count > 8192)
|
||||
throw new InvalidOperationException("HTTP header too large");
|
||||
}
|
||||
}
|
||||
done:;
|
||||
|
||||
var text = Encoding.ASCII.GetString(headerBytes.ToArray());
|
||||
var lines = text.Split("\r\n", StringSplitOptions.None);
|
||||
if (lines.Length < 1) throw new InvalidOperationException("invalid HTTP request");
|
||||
|
||||
var parts = lines[0].Split(' ');
|
||||
if (parts.Length < 3) throw new InvalidOperationException("invalid HTTP request line");
|
||||
var method = parts[0];
|
||||
var path = parts[1];
|
||||
|
||||
var headers = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
|
||||
for (int i = 1; i < lines.Length; i++)
|
||||
{
|
||||
var line = lines[i];
|
||||
if (string.IsNullOrEmpty(line)) break;
|
||||
var colonIdx = line.IndexOf(':');
|
||||
if (colonIdx > 0)
|
||||
{
|
||||
var name = line[..colonIdx].Trim();
|
||||
var value = line[(colonIdx + 1)..].Trim();
|
||||
headers[name] = value;
|
||||
}
|
||||
}
|
||||
|
||||
return (method, path, headers);
|
||||
}
|
||||
|
||||
private static bool HeaderContains(Dictionary<string, string> headers, string name, string value)
|
||||
{
|
||||
if (!headers.TryGetValue(name, out var headerValue))
|
||||
return false;
|
||||
foreach (var token in headerValue.Split(','))
|
||||
{
|
||||
if (string.Equals(token.Trim(), value, StringComparison.OrdinalIgnoreCase))
|
||||
return true;
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
private static Dictionary<string, string> ParseCookies(string cookieHeader)
|
||||
{
|
||||
var cookies = new Dictionary<string, string>(StringComparer.Ordinal);
|
||||
foreach (var pair in cookieHeader.Split(';'))
|
||||
{
|
||||
var trimmed = pair.Trim();
|
||||
var eqIdx = trimmed.IndexOf('=');
|
||||
if (eqIdx > 0)
|
||||
cookies[trimmed[..eqIdx].Trim()] = trimmed[(eqIdx + 1)..].Trim();
|
||||
}
|
||||
|
||||
return cookies;
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Result of a WebSocket upgrade handshake attempt.
|
||||
/// </summary>
|
||||
public readonly record struct WsUpgradeResult(
|
||||
bool Success,
|
||||
bool Compress,
|
||||
bool Browser,
|
||||
bool NoCompFrag,
|
||||
bool MaskRead,
|
||||
bool MaskWrite,
|
||||
string? CookieJwt,
|
||||
string? CookieUsername,
|
||||
string? CookiePassword,
|
||||
string? CookieToken,
|
||||
string? ClientIp,
|
||||
WsClientKind Kind)
|
||||
{
|
||||
public static readonly WsUpgradeResult Failed = new(
|
||||
Success: false, Compress: false, Browser: false, NoCompFrag: false,
|
||||
MaskRead: true, MaskWrite: false, CookieJwt: null, CookieUsername: null,
|
||||
CookiePassword: null, CookieToken: null, ClientIp: null, Kind: WsClientKind.Client);
|
||||
}
|
||||
Reference in New Issue
Block a user