(Very flawed) Implementation of System.Net.WebSockets.WebSocket over Owin async func spec

OwinWebSocket.cs

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() {
        }
    }
}

OwinWebSocketMessageTypes.cs

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;
    }
}