OwinWebSocketAcceptAdapter.cs

// Copyright (c) .NET Foundation. 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.Tasks;
using Microsoft.AspNetCore.Http;

namespace Microsoft.AspNetCore.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
        <
            WebSocketAcceptContext, // WebSocket Accept parameters
            Task<WebSocket>
        >;

    /// <summary>
    /// This adapts the OWIN WebSocket accept flow to match the ASP.NET Core WebSocket Accept flow.
    /// This enables ASP.NET Core components to use WebSockets on OWIN based servers.
    /// </summary>
    public class OwinWebSocketAcceptAdapter
    {
        private WebSocketAccept _owinWebSocketAccept;
        private TaskCompletionSource<int> _requestTcs = new TaskCompletionSource<int>();
        private TaskCompletionSource<WebSocket> _acceptTcs = new TaskCompletionSource<WebSocket>();
        private TaskCompletionSource<int> _upstreamWentAsync = new TaskCompletionSource<int>();
        private string _subProtocol = null;

        private OwinWebSocketAcceptAdapter(WebSocketAccept owinWebSocketAccept)
        {
            _owinWebSocketAccept = owinWebSocketAccept;
        }

        private Task RequestTask { get { return _requestTcs.Task; } }
        private Task UpstreamTask { get; set; }
        private TaskCompletionSource<int> UpstreamWentAsyncTcs { get { return _upstreamWentAsync; } }

        private async Task<WebSocket> AcceptWebSocketAsync(WebSocketAcceptContext context)
        {
            IDictionary<string, object> options = null;
            if (context is OwinWebSocketAcceptContext)
            {
                var acceptContext = context as OwinWebSocketAcceptContext;
                options = acceptContext.Options;
                _subProtocol = acceptContext.SubProtocol;
            }
            else if (context?.SubProtocol != null)
            {
                options = new Dictionary<string, object>(1)
                {
                    { OwinConstants.WebSocket.SubProtocol, context.SubProtocol }
                };
                _subProtocol = context.SubProtocol;
            }

            // Accept may have been called synchronously on the original request thread, we might not have a task yet. Go async.
            await _upstreamWentAsync.Task;

            _owinWebSocketAccept(options, OwinAcceptCallback);
            _requestTcs.TrySetResult(0); // Let the pipeline unwind.

            return await _acceptTcs.Task;
        }

        private Task OwinAcceptCallback(IDictionary<string, object> webSocketContext)
        {
            _acceptTcs.TrySetResult(new OwinWebSocketAdapter(webSocketContext, _subProtocol));
            return UpstreamTask;
        }

        // Make sure declined websocket requests complete. This is a no-op for accepted websocket requests.
        private void EnsureCompleted(Task task)
        {
            if (task.IsCanceled)
            {
                _requestTcs.TrySetCanceled();
            }
            else if (task.IsFaulted)
            {
                _requestTcs.TrySetException(task.Exception);
            }
            else
            {
                _requestTcs.TrySetResult(0);
            }
        }

        // 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 pipeline 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 creates a WebSocket adapter and completes the original 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 =>
            {
                object accept;
                if (environment.TryGetValue(OwinConstants.WebSocket.Accept, out accept) && accept is WebSocketAccept)
                {
                    var adapter = new OwinWebSocketAcceptAdapter((WebSocketAccept)accept);

                    environment[OwinConstants.WebSocket.AcceptAlt] = new WebSocketAcceptAlt(adapter.AcceptWebSocketAsync);

                    try
                    {
                        adapter.UpstreamTask = next(environment);
                        adapter.UpstreamWentAsyncTcs.TrySetResult(0);
                        adapter.UpstreamTask.ContinueWith(adapter.EnsureCompleted, TaskContinuationOptions.ExecuteSynchronously);
                    }
                    catch (Exception ex)
                    {
                        adapter.UpstreamWentAsyncTcs.TrySetException(ex);
                        throw;
                    }

                    return adapter.RequestTask;
                }
                else
                {
                    return next(environment);
                }
            };
        }
    }
}