#96 Enable AspNet->Owin WebSockets.

src/Microsoft.AspNet.Owin/Microsoft.AspNet.Owin.kproj

@@ -29,6 +29,8 @@
     <Compile Include="WebSockets\OwinWebSocketAcceptContext.cs" />
     <Compile Include="WebSockets\OwinWebSocketAdapter.cs" />
     <Compile Include="WebSockets\OwinWebSocketAcceptAdapter.cs" />
+    <Compile Include="WebSockets\WebSocketAcceptAdapter.cs" />
+    <Compile Include="WebSockets\WebSocketAdapter.cs" />
   </ItemGroup>
   <Import Project="$(VSToolsPath)\AspNet\Microsoft.Web.AspNet.targets" Condition="'$(VSToolsPath)' != ''" />
 </Project>

src/Microsoft.AspNet.Owin/OwinConstants.cs

@@ -130,6 +130,7 @@ namespace Microsoft.AspNet.Owin
             // 3.1. Startup
 
             public const string Version = "websocket.Version";
+            public const string VersionValue = "1.0";
 
             // 3.2. Per Request
 

src/Microsoft.AspNet.Owin/OwinEnvironment.cs

@@ -8,6 +8,7 @@ using System.Globalization;
 using System.IO;
 using System.Linq;
 using System.Net;
+using System.Net.WebSockets;
 using System.Security.Cryptography.X509Certificates;
 using System.Security.Principal;
 using System.Threading;
@@ -20,6 +21,12 @@ using Microsoft.AspNet.PipelineCore.Security;
 namespace Microsoft.AspNet.Owin
 {
     using SendFileFunc = Func<string, long, long?, CancellationToken, Task>;
+    using WebSocketAcceptAlt =
+        Func
+        <
+            IWebSocketAcceptContext, // WebSocket Accept parameters
+            Task<WebSocket>
+        >;
 
     public class OwinEnvironment : IDictionary<string, object>
     {
@@ -76,6 +83,11 @@ namespace Microsoft.AspNet.Owin
                     feature => new Func<Task>(() => feature.GetClientCertificateAsync(CancellationToken.None))));
             }
 
+            if (context.IsWebSocketRequest)
+            {
+                _entries.Add(OwinConstants.WebSocket.AcceptAlt, new FeatureMap<IHttpWebSocketFeature>(feature => new WebSocketAcceptAlt(feature.AcceptAsync)));
+            }
+
             _context.Items[typeof(HttpContext).FullName] = _context; // Store for lookup when we transition back out of OWIN
         }
 

src/Microsoft.AspNet.Owin/WebSockets/OwinWebSocketAcceptContext.cs

@@ -8,18 +8,23 @@ namespace Microsoft.AspNet.Owin
 {
     public class OwinWebSocketAcceptContext : IWebSocketAcceptContext
     {
-        private IDictionary<string, object> _options = new Dictionary<string, object>(1);
+        private IDictionary<string, object> _options;
 
-	    public OwinWebSocketAcceptContext()
+	    public OwinWebSocketAcceptContext() : this(new Dictionary<string, object>(1))
 	    {
         }
 
+        public OwinWebSocketAcceptContext(IDictionary<string, object> options)
+        {
+            _options = options;
+        }
+
         public string SubProtocol
         {
             get
             {
                 object obj;
-                if (_options.TryGetValue(OwinConstants.WebSocket.SubProtocol, out obj))
+                if (_options != null && _options.TryGetValue(OwinConstants.WebSocket.SubProtocol, out obj))
                 {
                     return (string)obj;
                 }
@@ -27,6 +32,10 @@ namespace Microsoft.AspNet.Owin
             }
             set
             {
+                if (_options == null)
+                {
+                    _options = new Dictionary<string, object>(1);
+                }
                 _options[OwinConstants.WebSocket.SubProtocol] = value;
             }
         }

src/Microsoft.AspNet.Owin/WebSockets/WebSocketAcceptAdapter.cs

@@ -0,0 +1,88 @@
+// Copyright (c) Microsoft Open Technologies, Inc. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Net.WebSockets;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNet.HttpFeature;
+
+namespace Microsoft.AspNet.Owin
+{
+    using AppFunc = Func<IDictionary<string, object>, Task>;
+    using WebSocketAccept =
+        Action
+        <
+            IDictionary<string, object>, // WebSocket Accept parameters
+            Func // WebSocketFunc callback
+            <
+                IDictionary<string, object>, // WebSocket environment
+                Task // Complete
+            >
+        >;
+    using WebSocketAcceptAlt =
+        Func
+        <
+            IWebSocketAcceptContext, // WebSocket Accept parameters
+            Task<WebSocket>
+        >;
+
+    public class WebSocketAcceptAdapter
+    {
+        private IDictionary<string, object> _env;
+        private WebSocketAcceptAlt _accept;
+        private AppFunc _callback;
+        private IDictionary<string, object> _options;
+
+        public WebSocketAcceptAdapter(IDictionary<string, object> env, WebSocketAcceptAlt accept)
+	    {
+            _env = env;
+            _accept = accept;
+        }
+
+        private void AcceptWebSocket(IDictionary<string, object> options, AppFunc callback)
+        {
+            _options = options;
+            _callback = callback;
+            _env[OwinConstants.ResponseStatusCode] = 101;
+        }
+
+        public static AppFunc AdaptWebSockets(AppFunc next)
+        {
+            return async environment =>
+            {
+                object accept;
+                if (environment.TryGetValue(OwinConstants.WebSocket.AcceptAlt, out accept) && accept is WebSocketAcceptAlt)
+                {
+                    var adapter = new WebSocketAcceptAdapter(environment, (WebSocketAcceptAlt)accept);
+
+                    environment[OwinConstants.WebSocket.Accept] = new WebSocketAccept(adapter.AcceptWebSocket);
+                    await next(environment);
+                    if ((int)environment[OwinConstants.ResponseStatusCode] == 101 && adapter._callback != null)
+                    {
+                        IWebSocketAcceptContext acceptContext = null;
+                        object obj;
+                        if (adapter._options.TryGetValue(typeof(IWebSocketAcceptContext).FullName, out obj))
+                        {
+                            acceptContext = obj as IWebSocketAcceptContext;
+                        }
+                        else if (adapter._options != null)
+                        {
+                            acceptContext = new OwinWebSocketAcceptContext(adapter._options);
+                        }
+
+                        var webSocket = await adapter._accept(acceptContext);
+                        var webSocketAdapter = new WebSocketAdapter(webSocket, (CancellationToken)environment[OwinConstants.CallCancelled]);
+                        await adapter._callback(webSocketAdapter.Environment);
+                        await webSocketAdapter.CleanupAsync();
+                    }
+                }
+                else
+                {
+                    await next(environment);
+                }
+            };
+        }
+    }
+}

src/Microsoft.AspNet.Owin/WebSockets/WebSocketAdapter.cs

@@ -0,0 +1,172 @@
+// Copyright (c) Microsoft Open Technologies, Inc. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Net.WebSockets;
+using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNet.Owin
+{
+    using WebSocketCloseAsync =
+        Func<int /* closeStatus */,
+            string /* closeDescription */,
+            CancellationToken /* cancel */,
+            Task>;
+    using WebSocketReceiveAsync =
+        Func<ArraySegment<byte> /* data */,
+            CancellationToken /* cancel */,
+            Task<Tuple<int /* messageType */,
+                bool /* endOfMessage */,
+                int /* count */>>>;
+    using WebSocketReceiveTuple =
+        Tuple<int /* messageType */,
+            bool /* endOfMessage */,
+            int /* count */>;
+    using WebSocketSendAsync =
+        Func<ArraySegment<byte> /* data */,
+            int /* messageType */,
+            bool /* endOfMessage */,
+            CancellationToken /* cancel */,
+            Task>;
+
+    public class WebSocketAdapter
+    {
+        private readonly WebSocket _webSocket;
+        private readonly IDictionary<string, object> _environment;
+        private readonly CancellationToken _cancellationToken;
+        private readonly WebSocket _context;
+
+        internal WebSocketAdapter(WebSocket webSocket, CancellationToken ct)
+        {
+            _webSocket = webSocket;
+            _cancellationToken = ct;
+
+            _environment = new Dictionary<string, object>();
+            _environment[OwinConstants.WebSocket.SendAsync] = new WebSocketSendAsync(SendAsync);
+            _environment[OwinConstants.WebSocket.ReceiveAsync] = new WebSocketReceiveAsync(ReceiveAsync);
+            _environment[OwinConstants.WebSocket.CloseAsync] = new WebSocketCloseAsync(CloseAsync);
+            _environment[OwinConstants.WebSocket.CallCancelled] = ct;
+            _environment[OwinConstants.WebSocket.Version] = OwinConstants.WebSocket.VersionValue;
+
+            _environment[typeof(WebSocket).FullName] = webSocket;
+        }
+
+        internal IDictionary<string, object> Environment
+        {
+            get { return _environment; }
+        }
+
+        internal Task SendAsync(ArraySegment<byte> buffer, int messageType, bool endOfMessage, CancellationToken cancel)
+        {
+            // Remap close messages to CloseAsync.  System.Net.WebSockets.WebSocket.SendAsync does not allow close messages.
+            if (messageType == 0x8)
+            {
+                return RedirectSendToCloseAsync(buffer, cancel);
+            }
+            else if (messageType == 0x9 || messageType == 0xA)
+            {
+                // Ping & Pong, not allowed by the underlying APIs, silently discard.
+                return Task.FromResult(0);
+            }
+
+            return _webSocket.SendAsync(buffer, OpCodeToEnum(messageType), endOfMessage, cancel);
+        }
+
+        internal async Task<WebSocketReceiveTuple> ReceiveAsync(ArraySegment<byte> buffer, CancellationToken cancel)
+        {
+            WebSocketReceiveResult nativeResult = await _webSocket.ReceiveAsync(buffer, cancel);
+
+            if (nativeResult.MessageType == WebSocketMessageType.Close)
+            {
+                _environment[OwinConstants.WebSocket.ClientCloseStatus] = (int)(nativeResult.CloseStatus ?? WebSocketCloseStatus.NormalClosure);
+                _environment[OwinConstants.WebSocket.ClientCloseDescription] = nativeResult.CloseStatusDescription ?? string.Empty;
+            }
+
+            return new WebSocketReceiveTuple(
+                EnumToOpCode(nativeResult.MessageType),
+                nativeResult.EndOfMessage,
+                nativeResult.Count);
+        }
+
+        internal Task CloseAsync(int status, string description, CancellationToken cancel)
+        {
+            return _webSocket.CloseOutputAsync((WebSocketCloseStatus)status, description, cancel);
+        }
+
+        private Task RedirectSendToCloseAsync(ArraySegment<byte> buffer, CancellationToken cancel)
+        {
+            if (buffer.Array == null || buffer.Count == 0)
+            {
+                return CloseAsync(1000, string.Empty, cancel);
+            }
+            else if (buffer.Count >= 2)
+            {
+                // Unpack the close message.
+                int statusCode =
+                    (buffer.Array[buffer.Offset] << 8)
+                        | buffer.Array[buffer.Offset + 1];
+                string description = Encoding.UTF8.GetString(buffer.Array, buffer.Offset + 2, buffer.Count - 2);
+
+                return CloseAsync(statusCode, description, cancel);
+            }
+            else
+            {
+                throw new ArgumentOutOfRangeException("buffer");
+            }
+        }
+
+        internal async Task CleanupAsync()
+        {
+            switch (_webSocket.State)
+            {
+                case WebSocketState.Closed: // Closed gracefully, no action needed. 
+                case WebSocketState.Aborted: // Closed abortively, no action needed.                       
+                    break;
+                case WebSocketState.CloseReceived:
+                    // Echo what the client said, if anything.
+                    await _webSocket.CloseAsync(_webSocket.CloseStatus ?? WebSocketCloseStatus.NormalClosure,
+                        _webSocket.CloseStatusDescription ?? string.Empty, _cancellationToken);
+                    break;
+                case WebSocketState.Open:
+                case WebSocketState.CloseSent: // No close received, abort so we don't have to drain the pipe.
+                    _webSocket.Abort();
+                    break;
+                default:
+                    throw new ArgumentOutOfRangeException("state", _webSocket.State, string.Empty);
+            }
+        }
+
+        private static WebSocketMessageType OpCodeToEnum(int messageType)
+        {
+            switch (messageType)
+            {
+                case 0x1:
+                    return WebSocketMessageType.Text;
+                case 0x2:
+                    return WebSocketMessageType.Binary;
+                case 0x8:
+                    return WebSocketMessageType.Close;
+                default:
+                    throw new ArgumentOutOfRangeException("messageType", messageType, string.Empty);
+            }
+        }
+
+        private static int EnumToOpCode(WebSocketMessageType webSocketMessageType)
+        {
+            switch (webSocketMessageType)
+            {
+                case WebSocketMessageType.Text:
+                    return 0x1;
+                case WebSocketMessageType.Binary:
+                    return 0x2;
+                case WebSocketMessageType.Close:
+                    return 0x8;
+                default:
+                    throw new ArgumentOutOfRangeException("webSocketMessageType", webSocketMessageType, string.Empty);
+            }
+        }
+    }
+}