feat: integrate WebSocket accept loop into NatsServer and NatsClient
Add WebSocket listener support to NatsServer alongside the existing TCP listener. When WebSocketOptions.Port >= 0, the server binds a second socket, performs HTTP upgrade via WsUpgrade.TryUpgradeAsync, wraps the connection in WsConnection for transparent frame/deframe, and hands it to the standard NatsClient pipeline. Changes: - NatsClient: add IsWebSocket and WsInfo properties - NatsServer: add RunWebSocketAcceptLoopAsync and AcceptWebSocketClientAsync, WS listener lifecycle in StartAsync/ShutdownAsync/Dispose - NatsOptions: change WebSocketOptions.Port default from 0 to -1 (disabled) - WsConnection.ReadAsync: fix premature end-of-stream when ReadFrames returns no payloads by looping until data is available - Add WsIntegration tests (connect, ping, pub/sub over WebSocket) - Add WsConnection masked frame and end-of-stream unit tests
This commit is contained in:
Joseph Doherty
@@ -5,10 +5,10 @@ namespace NATS.Server.Tests.WebSocket;
|
||||
public class WebSocketOptionsTests
|
||||
{
|
||||
[Fact]
|
||||
public void DefaultOptions_PortIsZero_Disabled()
|
||||
public void DefaultOptions_PortIsNegativeOne_Disabled()
|
||||
{
|
||||
var opts = new WebSocketOptions();
|
||||
opts.Port.ShouldBe(0);
|
||||
opts.Port.ShouldBe(-1);
|
||||
opts.Host.ShouldBe("0.0.0.0");
|
||||
opts.Compression.ShouldBeFalse();
|
||||
opts.NoTls.ShouldBeFalse();
|
||||
@@ -21,6 +21,6 @@ public class WebSocketOptionsTests
|
||||
{
|
||||
var opts = new NatsOptions();
|
||||
opts.WebSocket.ShouldNotBeNull();
|
||||
opts.WebSocket.Port.ShouldBe(0);
|
||||
opts.WebSocket.Port.ShouldBe(-1);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -75,6 +75,42 @@ public class WsConnectionTests
|
||||
(written[0] & WsConstants.Rsv1Bit).ShouldBe(0);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task ReadAsync_DecodesMaskedFrame()
|
||||
{
|
||||
var payload = "CONNECT {}\r\n"u8.ToArray();
|
||||
var (header, _) = WsFrameWriter.CreateFrameHeader(
|
||||
useMasking: true, compressed: false,
|
||||
opcode: WsConstants.BinaryMessage, payloadLength: payload.Length);
|
||||
var maskKey = header[^4..];
|
||||
WsFrameWriter.MaskBuf(maskKey, payload);
|
||||
|
||||
var frame = new byte[header.Length + payload.Length];
|
||||
header.CopyTo(frame, 0);
|
||||
payload.CopyTo(frame, header.Length);
|
||||
|
||||
var inner = new MemoryStream(frame);
|
||||
var ws = new WsConnection(inner, compress: false, maskRead: true, maskWrite: false, browser: false, noCompFrag: false);
|
||||
|
||||
var buf = new byte[256];
|
||||
int n = await ws.ReadAsync(buf);
|
||||
|
||||
n.ShouldBe("CONNECT {}\r\n".Length);
|
||||
System.Text.Encoding.ASCII.GetString(buf, 0, n).ShouldBe("CONNECT {}\r\n");
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task ReadAsync_ReturnsZero_OnEndOfStream()
|
||||
{
|
||||
// Empty stream should return 0 (true end of stream)
|
||||
var inner = new MemoryStream([]);
|
||||
var ws = new WsConnection(inner, compress: false, maskRead: false, maskWrite: false, browser: false, noCompFrag: false);
|
||||
|
||||
var buf = new byte[256];
|
||||
int n = await ws.ReadAsync(buf);
|
||||
n.ShouldBe(0);
|
||||
}
|
||||
|
||||
private static byte[] BuildUnmaskedFrame(byte[] payload)
|
||||
{
|
||||
var header = new byte[2];
|
||||
|
||||
162
tests/NATS.Server.Tests/WebSocket/WsIntegrationTests.cs
Normal file
162
tests/NATS.Server.Tests/WebSocket/WsIntegrationTests.cs
Normal file
@@ -0,0 +1,162 @@
|
||||
using System.Buffers.Binary;
|
||||
using System.Net;
|
||||
using System.Net.Sockets;
|
||||
using System.Security.Cryptography;
|
||||
using System.Text;
|
||||
using NATS.Server.WebSocket;
|
||||
|
||||
namespace NATS.Server.Tests.WebSocket;
|
||||
|
||||
public class WsIntegrationTests : IAsyncLifetime
|
||||
{
|
||||
private NatsServer _server = null!;
|
||||
private NatsOptions _options = null!;
|
||||
|
||||
public async Task InitializeAsync()
|
||||
{
|
||||
_options = new NatsOptions
|
||||
{
|
||||
Port = 0,
|
||||
WebSocket = new WebSocketOptions { Port = 0, NoTls = true },
|
||||
};
|
||||
var loggerFactory = Microsoft.Extensions.Logging.LoggerFactory.Create(b => { });
|
||||
_server = new NatsServer(_options, loggerFactory);
|
||||
_ = _server.StartAsync(CancellationToken.None);
|
||||
await _server.WaitForReadyAsync();
|
||||
}
|
||||
|
||||
public async Task DisposeAsync()
|
||||
{
|
||||
await _server.ShutdownAsync();
|
||||
_server.Dispose();
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task WebSocket_ConnectAndReceiveInfo()
|
||||
{
|
||||
using var socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
|
||||
await socket.ConnectAsync(new IPEndPoint(IPAddress.Loopback, _options.WebSocket.Port));
|
||||
using var stream = new NetworkStream(socket, ownsSocket: false);
|
||||
|
||||
await SendUpgradeRequest(stream);
|
||||
var response = await ReadHttpResponse(stream);
|
||||
response.ShouldContain("101");
|
||||
|
||||
var wsFrame = await ReadWsFrame(stream);
|
||||
var info = Encoding.ASCII.GetString(wsFrame);
|
||||
info.ShouldStartWith("INFO ");
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task WebSocket_ConnectAndPing()
|
||||
{
|
||||
using var client = await ConnectWsClient();
|
||||
|
||||
// Send CONNECT and PING together
|
||||
await SendWsText(client, "CONNECT {}\r\nPING\r\n");
|
||||
|
||||
// Read PONG WS frame
|
||||
using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
|
||||
var pong = await ReadWsFrameAsync(client, cts.Token);
|
||||
Encoding.ASCII.GetString(pong).ShouldContain("PONG");
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task WebSocket_PubSub()
|
||||
{
|
||||
using var sub = await ConnectWsClient();
|
||||
using var pub = await ConnectWsClient();
|
||||
|
||||
await SendWsText(sub, "CONNECT {}\r\nSUB test.ws 1\r\n");
|
||||
await Task.Delay(200);
|
||||
|
||||
await SendWsText(pub, "CONNECT {}\r\nPUB test.ws 5\r\nHello\r\n");
|
||||
|
||||
using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
|
||||
var msg = await ReadWsFrameAsync(sub, cts.Token);
|
||||
Encoding.ASCII.GetString(msg).ShouldContain("MSG test.ws 1 5");
|
||||
}
|
||||
|
||||
private async Task<NetworkStream> ConnectWsClient()
|
||||
{
|
||||
var socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
|
||||
await socket.ConnectAsync(new IPEndPoint(IPAddress.Loopback, _options.WebSocket.Port));
|
||||
var stream = new NetworkStream(socket, ownsSocket: true);
|
||||
|
||||
await SendUpgradeRequest(stream);
|
||||
var response = await ReadHttpResponse(stream);
|
||||
response.ShouldContain("101");
|
||||
|
||||
await ReadWsFrame(stream); // Read INFO frame
|
||||
return stream;
|
||||
}
|
||||
|
||||
private static async Task SendUpgradeRequest(NetworkStream stream)
|
||||
{
|
||||
var keyBytes = new byte[16];
|
||||
RandomNumberGenerator.Fill(keyBytes);
|
||||
var key = Convert.ToBase64String(keyBytes);
|
||||
|
||||
var request = $"GET / HTTP/1.1\r\nHost: localhost\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n\r\n";
|
||||
await stream.WriteAsync(Encoding.ASCII.GetBytes(request));
|
||||
await stream.FlushAsync();
|
||||
}
|
||||
|
||||
private static async Task<string> ReadHttpResponse(NetworkStream stream)
|
||||
{
|
||||
// Read one byte at a time to avoid consuming WS frame bytes that follow the HTTP response
|
||||
var sb = new StringBuilder();
|
||||
var buf = new byte[1];
|
||||
while (true)
|
||||
{
|
||||
int n = await stream.ReadAsync(buf);
|
||||
if (n == 0) break;
|
||||
sb.Append((char)buf[0]);
|
||||
if (sb.Length >= 4 &&
|
||||
sb[^4] == '\r' && sb[^3] == '\n' &&
|
||||
sb[^2] == '\r' && sb[^1] == '\n')
|
||||
break;
|
||||
}
|
||||
|
||||
return sb.ToString();
|
||||
}
|
||||
|
||||
private static Task<byte[]> ReadWsFrame(NetworkStream stream)
|
||||
=> ReadWsFrameAsync(stream, CancellationToken.None);
|
||||
|
||||
private static async Task<byte[]> ReadWsFrameAsync(NetworkStream stream, CancellationToken ct)
|
||||
{
|
||||
var header = new byte[2];
|
||||
await stream.ReadExactlyAsync(header, ct);
|
||||
int len = header[1] & 0x7F;
|
||||
if (len == 126)
|
||||
{
|
||||
var extLen = new byte[2];
|
||||
await stream.ReadExactlyAsync(extLen, ct);
|
||||
len = BinaryPrimitives.ReadUInt16BigEndian(extLen);
|
||||
}
|
||||
else if (len == 127)
|
||||
{
|
||||
var extLen = new byte[8];
|
||||
await stream.ReadExactlyAsync(extLen, ct);
|
||||
len = (int)BinaryPrimitives.ReadUInt64BigEndian(extLen);
|
||||
}
|
||||
|
||||
var payload = new byte[len];
|
||||
if (len > 0) await stream.ReadExactlyAsync(payload, ct);
|
||||
return payload;
|
||||
}
|
||||
|
||||
private static async Task SendWsText(NetworkStream stream, string text)
|
||||
{
|
||||
var payload = Encoding.ASCII.GetBytes(text);
|
||||
var (header, _) = WsFrameWriter.CreateFrameHeader(
|
||||
useMasking: true, compressed: false,
|
||||
opcode: WsConstants.BinaryMessage, payloadLength: payload.Length);
|
||||
var maskKey = header[^4..];
|
||||
WsFrameWriter.MaskBuf(maskKey, payload);
|
||||
await stream.WriteAsync(header);
|
||||
await stream.WriteAsync(payload);
|
||||
await stream.FlushAsync();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user