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");