This commit is contained in:
Stef Heyenrath
2026-02-16 09:37:51 +01:00
parent f29a281b55
commit fa3a33dcc6
4 changed files with 141 additions and 224 deletions

View File

@@ -50,5 +50,20 @@ public enum BodyType
/// <summary>
/// Use Server-Sent Events (string)
/// </summary>
SseString
SseString,
/// <summary>
/// WebSocket message in clear text.
/// </summary>
WebSocketText,
/// <summary>
/// WebSocket message in binary format.
/// </summary>
WebSocketBinary,
/// <summary>
/// WebSocket close message.
/// </summary>
WebSocketClose
}

View File

@@ -2,21 +2,14 @@
using System.Buffers;
using System.Diagnostics;
using System.Drawing;
using System.Net;
using System.Net.WebSockets;
using System.Text;
using Microsoft.AspNetCore.Http;
using WireMock.Constants;
using WireMock.Logging;
using WireMock.Matchers;
using WireMock.Matchers.Request;
using WireMock.Models;
using WireMock.Owin;
using WireMock.Owin.ActivityTracing;
using WireMock.Settings;
using WireMock.Types;
using WireMock.Util;
using WireMock.WebSockets;
namespace WireMock.ResponseProviders;
@@ -175,7 +168,7 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder) : IResponsePr
);
}
LogWebSocketMessage(context, WebSocketMessageDirection.Receive, result.MessageType, null, receiveActivity);
context.LogWebSocketMessage(WebSocketMessageDirection.Receive, result.MessageType, null, receiveActivity);
await context.CloseAsync(
WebSocketCloseStatus.NormalClosure,
@@ -204,7 +197,7 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder) : IResponsePr
}
// Log the receive operation
LogWebSocketMessage(context, WebSocketMessageDirection.Receive, result.MessageType, textContent, receiveActivity);
context.LogWebSocketMessage(WebSocketMessageDirection.Receive, result.MessageType, textContent, receiveActivity);
// Echo back (this will be logged by context.SendAsync)
await context.WebSocket.SendAsync(
@@ -276,7 +269,7 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder) : IResponsePr
);
}
LogWebSocketMessage(context, WebSocketMessageDirection.Receive, result.MessageType, null, receiveActivity);
context.LogWebSocketMessage(WebSocketMessageDirection.Receive, result.MessageType, null, receiveActivity);
await context.CloseAsync(
WebSocketCloseStatus.NormalClosure,
@@ -302,7 +295,7 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder) : IResponsePr
// Log the receive operation
object? data = message.Text != null ? message.Text : message.Bytes;
LogWebSocketMessage(context, WebSocketMessageDirection.Receive, result.MessageType, data, receiveActivity);
context.LogWebSocketMessage(WebSocketMessageDirection.Receive, result.MessageType, data, receiveActivity);
// Call custom handler
await handler(message, context).ConfigureAwait(false);
@@ -396,7 +389,7 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder) : IResponsePr
);
}
LogWebSocketMessage(context, direction, result.MessageType, null, activity);
context.LogWebSocketMessage(direction, result.MessageType, null, activity);
await destination.CloseAsync(
result.CloseStatus ?? WebSocketCloseStatus.NormalClosure,
@@ -431,7 +424,7 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder) : IResponsePr
}
// Log the proxy operation
LogWebSocketMessage(context, direction, result.MessageType, data, activity);
context.LogWebSocketMessage(direction, result.MessageType, data, activity);
await destination.SendAsync(
new ArraySegment<byte>(buffer, 0, result.Count),
@@ -501,7 +494,7 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder) : IResponsePr
Array.Copy(buffer, (byte[])data, result.Count);
}
LogWebSocketMessage(context, WebSocketMessageDirection.Receive, result.MessageType, data, receiveActivity);
context.LogWebSocketMessage(WebSocketMessageDirection.Receive, result.MessageType, data, receiveActivity);
if (result.MessageType == WebSocketMessageType.Close)
{
@@ -529,101 +522,6 @@ internal class WebSocketResponseProvider(WebSocketBuilder builder) : IResponsePr
}
}
private static void LogWebSocketMessage(
WireMockWebSocketContext context,
WebSocketMessageDirection direction,
WebSocketMessageType messageType,
object? data,
Activity? activity)
{
// Skip logging if log count limit is disabled
if (context.Options.MaxRequestLogCount == 0)
{
return;
}
// Create body data
IBodyData bodyData;
if (messageType == WebSocketMessageType.Text && data is string textContent)
{
bodyData = new BodyData
{
BodyAsString = textContent,
DetectedBodyType = BodyType.String
};
}
else if (messageType == WebSocketMessageType.Binary && data is byte[] binary)
{
bodyData = new BodyData
{
BodyAsBytes = binary,
DetectedBodyType = BodyType.Bytes
};
}
else
{
bodyData = new BodyData
{
BodyAsString = messageType.ToString(),
DetectedBodyType = BodyType.String
};
}
// Create a pseudo-request or pseudo-response depending on direction
RequestMessage? requestMessage = null;
IResponseMessage? responseMessage = null;
var method = $"WS_{direction.ToString().ToUpperInvariant()}";
if (direction == WebSocketMessageDirection.Receive)
{
// Received message - log as request
requestMessage = new RequestMessage(
new UrlDetails(context.RequestMessage.Url),
method,
context.RequestMessage.ClientIP,
bodyData,
null,
null
)
{
DateTime = DateTime.UtcNow
};
}
else
{
// Sent message - log as response
responseMessage = new ResponseMessage
{
Method = method,
StatusCode = HttpStatusCode.SwitchingProtocols, // WebSocket status
BodyData = bodyData,
DateTime = DateTime.UtcNow
};
}
// Create log entry
var logEntry = new LogEntry
{
Guid = Guid.NewGuid(),
RequestMessage = requestMessage,
ResponseMessage = responseMessage,
MappingGuid = context.Mapping.Guid,
MappingTitle = context.Mapping.Title
};
// Enrich activity if present
if (activity != null && context.Options.ActivityTracingOptions != null)
{
WireMockActivitySource.EnrichWithLogEntry(activity, logEntry, context.Options.ActivityTracingOptions);
}
// Log using LogLogEntry
context.Logger.LogLogEntry(logEntry, context.Options.MaxRequestLogCount is null or > 0);
activity?.Dispose();
}
private static WebSocketMessage CreateWebSocketMessage(WebSocketReceiveResult result, byte[] buffer)
{
var message = new WebSocketMessage

View File

@@ -42,6 +42,8 @@ internal class LogEntryMapper(IWireMockMiddlewareOptions options)
{
case BodyType.String:
case BodyType.FormUrlEncoded:
case BodyType.WebSocketText:
case BodyType.WebSocketClose:
logRequestModel.Body = logEntry.RequestMessage.BodyData.BodyAsString;
break;
@@ -51,6 +53,7 @@ internal class LogEntryMapper(IWireMockMiddlewareOptions options)
break;
case BodyType.Bytes:
case BodyType.WebSocketBinary:
logRequestModel.BodyAsBytes = logEntry.RequestMessage.BodyData.BodyAsBytes;
break;
}
@@ -126,6 +129,8 @@ internal class LogEntryMapper(IWireMockMiddlewareOptions options)
{
case BodyType.String:
case BodyType.FormUrlEncoded:
case BodyType.WebSocketText:
case BodyType.WebSocketClose:
if (!string.IsNullOrEmpty(logEntry.ResponseMessage.BodyData.IsFuncUsed) && options.DoNotSaveDynamicResponseInLogEntry == true)
{
logResponseModel.Body = logEntry.ResponseMessage.BodyData.IsFuncUsed;
@@ -141,6 +146,7 @@ internal class LogEntryMapper(IWireMockMiddlewareOptions options)
break;
case BodyType.Bytes:
case BodyType.WebSocketBinary:
logResponseModel.BodyAsBytes = logEntry.ResponseMessage.BodyData.BodyAsBytes;
break;

View File

@@ -7,8 +7,6 @@ using System.Text;
using Microsoft.AspNetCore.Http;
using Stef.Validation;
using WireMock.Logging;
using WireMock.Matchers;
using WireMock.Matchers.Request;
using WireMock.Models;
using WireMock.Owin;
using WireMock.Owin.ActivityTracing;
@@ -93,51 +91,97 @@ public class WireMockWebSocketContext : IWebSocketContext
);
}
private async Task SendAsyncInternal(
ArraySegment<byte> buffer,
WebSocketMessageType messageType,
bool endOfMessage,
object? data,
CancellationToken cancellationToken)
/// <inheritdoc />
public async Task CloseAsync(WebSocketCloseStatus closeStatus, string statusDescription)
{
Activity? activity = null;
var shouldTrace = Options.ActivityTracingOptions is not null;
await WebSocket.CloseAsync(closeStatus, statusDescription, CancellationToken.None);
if (shouldTrace)
LogWebSocketMessage(WebSocketMessageDirection.Send, WebSocketMessageType.Close, $"CloseStatus: {closeStatus}, Description: {statusDescription}", null);
}
/// <inheritdoc />
public void SetScenarioState(string nextState)
{
SetScenarioState(nextState, null);
}
/// <inheritdoc />
public void SetScenarioState(string nextState, string? description)
{
if (Mapping.Scenario == null)
{
activity = WireMockActivitySource.StartWebSocketMessageActivity(WebSocketMessageDirection.Send, Mapping.Guid);
WireMockActivitySource.EnrichWithWebSocketMessage(
activity,
messageType,
buffer.Count,
endOfMessage,
data as string,
Options.ActivityTracingOptions
);
return;
}
try
// Use the same logic as WireMockMiddleware
if (Options.Scenarios.TryGetValue(Mapping.Scenario, out var scenarioState))
{
await WebSocket.SendAsync(buffer, messageType, endOfMessage, cancellationToken).ConfigureAwait(false);
// Directly set the next state (bypass counter logic for manual WebSocket state changes)
scenarioState.NextState = nextState;
scenarioState.Started = true;
scenarioState.Finished = nextState == null;
// Log the send operation
if (Options.MaxRequestLogCount is null or > 0)
// Reset counter when manually setting state
scenarioState.Counter = 0;
}
else
{
// Create new scenario state if it doesn't exist
Options.Scenarios.TryAdd(Mapping.Scenario, new ScenarioState
{
LogWebSocketMessage(WebSocketMessageDirection.Send, messageType, data, activity);
}
}
catch (Exception ex)
{
WireMockActivitySource.RecordException(activity, ex);
throw;
}
finally
{
activity?.Dispose();
Name = Mapping.Scenario,
NextState = nextState,
Started = true,
Finished = nextState == null,
Counter = 0
});
}
}
private void LogWebSocketMessage(
/// <inheritdoc />
public async Task BroadcastTextAsync(string text, CancellationToken cancellationToken = default)
{
if (Registry != null)
{
await Registry.BroadcastTextAsync(text, cancellationToken);
}
}
/// <summary>
/// Update scenario state following the same pattern as WireMockMiddleware.UpdateScenarioState
/// This is called automatically when the WebSocket connection is established.
/// </summary>
internal void UpdateScenarioState()
{
if (Mapping.Scenario == null)
{
return;
}
// Ensure scenario exists
if (!Options.Scenarios.TryGetValue(Mapping.Scenario, out var scenario))
{
return;
}
// Follow exact same logic as WireMockMiddleware.UpdateScenarioState
// Increase the number of times this state has been executed
scenario.Counter++;
// Only if the number of times this state is executed equals the required StateTimes,
// proceed to next state and reset the counter to 0
if (scenario.Counter == (Mapping.TimesInSameState ?? 1))
{
scenario.NextState = Mapping.NextState;
scenario.Counter = 0;
}
// Else just update Started and Finished
scenario.Started = true;
scenario.Finished = Mapping.NextState == null;
}
internal void LogWebSocketMessage(
WebSocketMessageDirection direction,
WebSocketMessageType messageType,
object? data,
@@ -150,7 +194,7 @@ public class WireMockWebSocketContext : IWebSocketContext
bodyData = new BodyData
{
BodyAsString = textContent,
DetectedBodyType = BodyType.String
DetectedBodyType = BodyType.WebSocketText
};
}
else if (messageType == WebSocketMessageType.Binary && data is byte[] binary)
@@ -158,7 +202,7 @@ public class WireMockWebSocketContext : IWebSocketContext
bodyData = new BodyData
{
BodyAsBytes = binary,
DetectedBodyType = BodyType.Bytes
DetectedBodyType = BodyType.WebSocketBinary
};
}
else
@@ -166,7 +210,7 @@ public class WireMockWebSocketContext : IWebSocketContext
bodyData = new BodyData
{
BodyAsString = messageType.ToString(),
DetectedBodyType = BodyType.String
DetectedBodyType = BodyType.WebSocketClose
};
}
@@ -225,93 +269,47 @@ public class WireMockWebSocketContext : IWebSocketContext
activity?.Dispose();
}
/// <inheritdoc />
public async Task CloseAsync(WebSocketCloseStatus closeStatus, string statusDescription)
private async Task SendAsyncInternal(
ArraySegment<byte> buffer,
WebSocketMessageType messageType,
bool endOfMessage,
object? data,
CancellationToken cancellationToken)
{
await WebSocket.CloseAsync(closeStatus, statusDescription, CancellationToken.None);
Activity? activity = null;
var shouldTrace = Options.ActivityTracingOptions is not null;
LogWebSocketMessage(WebSocketMessageDirection.Send, WebSocketMessageType.Close, $"CloseStatus: {closeStatus}, Description: {statusDescription}", null);
}
/// <inheritdoc />
public void SetScenarioState(string nextState)
{
SetScenarioState(nextState, null);
}
/// <inheritdoc />
public void SetScenarioState(string nextState, string? description)
{
if (Mapping.Scenario == null)
if (shouldTrace)
{
return;
activity = WireMockActivitySource.StartWebSocketMessageActivity(WebSocketMessageDirection.Send, Mapping.Guid);
WireMockActivitySource.EnrichWithWebSocketMessage(
activity,
messageType,
buffer.Count,
endOfMessage,
data as string,
Options.ActivityTracingOptions
);
}
// Use the same logic as WireMockMiddleware
if (Options.Scenarios.TryGetValue(Mapping.Scenario, out var scenarioState))
try
{
// Directly set the next state (bypass counter logic for manual WebSocket state changes)
scenarioState.NextState = nextState;
scenarioState.Started = true;
scenarioState.Finished = nextState == null;
await WebSocket.SendAsync(buffer, messageType, endOfMessage, cancellationToken).ConfigureAwait(false);
// Reset counter when manually setting state
scenarioState.Counter = 0;
}
else
{
// Create new scenario state if it doesn't exist
Options.Scenarios.TryAdd(Mapping.Scenario, new ScenarioState
// Log the send operation
if (Options.MaxRequestLogCount is null or > 0)
{
Name = Mapping.Scenario,
NextState = nextState,
Started = true,
Finished = nextState == null,
Counter = 0
});
LogWebSocketMessage(WebSocketMessageDirection.Send, messageType, data, activity);
}
}
}
/// <summary>
/// Update scenario state following the same pattern as WireMockMiddleware.UpdateScenarioState
/// This is called automatically when the WebSocket connection is established.
/// </summary>
internal void UpdateScenarioState()
{
if (Mapping.Scenario == null)
catch (Exception ex)
{
return;
WireMockActivitySource.RecordException(activity, ex);
throw;
}
// Ensure scenario exists
if (!Options.Scenarios.TryGetValue(Mapping.Scenario, out var scenario))
finally
{
return;
}
// Follow exact same logic as WireMockMiddleware.UpdateScenarioState
// Increase the number of times this state has been executed
scenario.Counter++;
// Only if the number of times this state is executed equals the required StateTimes,
// proceed to next state and reset the counter to 0
if (scenario.Counter == (Mapping.TimesInSameState ?? 1))
{
scenario.NextState = Mapping.NextState;
scenario.Counter = 0;
}
// Else just update Started and Finished
scenario.Started = true;
scenario.Finished = Mapping.NextState == null;
}
/// <inheritdoc />
public async Task BroadcastTextAsync(string text, CancellationToken cancellationToken = default)
{
if (Registry != null)
{
await Registry.BroadcastTextAsync(text, cancellationToken);
activity?.Dispose();
}
}
}