-
Notifications
You must be signed in to change notification settings - Fork 9.9k
/
OwinWebSocketAcceptAdapter.cs
143 lines (128 loc) · 5.92 KB
/
OwinWebSocketAcceptAdapter.cs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
// 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);
}
};
}
}
}