feat: add leaf node WebSocket support with stream adapter (Gap 12.5)
Implements WebSocketStreamAdapter — a Stream subclass that wraps System.Net.WebSockets.WebSocket for use by LeafConnection. Handles message framing (per-message receive/send), tracks BytesRead/BytesWritten and MessagesRead/MessagesWritten counters, and exposes IsConnected. Ten NSubstitute-based unit tests cover all capability flags, delegation, and telemetry (10/10 pass).
This commit is contained in:
Joseph Doherty
235
src/NATS.Server/LeafNodes/WebSocketStreamAdapter.cs
Normal file
235
src/NATS.Server/LeafNodes/WebSocketStreamAdapter.cs
Normal file
@@ -0,0 +1,235 @@
|
||||
using SystemWebSocket = System.Net.WebSockets.WebSocket;
|
||||
using System.Net.WebSockets;
|
||||
|
||||
namespace NATS.Server.LeafNodes;
|
||||
|
||||
/// <summary>
|
||||
/// Adapts a System.Net.WebSockets.WebSocket into a Stream suitable for use
|
||||
/// by LeafConnection. Handles message framing: reads aggregate WebSocket messages
|
||||
/// into a contiguous byte stream, and writes flush as single WebSocket messages.
|
||||
/// Go reference: leafnode.go wsCreateLeafConnection, client.go wsRead/wsWrite.
|
||||
/// </summary>
|
||||
public sealed class WebSocketStreamAdapter : Stream
|
||||
{
|
||||
private readonly SystemWebSocket _ws;
|
||||
private byte[] _readBuffer;
|
||||
private int _readOffset;
|
||||
private int _readCount;
|
||||
private bool _disposed;
|
||||
|
||||
public WebSocketStreamAdapter(SystemWebSocket ws, int initialBufferSize = 4096)
|
||||
{
|
||||
_ws = ws ?? throw new ArgumentNullException(nameof(ws));
|
||||
_readBuffer = new byte[Math.Max(initialBufferSize, 64)];
|
||||
_readOffset = 0;
|
||||
_readCount = 0;
|
||||
}
|
||||
|
||||
// Stream capability overrides
|
||||
public override bool CanRead => true;
|
||||
public override bool CanWrite => true;
|
||||
public override bool CanSeek => false;
|
||||
|
||||
// Telemetry properties
|
||||
public bool IsConnected => _ws.State == WebSocketState.Open;
|
||||
public long BytesRead { get; private set; }
|
||||
public long BytesWritten { get; private set; }
|
||||
public int MessagesRead { get; private set; }
|
||||
public int MessagesWritten { get; private set; }
|
||||
|
||||
/// <summary>
|
||||
/// Reads data from the WebSocket into <paramref name="buffer"/>.
|
||||
/// If the internal read buffer has buffered data from a previous message,
|
||||
/// that is served first. Otherwise a new WebSocket message is received.
|
||||
/// Go reference: client.go wsRead.
|
||||
/// </summary>
|
||||
public override async Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken ct)
|
||||
{
|
||||
ObjectDisposedException.ThrowIf(_disposed, this);
|
||||
|
||||
// Drain any leftover data from the previous WebSocket message first.
|
||||
if (_readCount > 0)
|
||||
{
|
||||
var fromBuffer = Math.Min(_readCount, count);
|
||||
_readBuffer.AsSpan(_readOffset, fromBuffer).CopyTo(buffer.AsSpan(offset, fromBuffer));
|
||||
_readOffset += fromBuffer;
|
||||
_readCount -= fromBuffer;
|
||||
if (_readCount == 0)
|
||||
_readOffset = 0;
|
||||
return fromBuffer;
|
||||
}
|
||||
|
||||
// Receive the next WebSocket message, growing the buffer as needed.
|
||||
var totalReceived = 0;
|
||||
while (true)
|
||||
{
|
||||
EnsureReadBufferCapacity(totalReceived + 1024);
|
||||
var result = await _ws.ReceiveAsync(
|
||||
_readBuffer.AsMemory(totalReceived),
|
||||
ct).ConfigureAwait(false);
|
||||
|
||||
if (result.MessageType == WebSocketMessageType.Close)
|
||||
return 0;
|
||||
|
||||
totalReceived += result.Count;
|
||||
|
||||
if (result.EndOfMessage)
|
||||
{
|
||||
MessagesRead++;
|
||||
BytesRead += totalReceived;
|
||||
|
||||
// Copy what fits into the caller's buffer; remainder stays in _readBuffer.
|
||||
var toCopy = Math.Min(totalReceived, count);
|
||||
_readBuffer.AsSpan(0, toCopy).CopyTo(buffer.AsSpan(offset, toCopy));
|
||||
|
||||
var remaining = totalReceived - toCopy;
|
||||
if (remaining > 0)
|
||||
{
|
||||
_readOffset = toCopy;
|
||||
_readCount = remaining;
|
||||
}
|
||||
else
|
||||
{
|
||||
_readOffset = 0;
|
||||
_readCount = 0;
|
||||
}
|
||||
|
||||
return toCopy;
|
||||
}
|
||||
|
||||
// Partial message — make sure buffer has room for more data.
|
||||
EnsureReadBufferCapacity(totalReceived + 1024);
|
||||
}
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken ct = default)
|
||||
{
|
||||
ObjectDisposedException.ThrowIf(_disposed, this);
|
||||
|
||||
// Drain buffered data first.
|
||||
if (_readCount > 0)
|
||||
{
|
||||
var fromBuffer = Math.Min(_readCount, buffer.Length);
|
||||
_readBuffer.AsMemory(_readOffset, fromBuffer).CopyTo(buffer[..fromBuffer]);
|
||||
_readOffset += fromBuffer;
|
||||
_readCount -= fromBuffer;
|
||||
if (_readCount == 0)
|
||||
_readOffset = 0;
|
||||
return fromBuffer;
|
||||
}
|
||||
|
||||
// Receive the next WebSocket message into a temporary staging area.
|
||||
var totalReceived = 0;
|
||||
while (true)
|
||||
{
|
||||
EnsureReadBufferCapacity(totalReceived + 1024);
|
||||
var result = await _ws.ReceiveAsync(
|
||||
_readBuffer.AsMemory(totalReceived),
|
||||
ct).ConfigureAwait(false);
|
||||
|
||||
if (result.MessageType == WebSocketMessageType.Close)
|
||||
return 0;
|
||||
|
||||
totalReceived += result.Count;
|
||||
|
||||
if (result.EndOfMessage)
|
||||
{
|
||||
MessagesRead++;
|
||||
BytesRead += totalReceived;
|
||||
|
||||
var toCopy = Math.Min(totalReceived, buffer.Length);
|
||||
_readBuffer.AsMemory(0, toCopy).CopyTo(buffer[..toCopy]);
|
||||
|
||||
var remaining = totalReceived - toCopy;
|
||||
if (remaining > 0)
|
||||
{
|
||||
_readOffset = toCopy;
|
||||
_readCount = remaining;
|
||||
}
|
||||
else
|
||||
{
|
||||
_readOffset = 0;
|
||||
_readCount = 0;
|
||||
}
|
||||
|
||||
return toCopy;
|
||||
}
|
||||
|
||||
EnsureReadBufferCapacity(totalReceived + 1024);
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Sends <paramref name="buffer"/> as a single binary WebSocket message.
|
||||
/// Go reference: client.go wsWrite.
|
||||
/// </summary>
|
||||
public override async Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken ct)
|
||||
{
|
||||
ObjectDisposedException.ThrowIf(_disposed, this);
|
||||
await _ws.SendAsync(
|
||||
buffer.AsMemory(offset, count),
|
||||
WebSocketMessageType.Binary,
|
||||
endOfMessage: true,
|
||||
ct).ConfigureAwait(false);
|
||||
BytesWritten += count;
|
||||
MessagesWritten++;
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken ct = default)
|
||||
{
|
||||
ObjectDisposedException.ThrowIf(_disposed, this);
|
||||
await _ws.SendAsync(
|
||||
buffer,
|
||||
WebSocketMessageType.Binary,
|
||||
endOfMessage: true,
|
||||
ct).ConfigureAwait(false);
|
||||
BytesWritten += buffer.Length;
|
||||
MessagesWritten++;
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public override Task FlushAsync(CancellationToken ct) => Task.CompletedTask;
|
||||
|
||||
// Not-supported synchronous and seeking members
|
||||
public override long Length => throw new NotSupportedException();
|
||||
public override long Position
|
||||
{
|
||||
get => throw new NotSupportedException();
|
||||
set => throw new NotSupportedException();
|
||||
}
|
||||
public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
|
||||
public override void SetLength(long value) => throw new NotSupportedException();
|
||||
public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException("Use async methods");
|
||||
public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException("Use async methods");
|
||||
public override void Flush() { }
|
||||
|
||||
protected override void Dispose(bool disposing)
|
||||
{
|
||||
if (_disposed)
|
||||
return;
|
||||
_disposed = true;
|
||||
if (disposing)
|
||||
_ws.Dispose();
|
||||
base.Dispose(disposing);
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// -------------------------------------------------------------------------
|
||||
|
||||
private void EnsureReadBufferCapacity(int required)
|
||||
{
|
||||
if (_readBuffer.Length >= required)
|
||||
return;
|
||||
|
||||
var newSize = Math.Max(required, _readBuffer.Length * 2);
|
||||
var next = new byte[newSize];
|
||||
if (_readCount > 0)
|
||||
_readBuffer.AsSpan(_readOffset, _readCount).CopyTo(next);
|
||||
_readBuffer = next;
|
||||
_readOffset = 0;
|
||||
// _readCount unchanged
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user