diff --git a/src/Microsoft.Extensions.WebSockets/IWebSocketConnection.cs b/src/Microsoft.Extensions.WebSockets/IWebSocketConnection.cs
index fa53a74d49..a1d0f50df4 100644
--- a/src/Microsoft.Extensions.WebSockets/IWebSocketConnection.cs
+++ b/src/Microsoft.Extensions.WebSockets/IWebSocketConnection.cs
@@ -48,8 +48,9 @@ namespace Microsoft.Extensions.WebSockets
/// Runs the WebSocket receive loop, using the provided message handler.
///
/// The callback that will be invoked for each new frame
+ /// A state parameter that will be passed to each invocation of
/// A that will complete when the client has sent a close frame, or the connection has been terminated
- Task ExecuteAsync(Func messageHandler);
+ Task ExecuteAsync(Func messageHandler, object state);
}
public static class WebSocketConnectionExtensions
@@ -74,9 +75,28 @@ namespace Microsoft.Extensions.WebSockets
/// The callback that will be invoked for each new frame
/// A that will complete when the client has sent a close frame, or the connection has been terminated
public static Task ExecuteAsync(this IWebSocketConnection self, Action messageHandler) =>
- self.ExecuteAsync(frame => {
+ self.ExecuteAsync((frame, _) => {
messageHandler(frame);
return Task.CompletedTask;
- });
+ }, null);
+
+ ///
+ /// Runs the WebSocket receive loop, using the provided message handler.
+ ///
+ /// The callback that will be invoked for each new frame
+ /// A that will complete when the client has sent a close frame, or the connection has been terminated
+ public static Task ExecuteAsync(this IWebSocketConnection self, Action messageHandler, object state) =>
+ self.ExecuteAsync((frame, s) => {
+ messageHandler(frame, s);
+ return Task.CompletedTask;
+ }, state);
+
+ ///
+ /// Runs the WebSocket receive loop, using the provided message handler.
+ ///
+ /// The callback that will be invoked for each new frame
+ /// A that will complete when the client has sent a close frame, or the connection has been terminated
+ public static Task ExecuteAsync(this IWebSocketConnection self, Func messageHandler) =>
+ self.ExecuteAsync((frame, _) => messageHandler(frame), null);
}
}
diff --git a/src/Microsoft.Extensions.WebSockets/WebSocketCloseResult.cs b/src/Microsoft.Extensions.WebSockets/WebSocketCloseResult.cs
index 6aa21bca3b..ad14a1b2db 100644
--- a/src/Microsoft.Extensions.WebSockets/WebSocketCloseResult.cs
+++ b/src/Microsoft.Extensions.WebSockets/WebSocketCloseResult.cs
@@ -34,13 +34,13 @@ namespace Microsoft.Extensions.WebSockets
public static bool TryParse(ReadableBuffer payload, out WebSocketCloseResult result)
{
- if(payload.Length == 0)
+ if (payload.Length == 0)
{
// Empty payload is OK
result = new WebSocketCloseResult(WebSocketCloseStatus.Empty, string.Empty);
return true;
}
- else if(payload.Length < 2)
+ else if (payload.Length < 2)
{
result = default(WebSocketCloseResult);
return false;
@@ -50,7 +50,7 @@ namespace Microsoft.Extensions.WebSockets
var status = payload.ReadBigEndian();
var description = string.Empty;
payload = payload.Slice(2);
- if(payload.Length > 0)
+ if (payload.Length > 0)
{
description = payload.GetUtf8String();
}
diff --git a/src/Microsoft.Extensions.WebSockets/WebSocketConnection.cs b/src/Microsoft.Extensions.WebSockets/WebSocketConnection.cs
index 1501da1dfc..4a04493fb8 100644
--- a/src/Microsoft.Extensions.WebSockets/WebSocketConnection.cs
+++ b/src/Microsoft.Extensions.WebSockets/WebSocketConnection.cs
@@ -79,7 +79,7 @@ namespace Microsoft.Extensions.WebSockets
_terminateReceiveCts.Cancel();
}
- public Task ExecuteAsync(Func messageHandler)
+ public Task ExecuteAsync(Func messageHandler, object state)
{
if (State == WebSocketConnectionState.Closed)
{
@@ -91,7 +91,7 @@ namespace Microsoft.Extensions.WebSockets
throw new InvalidOperationException("Connection is already running.");
}
State = WebSocketConnectionState.Connected;
- return Task.Run(() => ReceiveLoop(messageHandler, _terminateReceiveCts.Token));
+ return ReceiveLoop(messageHandler, state, _terminateReceiveCts.Token);
}
///
@@ -177,7 +177,7 @@ namespace Microsoft.Extensions.WebSockets
buffer.Set(_maskingKey);
}
- private async Task ReceiveLoop(Func messageHandler, CancellationToken cancellationToken)
+ private async Task ReceiveLoop(Func messageHandler, object state, CancellationToken cancellationToken)
{
while (!cancellationToken.IsCancellationRequested)
{
@@ -309,7 +309,7 @@ namespace Microsoft.Extensions.WebSockets
}
else
{
- await messageHandler(frame);
+ await messageHandler(frame, state);
}
// Mark the payload as consumed
@@ -376,6 +376,7 @@ namespace Microsoft.Extensions.WebSockets
// Allocate a buffer
var buffer = _outbound.Alloc(minimumSize: allocSize);
+ Debug.Assert(buffer.Memory.Length >= allocSize);
if (buffer.Memory.Length < allocSize)
{
throw new InvalidOperationException("Couldn't allocate enough data from the channel to write the header");