-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathWebSocket.cs
More file actions
338 lines (274 loc) · 12.3 KB
/
Copy pathWebSocket.cs
File metadata and controls
338 lines (274 loc) · 12.3 KB
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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
using System.Buffers;
using System.Buffers.Binary;
using System.IO.Pipelines;
using System.Text;
using Shiny.Net.HttpServer.Transports;
namespace Shiny.Net.HttpServer.WebSockets;
/// <summary>A received message.</summary>
public sealed class WebSocketMessage
{
internal WebSocketMessage(WebSocketOpcode type, byte[] payload)
{
this.Type = type;
this.Payload = payload;
}
/// <summary><see cref="WebSocketOpcode.Text"/> or <see cref="WebSocketOpcode.Binary"/>.</summary>
public WebSocketOpcode Type { get; }
public byte[] Payload { get; }
public bool IsText => this.Type == WebSocketOpcode.Text;
/// <summary>The payload decoded as UTF-8. Only meaningful for a text message.</summary>
public string Text => Encoding.UTF8.GetString(this.Payload);
}
/// <summary>Why a socket ended.</summary>
public sealed class WebSocketCloseResult(WebSocketCloseStatus status, string? description)
{
public WebSocketCloseStatus Status { get; } = status;
public string? Description { get; } = description;
}
/// <summary>
/// A live WebSocket.
/// <para>
/// Reading and writing are independent — a socket is full duplex and treating it as
/// request/response is the usual way to deadlock one. Control frames are handled here: a ping is
/// answered with a pong automatically, and a close is acknowledged, so a handler only ever sees the
/// messages it actually cares about.
/// </para>
/// </summary>
public sealed class WebSocket : IAsyncDisposable
{
readonly IConnection connection;
readonly PipeReader input;
readonly PipeWriter output;
readonly SemaphoreSlim writeGate = new(1, 1);
readonly long maxMessageLength;
bool closeSent;
bool closeReceived;
int disposed;
internal WebSocket(IConnection connection, long maxMessageLength)
{
this.connection = connection;
this.input = connection.Input;
this.output = connection.Output;
this.maxMessageLength = maxMessageLength;
}
/// <summary>The sub-protocol agreed during the handshake, if any.</summary>
public string? SubProtocol { get; internal init; }
/// <summary>True until a close frame has been both sent and received.</summary>
public bool IsOpen => !this.closeSent || !this.closeReceived;
/// <summary>How the peer closed, once it has.</summary>
public WebSocketCloseResult? CloseResult { get; private set; }
/// <summary>
/// Reads the next message, returning null when the socket closes. Pings and pongs are handled
/// internally and never surface here.
/// </summary>
public async ValueTask<WebSocketMessage?> ReceiveAsync(CancellationToken cancellationToken = default)
{
var payload = new ArrayBufferWriter<byte>();
var messageType = WebSocketOpcode.Continuation;
var fragmented = false;
while (true)
{
var header = await this.ReadHeaderAsync(cancellationToken).ConfigureAwait(false);
if (header is not { } frame)
return null;
// RFC 6455 §5.1: every frame from a client must be masked. An unmasked one is either a
// broken client or a proxy rewriting traffic, and neither should be trusted.
if (!frame.Masked)
throw new WebSocketProtocolException(WebSocketCloseStatus.ProtocolError, "Client frames must be masked.");
if (frame.IsControl)
{
// Control frames interleave with a fragmented message, so this cannot wait until
// the message is complete.
if (await this.HandleControlAsync(frame, cancellationToken).ConfigureAwait(false))
return null;
continue;
}
switch (frame.Opcode)
{
case WebSocketOpcode.Text or WebSocketOpcode.Binary when fragmented:
throw new WebSocketProtocolException(
WebSocketCloseStatus.ProtocolError,
"A new message started before the previous one finished."
);
case WebSocketOpcode.Text or WebSocketOpcode.Binary:
messageType = frame.Opcode;
break;
case WebSocketOpcode.Continuation when !fragmented:
throw new WebSocketProtocolException(
WebSocketCloseStatus.ProtocolError,
"A continuation frame arrived with no message to continue."
);
}
if (payload.WrittenCount + frame.PayloadLength > this.maxMessageLength)
throw new WebSocketProtocolException(WebSocketCloseStatus.MessageTooBig, "The message is too large.");
await this.ReadPayloadAsync(frame, payload, cancellationToken).ConfigureAwait(false);
if (frame.Fin)
return new WebSocketMessage(messageType, payload.WrittenSpan.ToArray());
fragmented = true;
}
}
/// <summary>Sends a text message.</summary>
public ValueTask SendAsync(string text, CancellationToken cancellationToken = default)
{
ArgumentNullException.ThrowIfNull(text);
return this.SendAsync(Encoding.UTF8.GetBytes(text), WebSocketOpcode.Text, cancellationToken);
}
/// <summary>Sends a binary message.</summary>
public ValueTask SendAsync(ReadOnlyMemory<byte> payload, CancellationToken cancellationToken = default)
=> this.SendAsync(payload, WebSocketOpcode.Binary, cancellationToken);
async ValueTask SendAsync(ReadOnlyMemory<byte> payload, WebSocketOpcode opcode, CancellationToken cancellationToken)
{
ObjectDisposedException.ThrowIf(this.disposed != 0, this);
if (this.closeSent)
throw new InvalidOperationException("The socket has already been closed.");
await this.WriteFrameAsync(opcode, payload, cancellationToken).ConfigureAwait(false);
}
/// <summary>Sends a ping. The peer is expected to pong; nothing here waits for it.</summary>
public ValueTask PingAsync(ReadOnlyMemory<byte> payload = default, CancellationToken cancellationToken = default)
=> this.WriteFrameAsync(WebSocketOpcode.Ping, payload, cancellationToken);
/// <summary>
/// Sends a close frame. The socket is not fully closed until the peer's close comes back, which
/// <see cref="ReceiveAsync"/> reports by returning null.
/// </summary>
public async ValueTask CloseAsync(
WebSocketCloseStatus status = WebSocketCloseStatus.NormalClosure,
string? description = null,
CancellationToken cancellationToken = default
)
{
if (this.closeSent || this.disposed != 0)
return;
this.closeSent = true;
var reason = description is { Length: > 0 } ? Encoding.UTF8.GetBytes(description) : [];
var payload = new byte[2 + reason.Length];
BinaryPrimitives.WriteUInt16BigEndian(payload, (ushort)status);
reason.CopyTo(payload, 2);
try
{
await this.WriteFrameAsync(WebSocketOpcode.Close, payload, cancellationToken).ConfigureAwait(false);
}
catch (Exception ex) when (IsDisconnect(ex))
{
// The peer went away before the courtesy close reached it. Nothing to salvage.
}
}
async ValueTask<WebSocketFrameHeader?> ReadHeaderAsync(CancellationToken cancellationToken)
{
while (true)
{
var result = await this.input.ReadAsync(cancellationToken).ConfigureAwait(false);
var buffer = result.Buffer;
try
{
var remaining = buffer;
if (WebSocketFrameCodec.TryReadHeader(ref remaining, out var header))
{
this.input.AdvanceTo(remaining.Start);
return header;
}
}
catch
{
this.input.AdvanceTo(buffer.Start, buffer.End);
throw;
}
this.input.AdvanceTo(buffer.Start, buffer.End);
if (result.IsCompleted || result.IsCanceled)
return null;
}
}
async ValueTask ReadPayloadAsync(
WebSocketFrameHeader header,
IBufferWriter<byte> destination,
CancellationToken cancellationToken
)
{
var remaining = header.PayloadLength;
var maskOffset = 0;
while (remaining > 0)
{
var result = await this.input.ReadAsync(cancellationToken).ConfigureAwait(false);
var buffer = result.Buffer;
if (buffer.IsEmpty && (result.IsCompleted || result.IsCanceled))
throw new WebSocketProtocolException(WebSocketCloseStatus.ProtocolError, "The frame payload was truncated.");
var take = (int)Math.Min(buffer.Length, remaining);
var chunk = buffer.Slice(0, take);
var span = destination.GetSpan(take)[..take];
chunk.CopyTo(span);
// The mask cycles every four bytes across the whole payload, so a chunk that starts
// mid-cycle has to say where it starts.
WebSocketFrameCodec.Unmask(span, header.MaskingKey, maskOffset);
destination.Advance(take);
this.input.AdvanceTo(chunk.End);
remaining -= take;
maskOffset += take;
}
}
/// <summary>Handles a control frame. Returns true when the socket should stop reading.</summary>
async ValueTask<bool> HandleControlAsync(WebSocketFrameHeader header, CancellationToken cancellationToken)
{
var payload = new ArrayBufferWriter<byte>((int)Math.Max(header.PayloadLength, 1));
await this.ReadPayloadAsync(header, payload, cancellationToken).ConfigureAwait(false);
switch (header.Opcode)
{
case WebSocketOpcode.Ping:
// Answered with the same payload, which is what the spec requires and what
// keepalive implementations check for.
await this.WriteFrameAsync(WebSocketOpcode.Pong, payload.WrittenMemory, cancellationToken)
.ConfigureAwait(false);
return false;
case WebSocketOpcode.Pong:
return false;
case WebSocketOpcode.Close:
this.closeReceived = true;
this.CloseResult = ParseClose(payload.WrittenSpan);
// Echo the close back so the peer can shut down cleanly, then stop.
await this.CloseAsync(
this.CloseResult.Status == WebSocketCloseStatus.NoStatusReceived
? WebSocketCloseStatus.NormalClosure
: this.CloseResult.Status,
cancellationToken: cancellationToken
).ConfigureAwait(false);
return true;
default:
return false;
}
}
static WebSocketCloseResult ParseClose(ReadOnlySpan<byte> payload)
{
if (payload.Length < 2)
return new WebSocketCloseResult(WebSocketCloseStatus.NoStatusReceived, null);
var status = (WebSocketCloseStatus)BinaryPrimitives.ReadUInt16BigEndian(payload);
var description = payload.Length > 2 ? Encoding.UTF8.GetString(payload[2..]) : null;
return new WebSocketCloseResult(status, description);
}
async ValueTask WriteFrameAsync(WebSocketOpcode opcode, ReadOnlyMemory<byte> payload, CancellationToken cancellationToken)
{
// Serialized: two concurrent sends would interleave their frames and corrupt both.
await this.writeGate.WaitAsync(cancellationToken).ConfigureAwait(false);
try
{
WebSocketFrameCodec.WriteHeader(this.output, opcode, fin: true, payload.Length);
if (!payload.IsEmpty)
this.output.Write(payload.Span);
await this.output.FlushAsync(cancellationToken).ConfigureAwait(false);
}
finally
{
this.writeGate.Release();
}
}
internal static bool IsDisconnect(Exception ex) => ex
is OperationCanceledException
or System.Net.Sockets.SocketException
or ObjectDisposedException
or InvalidOperationException
or IOException;
public async ValueTask DisposeAsync()
{
if (Interlocked.Exchange(ref this.disposed, 1) != 0)
return;
this.writeGate.Dispose();
await this.connection.DisposeAsync().ConfigureAwait(false);
}
}