Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
(Very flawed) Implementation of System.Net.WebSockets.WebSocket over Owin async func spec
using System;
using System.Collections.Generic;
using System.Threading.Tasks;
using System.Net.WebSockets;
using System.Threading;
// This is bad and potentially unreliable code, which only serves as a base for
// further improvements. Please do not use it as is.
namespace MirrorSharp.Owin.Internal {
using WebSocketSendAsync = Func<
ArraySegment<byte> /* data */,
int /* messageType */,
bool /* endOfMessage */,
CancellationToken /* cancel */,
Task
>;
using WebSocketReceiveAsync = Func<
ArraySegment<byte> /* data */,
CancellationToken /* cancel */,
Task<Tuple<
int /* messageType */,
bool /* endOfMessage */,
int /* count */
>>
>;
using WebSocketCloseAsync = Func<
int /* closeStatus */,
string /* closeDescription */,
CancellationToken /* cancel */,
Task
>;
internal class OwinWebSocket : WebSocket {
private readonly IDictionary<string, object> _environment;
private readonly WebSocketSendAsync _sendAsync;
private readonly WebSocketReceiveAsync _receiveAsync;
private readonly WebSocketCloseAsync _closeAsync;
private readonly TaskCompletionSource<object> _abortedTaskSource;
private WebSocketState _state;
private WebSocketCloseStatus? _closeStatus;
private string _closeDescription;
public OwinWebSocket(IDictionary<string, object> environment) {
_environment = environment;
_state = WebSocketState.Open;
_sendAsync = (WebSocketSendAsync)environment["websocket.SendAsync"];
_receiveAsync = (WebSocketReceiveAsync)environment["websocket.ReceiveAsync"];
_closeAsync = (WebSocketCloseAsync)environment["websocket.CloseAsync"];
var callCancelledToken = (CancellationToken)environment["websocket.CallCancelled"];
callCancelledToken.Register(Abort);
_abortedTaskSource = new TaskCompletionSource<object>();
}
public override WebSocketCloseStatus? CloseStatus => _closeStatus;
public override string CloseStatusDescription => _closeDescription;
public override WebSocketState State => _state;
public Task AbortedTask => _abortedTaskSource.Task;
public override string SubProtocol {
get { throw new NotSupportedException(); }
}
public override async Task<WebSocketReceiveResult> ReceiveAsync(ArraySegment<byte> buffer, CancellationToken cancellationToken) {
if (_state != WebSocketState.Open && _state != WebSocketState.CloseSent)
throw new WebSocketException(WebSocketError.InvalidState);
int count;
WebSocketMessageType messageType;
bool endOfMessage;
try {
var tuple = await _receiveAsync(buffer, cancellationToken).ConfigureAwait(false);
messageType = MapMessageTypeToEnum(tuple.Item1);
if (_state == WebSocketState.CloseSent && messageType != WebSocketMessageType.Close)
throw new WebSocketException(WebSocketError.InvalidMessageType);
count = tuple.Item3;
endOfMessage = tuple.Item2;
}
catch (WebSocketException) {
Abort();
throw;
}
if (messageType == WebSocketMessageType.Close) {
if (_state == WebSocketState.Open) {
_state = WebSocketState.CloseReceived;
_closeStatus = (WebSocketCloseStatus?)(int?)_environment.GetValueOrDefault("websocket.ClientCloseStatus");
_closeDescription = (string)_environment.GetValueOrDefault("websocket.ClientCloseDescription");
return new WebSocketReceiveResult(
count, messageType, endOfMessage,
_closeStatus, _closeDescription
);
}
_state = WebSocketState.Closed;
}
return new WebSocketReceiveResult(count, messageType, endOfMessage);
}
public override async Task SendAsync(ArraySegment<byte> buffer, WebSocketMessageType messageType, bool endOfMessage, CancellationToken cancellationToken) {
if (_state != WebSocketState.Open && _state != WebSocketState.CloseReceived)
throw new WebSocketException(WebSocketError.InvalidState);
try {
await _sendAsync(buffer, MapMessageTypeFromEnum(messageType), endOfMessage, cancellationToken)
.ConfigureAwait(false);
}
catch (Exception ex) when (messageType == WebSocketMessageType.Close || ex is WebSocketException) {
Abort();
throw;
}
if (messageType == WebSocketMessageType.Close)
_state = (_state == WebSocketState.CloseReceived) ? WebSocketState.Closed : WebSocketState.CloseSent;
}
public override void Abort() {
if (_state == WebSocketState.Aborted)
return;
_state = WebSocketState.Aborted;
_abortedTaskSource.SetResult(null);
}
public override async Task CloseAsync(WebSocketCloseStatus closeStatus, string statusDescription, CancellationToken cancellationToken) {
if (_state != WebSocketState.Open && _state != WebSocketState.CloseReceived)
throw new WebSocketException(WebSocketError.InvalidState);
try {
await CloseOutputAsync(closeStatus, statusDescription, cancellationToken).ConfigureAwait(false);
await ReceiveAsync(new ArraySegment<byte>(new byte[123]), cancellationToken).ConfigureAwait(false);
}
catch (Exception) {
Abort();
throw;
}
}
public override async Task CloseOutputAsync(WebSocketCloseStatus closeStatus, string statusDescription, CancellationToken cancellationToken) {
if (_state != WebSocketState.Open && _state != WebSocketState.CloseReceived)
throw new WebSocketException(WebSocketError.InvalidState);
try {
await _closeAsync((int)closeStatus, statusDescription, cancellationToken).ConfigureAwait(false);
}
catch (Exception) {
Abort();
throw;
}
_closeStatus = closeStatus;
_closeDescription = statusDescription;
_state = (_state == WebSocketState.CloseReceived) ? WebSocketState.Closed : WebSocketState.CloseSent;
}
private int MapMessageTypeFromEnum(WebSocketMessageType messageType) {
switch (messageType) {
case WebSocketMessageType.Binary: return OwinWebSocketMessageTypes.Binary;
case WebSocketMessageType.Text: return OwinWebSocketMessageTypes.Text;
case WebSocketMessageType.Close: return OwinWebSocketMessageTypes.Close;
default: throw new ArgumentException($"Unknown message type: {messageType}.", nameof(messageType));
}
}
private WebSocketMessageType MapMessageTypeToEnum(int messageType) {
switch (messageType) {
case OwinWebSocketMessageTypes.Binary: return WebSocketMessageType.Binary;
case OwinWebSocketMessageTypes.Text: return WebSocketMessageType.Text;
case OwinWebSocketMessageTypes.Close: return WebSocketMessageType.Close;
default: throw new WebSocketException(WebSocketError.InvalidMessageType);
}
}
public override void Dispose() {
}
}
}
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;
namespace MirrorSharp.Owin.Internal {
internal static class OwinWebSocketMessageTypes {
public const int Text = 1;
public const int Binary = 2;
public const int Close = 8;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment