OWIN WebSockets: Cleanup, docs, extension methods.

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

@@ -23,6 +23,7 @@
     <Compile Include="IOwinEnvironmentFeature.cs" />
     <Compile Include="OwinConstants.cs" />
     <Compile Include="OwinEnvironment.cs" />
+    <Compile Include="OwinEnvironmentFeature.cs" />
     <Compile Include="OwinExtensions.cs" />
     <Compile Include="OwinFeatureCollection.cs" />
     <Compile Include="Utilities.cs" />

src/Microsoft.AspNet.Owin/OwinEnvironmentFeature.cs

@@ -0,0 +1,12 @@
+// 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.Collections.Generic;
+
+namespace Microsoft.AspNet.Owin
+{
+    public class OwinEnvironmentFeature : IOwinEnvironmentFeature
+    {
+        public IDictionary<string, object> Environment { get; set; }
+    }
+}

src/Microsoft.AspNet.Owin/OwinExtensions.cs

@@ -25,7 +25,7 @@ namespace Microsoft.AspNet.Builder
     {
         public static AddMiddleware UseOwin(this IBuilder builder)
         {
-            return middleware =>
+            AddMiddleware add = middleware =>
             {
                 Func<RequestDelegate, RequestDelegate> middleware1 = next1 =>
                 {
@@ -36,11 +36,26 @@ namespace Microsoft.AspNet.Builder
                     var app = middleware(exitMiddlware);
                     return httpContext =>
                     {
-                        return app.Invoke(new OwinEnvironment(httpContext));
+                        // Use the existing OWIN env if there is one.
+                        IDictionary<string, object> env;
+                        var owinEnvFeature = httpContext.GetFeature<IOwinEnvironmentFeature>();
+                        if (owinEnvFeature != null)
+                        {
+                            env = owinEnvFeature.Environment;
+                            env[typeof(HttpContext).FullName] = httpContext;
+                        }
+                        else
+                        {
+                            env = new OwinEnvironment(httpContext);
+                        }
+                        return app.Invoke(env);
                     };
                 };
                 builder.Use(middleware1);
             };
+            // Adapt WebSockets by default.
+            add(WebSocketAcceptAdapter.AdaptWebSockets);
+            return add;
         }
 
         public static IBuilder UseOwin(this IBuilder builder, Action<AddMiddleware> pipeline)
@@ -51,6 +66,8 @@ namespace Microsoft.AspNet.Builder
 
         public static IBuilder UseBuilder(this AddMiddleware app)
         {
+            // Adapt WebSockets by default.
+            app(OwinWebSocketAcceptAdapter.AdaptWebSockets);
             var builder = new Builder(serviceProvider: null);
 
             CreateMiddleware middleware = CreateMiddlewareFactory(exit =>
@@ -74,10 +91,22 @@ namespace Microsoft.AspNet.Builder
 
                 return env =>
                 {
-                    return app.Invoke(
+                    // Use the existing HttpContext if there is one.
-                        new DefaultHttpContext(
+                    HttpContext context;
-                            new FeatureCollection(
+                    object obj;
-                                    new OwinFeatureCollection(env))));
+                    if (env.TryGetValue(typeof(HttpContext).FullName, out obj))
+                    {
+                        context = (HttpContext)obj;
+                        context.SetFeature<IOwinEnvironmentFeature>(new OwinEnvironmentFeature() { Environment = env });
+                    }
+                    else
+                    {
+                        context = new DefaultHttpContext(
+                                    new FeatureCollection(
+                                        new OwinFeatureCollection(env)));
+                    }
+
+                    return app.Invoke(context);
                 };
             };
         }

src/Microsoft.AspNet.Owin/OwinFeatureCollection.cs

@@ -39,6 +39,7 @@ namespace Microsoft.AspNet.Owin
         public OwinFeatureCollection(IDictionary<string, object> environment)
         {
             Environment = environment;
+            SupportsWebSockets = true;
         }
 
         T Prop<T>(string key)
@@ -239,7 +240,7 @@ namespace Microsoft.AspNet.Owin
         IAuthenticationHandler IHttpAuthenticationFeature.Handler { get; set; }
 
         /// <summary>
-        /// Gets or sets if the underlying server supports WebSockets. This is disabled by default.
+        /// Gets or sets if the underlying server supports WebSockets. This is enabled by default.
         /// The value should be consistant across requests.
         /// </summary>
         public bool SupportsWebSockets { get; set; }

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

@@ -27,6 +27,10 @@ namespace Microsoft.AspNet.Owin
             Task<WebSocket>
         >;
 
+    /// <summary>
+    /// This adapts the OWIN WebSocket accept flow to match the ASP.NET WebSocket Accept flow.
+    /// This enables ASP.NET components to use WebSockets on OWIN based servers.
+    /// </summary>
     public class OwinWebSocketAcceptAdapter
     {
         private WebSocketAccept _owinWebSocketAccept;
@@ -94,6 +98,16 @@ namespace Microsoft.AspNet.Owin
             }
         }
 
+        // Order of operations:
+        // 1. A WebSocket handshake request is received by the middleware.
+        // 2. The middleware inserts an alternate Accept signature into the OWIN environment.
+        // 3. The middleware invokes Next and stores Next's Task locally. It then returns an alternate Task to the server.
+        // 4. The OwinFeatureCollection adapts the alternate Accept signature to IHttpWebSocketFeature.AcceptAsync.
+        // 5. A component later in the pipleline invokes IHttpWebSocketFeature.AcceptAsync (mapped to AcceptWebSocketAsync).
+        // 6. The middleware calls the OWIN Accept, providing a local callback, and returns an incomplete Task.
+        // 7. The middleware completes the alternate Task it returned from Invoke, telling the server that the request pipeline has completed.
+        // 8. The server invokes the middleware's callback, which creats a WebSocket adapter complete's the orriginal Accept Task with it.
+        // 9. The middleware waits while the application uses the WebSocket, where the end is signaled by the Next's Task completion.
         public static AppFunc AdaptWebSockets(AppFunc next)
         {
             return environment =>

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

@@ -10,8 +10,8 @@ namespace Microsoft.AspNet.Owin
     {
         private IDictionary<string, object> _options;
 
-	    public OwinWebSocketAcceptContext() : this(new Dictionary<string, object>(1))
+        public OwinWebSocketAcceptContext() : this(new Dictionary<string, object>(1))
-	    {
+        {
         }
 
         public OwinWebSocketAcceptContext(IDictionary<string, object> options)

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

@@ -28,6 +28,10 @@ namespace Microsoft.AspNet.Owin
             Task<WebSocket>
         >;
 
+    /// <summary>
+    /// This adapts the ASP.NET WebSocket Accept flow to match the OWIN WebSocket accept flow.
+    /// This enables OWIN based components to use WebSockets on ASP.NET servers.
+    /// </summary>
     public class WebSocketAcceptAdapter
     {
         private IDictionary<string, object> _env;
@@ -63,7 +67,7 @@ namespace Microsoft.AspNet.Owin
                     {
                         IWebSocketAcceptContext acceptContext = null;
                         object obj;
-                        if (adapter._options.TryGetValue(typeof(IWebSocketAcceptContext).FullName, out obj))
+                        if (adapter._options != null && adapter._options.TryGetValue(typeof(IWebSocketAcceptContext).FullName, out obj))
                         {
                             acceptContext = obj as IWebSocketAcceptContext;
                         }