This commit is contained in:
Stef Heyenrath
2026-02-22 21:48:33 +01:00
parent 8287ae79ec
commit ba3d1d758c
9 changed files with 629 additions and 98 deletions

View File

@@ -48,7 +48,6 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder, IGuidUtils gu
{
SubProtocol = builder.AcceptProtocol,
KeepAliveInterval = builder.KeepAliveIntervalSeconds ?? TimeSpan.FromSeconds(WebSocketConstants.DefaultKeepAliveIntervalSeconds)
};
var webSocket = await context.WebSockets.AcceptWebSocketAsync(acceptContext).ConfigureAwait(false);
#else
@@ -56,9 +55,7 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder, IGuidUtils gu
#endif
// Get or create registry from options
var registry = builder.IsBroadcast
? options.WebSocketRegistries.GetOrAdd(mapping.Guid, _ => new WebSocketConnectionRegistry())
: null;
var registry = options.WebSocketRegistries.GetOrAdd(mapping.Guid, _ => new WebSocketConnectionRegistry());
// Create WebSocket context
var wsContext = new WireMockWebSocketContext(
@@ -73,8 +70,8 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder, IGuidUtils gu
guidUtils
);
// Add to registry if broadcast is enabled
registry?.AddConnection(wsContext);
// Add to registry
registry.AddConnection(wsContext);
try
{
@@ -100,7 +97,7 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder, IGuidUtils gu
finally
{
// Remove from registry
registry?.RemoveConnection(wsContext.ConnectionId);
registry.RemoveConnection(wsContext.ConnectionId);
}
// Return special marker to indicate WebSocket was handled
@@ -213,9 +210,7 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder, IGuidUtils gu
}
}
private static async Task HandleCustomAsync(
WireMockWebSocketContext context,
Func<WebSocketMessage, IWebSocketContext, Task> handler)
private static async Task HandleCustomAsync(WireMockWebSocketContext context, Func<WebSocketMessage, IWebSocketContext, Task> handler)
{
var bufferSize = context.Builder.MaxMessageSize ?? WebSocketConstants.DefaultReceiveBufferSize;
using var buffer = ArrayPool<byte>.Shared.Lease(bufferSize);
@@ -236,10 +231,7 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder, IGuidUtils gu
try
{
var result = await context.WebSocket.ReceiveAsync(
new ArraySegment<byte>(buffer),
cts.Token
).ConfigureAwait(false);
var result = await context.WebSocket.ReceiveAsync(new ArraySegment<byte>(buffer), cts.Token).ConfigureAwait(false);
if (result.MessageType == WebSocketMessageType.Close)
{
@@ -257,10 +249,7 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder, IGuidUtils gu
context.LogWebSocketMessage(WebSocketMessageDirection.Receive, result.MessageType, null, receiveActivity);
await context.CloseAsync(
WebSocketCloseStatus.NormalClosure,
"Closed by client"
).ConfigureAwait(false);
await context.CloseAsync(WebSocketCloseStatus.NormalClosure, "Closed by client").ConfigureAwait(false);
break;
}
@@ -331,10 +320,7 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder, IGuidUtils gu
}
}
private static async Task ForwardMessagesAsync(
WireMockWebSocketContext context,
ClientWebSocket clientWebSocket,
WebSocketMessageDirection direction)
private static async Task ForwardMessagesAsync(WireMockWebSocketContext context, ClientWebSocket clientWebSocket, WebSocketMessageDirection direction)
{
using var buffer = ArrayPool<byte>.Shared.Lease(WebSocketConstants.ProxyForwardBufferSize);

View File

@@ -1,6 +1,5 @@
// Copyright © WireMock.Net
using System.Net.WebSockets;
using JetBrains.Annotations;
using WireMock.WebSockets;
@@ -32,17 +31,16 @@ public partial class WireMockServer
/// Close a specific WebSocket connection
/// </summary>
[PublicAPI]
public async Task CloseWebSocketConnectionAsync(
Guid connectionId,
WebSocketCloseStatus closeStatus = WebSocketCloseStatus.NormalClosure,
string statusDescription = "Closed by server",
CancellationToken cancellationToken = default)
public async Task AbortWebSocketConnectionAsync(Guid connectionId, string statusDescription = "Closed by server", CancellationToken cancellationToken = default)
{
foreach (var registry in _options.WebSocketRegistries.Values)
{
if (registry.TryGetConnection(connectionId, out var connection) && !cancellationToken.IsCancellationRequested)
if (registry.TryGetConnection(connectionId, out var connection))
{
await connection.CloseAsync(closeStatus, statusDescription, cancellationToken);
connection.Abort(statusDescription);
registry.RemoveConnection(connectionId);
await Task.Delay(100, cancellationToken); // Give the connection some time to close gracefully
return;
}
}
@@ -52,11 +50,11 @@ public partial class WireMockServer
/// Broadcast a text message to all WebSocket connections in a specific mapping
/// </summary>
[PublicAPI]
public async Task BroadcastToWebSocketsAsync(Guid mappingGuid, string text, Guid? excludeConnectionId = null, CancellationToken cancellationToken = default)
public async Task BroadcastToWebSocketsAsync(Guid mappingGuid, string text, CancellationToken cancellationToken = default)
{
if (_options.WebSocketRegistries.TryGetValue(mappingGuid, out var registry))
{
await registry.BroadcastAsync(text, excludeConnectionId, cancellationToken);
await registry.BroadcastAsync(text, null, cancellationToken);
}
}
@@ -64,11 +62,11 @@ public partial class WireMockServer
/// Broadcast a text message to all WebSocket connections
/// </summary>
[PublicAPI]
public async Task BroadcastToAllWebSocketsAsync(string text, Guid? excludeConnectionId = null, CancellationToken cancellationToken = default)
public async Task BroadcastToAllWebSocketsAsync(string text, CancellationToken cancellationToken = default)
{
foreach (var registry in _options.WebSocketRegistries.Values)
{
await registry.BroadcastAsync(text, excludeConnectionId, cancellationToken);
await registry.BroadcastAsync(text, null, cancellationToken);
}
}
@@ -76,11 +74,11 @@ public partial class WireMockServer
/// Broadcast a binary message to all WebSocket connections in a specific mapping
/// </summary>
[PublicAPI]
public async Task BroadcastToWebSocketsAsync(Guid mappingGuid, byte[] bytes, Guid? excludeConnectionId = null, CancellationToken cancellationToken = default)
public async Task BroadcastToWebSocketsAsync(Guid mappingGuid, byte[] bytes, CancellationToken cancellationToken = default)
{
if (_options.WebSocketRegistries.TryGetValue(mappingGuid, out var registry))
{
await registry.BroadcastAsync(bytes, excludeConnectionId, cancellationToken);
await registry.BroadcastAsync(bytes, null, cancellationToken);
}
}
@@ -88,11 +86,11 @@ public partial class WireMockServer
/// Broadcast a binary message to all WebSocket connections
/// </summary>
[PublicAPI]
public async Task BroadcastToAllWebSocketsAsync(byte[] bytes, Guid? excludeConnectionId = null, CancellationToken cancellationToken = default)
public async Task BroadcastToAllWebSocketsAsync(byte[] bytes, CancellationToken cancellationToken = default)
{
foreach (var registry in _options.WebSocketRegistries.Values)
{
await registry.BroadcastAsync(bytes, excludeConnectionId, cancellationToken);
await registry.BroadcastAsync(bytes, null, cancellationToken);
}
}
}

View File

@@ -13,34 +13,22 @@ internal class WebSocketBuilder(Response response) : IWebSocketBuilder
{
private readonly List<(IMatcher matcher, List<WebSocketMessageBuilder> messages)> _conditionalMessages = [];
/// <inheritdoc />
public string? AcceptProtocol { get; private set; }
/// <inheritdoc />
public bool IsEcho { get; private set; }
/// <inheritdoc />
public bool IsBroadcast { get; private set; }
/// <inheritdoc />
public Func<WebSocketMessage, IWebSocketContext, Task>? MessageHandler { get; private set; }
/// <inheritdoc />
public ProxyAndRecordSettings? ProxySettings { get; private set; }
/// <inheritdoc />
public TimeSpan? CloseTimeout { get; private set; }
/// <inheritdoc />
public int? MaxMessageSize { get; private set; }
/// <inheritdoc />
public int? ReceiveBufferSize { get; private set; }
/// <inheritdoc />
public TimeSpan? KeepAliveIntervalSeconds { get; private set; }
/// <inheritdoc />
public IWebSocketBuilder WithAcceptProtocol(string protocol)
{
AcceptProtocol = Guard.NotNull(protocol);
@@ -117,12 +105,6 @@ internal class WebSocketBuilder(Response response) : IWebSocketBuilder
return this;
}
public IWebSocketBuilder WithBroadcast()
{
IsBroadcast = true;
return this;
}
public IWebSocketBuilder WithProxy(ProxyAndRecordSettings settings)
{
ProxySettings = Guard.NotNull(settings);

View File

@@ -3,7 +3,6 @@
using System.Collections.Concurrent;
using System.Diagnostics.CodeAnalysis;
using System.Net.WebSockets;
using static System.Net.Mime.MediaTypeNames;
namespace WireMock.WebSockets;
@@ -27,7 +26,7 @@ internal class WebSocketConnectionRegistry
/// </summary>
public void RemoveConnection(Guid connectionId)
{
_connections.TryRemove(connectionId, out _);
_ = _connections.TryRemove(connectionId, out _);
}
/// <summary>
@@ -67,6 +66,6 @@ internal class WebSocketConnectionRegistry
private IEnumerable<WireMockWebSocketContext> Filter(Guid? excludeConnectionId)
{
return _connections.Values
.Where(c =>c.WebSocket.State == WebSocketState.Open && (!excludeConnectionId.HasValue || c.ConnectionId != excludeConnectionId));
.Where(c => c.WebSocket.State == WebSocketState.Open && (!excludeConnectionId.HasValue || c.ConnectionId != excludeConnectionId));
}
}

View File

@@ -1,5 +1,6 @@
// Copyright © WireMock.Net
using System.Buffers;
using System.Diagnostics;
using System.Net.WebSockets;
using System.Text;
@@ -33,7 +34,7 @@ public class WireMockWebSocketContext : IWebSocketContext
/// <inheritdoc />
public IMapping Mapping { get; }
internal WebSocketConnectionRegistry? Registry { get; }
internal WebSocketConnectionRegistry Registry { get; }
internal WebSocketBuilder Builder { get; }
@@ -49,7 +50,7 @@ public class WireMockWebSocketContext : IWebSocketContext
WebSocket webSocket,
IRequestMessage requestMessage,
IMapping mapping,
WebSocketConnectionRegistry? registry,
WebSocketConnectionRegistry registry,
WebSocketBuilder builder,
IWireMockMiddlewareOptions options,
IWireMockMiddlewareLogger logger,
@@ -96,29 +97,31 @@ public class WireMockWebSocketContext : IWebSocketContext
/// <inheritdoc />
public async Task CloseAsync(WebSocketCloseStatus closeStatus, string statusDescription, CancellationToken cancellationToken = default)
{
await WebSocket.CloseAsync(closeStatus, statusDescription, cancellationToken);
await WebSocket.CloseAsync(closeStatus, statusDescription, cancellationToken).ConfigureAwait(false);
LogWebSocketMessage(WebSocketMessageDirection.Send, WebSocketMessageType.Close, $"CloseStatus: {closeStatus}, Description: {statusDescription}", null);
}
/// <inheritdoc />
public void Abort(string? statusDescription = null)
{
WebSocket.Abort();
LogWebSocketMessage(WebSocketMessageDirection.Send, WebSocketMessageType.Close, $"CloseStatus: Abort, Description: {statusDescription}", null);
}
/// <inheritdoc />
public async Task BroadcastAsync(string text, bool excludeSender = false, CancellationToken cancellationToken = default)
{
if (Registry != null)
{
Guid? excludeConnectionId = excludeSender ? ConnectionId : null;
await Registry.BroadcastAsync(text, excludeConnectionId, cancellationToken);
}
Guid? excludeConnectionId = excludeSender ? ConnectionId : null;
await Registry.BroadcastAsync(text, excludeConnectionId, cancellationToken);
}
/// <inheritdoc />
public async Task BroadcastAsync(byte[] bytes, bool excludeSender = false, CancellationToken cancellationToken = default)
{
if (Registry != null)
{
Guid? excludeConnectionId = excludeSender ? ConnectionId : null;
await Registry.BroadcastAsync(bytes, excludeConnectionId, cancellationToken);
}
Guid? excludeConnectionId = excludeSender ? ConnectionId : null;
await Registry.BroadcastAsync(bytes, excludeConnectionId, cancellationToken);
}
internal void LogWebSocketMessage(

View File

@@ -64,12 +64,6 @@ public interface IWebSocketBuilder
[PublicAPI]
IWebSocketBuilder WithMessageHandler(Func<WebSocketMessage, IWebSocketContext, Task> handler);
/// <summary>
/// Enable broadcast mode for this mapping
/// </summary>
[PublicAPI]
IWebSocketBuilder WithBroadcast();
/// <summary>
/// Proxy to another WebSocket server
/// </summary>

View File

@@ -50,6 +50,11 @@ public interface IWebSocketContext
/// </summary>
Task CloseAsync(WebSocketCloseStatus closeStatus, string statusDescription, CancellationToken cancellationToken = default);
/// <summary>
/// Abort the WebSocket connection to immediately close the connection without waiting for the close handshake
/// </summary>
void Abort(string? statusDescription = null);
/// <summary>
/// Broadcast text message to all connections in this mapping
/// </summary>