Skip to content

Commit

Permalink
[release/8.0][browser] BrowserWebSocket.ReceiveAsync after server ini…
Browse files Browse the repository at this point in the history
…tiated close (#97002)
  • Loading branch information
pavelsavara authored Jan 23, 2024
1 parent e24179e commit 5e9e29f
Show file tree
Hide file tree
Showing 6 changed files with 205 additions and 61 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,15 @@ internal static partial class WebSocketValidate
internal const int MaxDeflateWindowBits = 15;

internal const int MaxControlFramePayloadLength = 123;
#if TARGET_BROWSER
private const int ValidCloseStatusCodesFrom = 3000;
private const int ValidCloseStatusCodesTo = 4999;
#else
private const int CloseStatusCodeAbort = 1006;
private const int CloseStatusCodeFailedTLSHandshake = 1015;
private const int InvalidCloseStatusCodesFrom = 0;
private const int InvalidCloseStatusCodesTo = 999;
#endif

// [0x21, 0x7E] except separators "()<>@,;:\\\"/[]?={} ".
private static readonly SearchValues<char> s_validSubprotocolChars =
Expand Down Expand Up @@ -84,11 +89,15 @@ internal static void ValidateCloseStatus(WebSocketCloseStatus closeStatus, strin
}

int closeStatusCode = (int)closeStatus;

#if TARGET_BROWSER
// as defined in https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close#code
if (closeStatus != WebSocketCloseStatus.NormalClosure && (closeStatusCode < ValidCloseStatusCodesFrom || closeStatusCode > ValidCloseStatusCodesTo))
#else
if ((closeStatusCode >= InvalidCloseStatusCodesFrom &&
closeStatusCode <= InvalidCloseStatusCodesTo) ||
closeStatusCode == CloseStatusCodeAbort ||
closeStatusCode == CloseStatusCodeFailedTLSHandshake)
#endif
{
// CloseStatus 1006 means Aborted - this will never appear on the wire and is reflected by calling WebSocket.Abort
throw new ArgumentException(SR.Format(SR.net_WebSockets_InvalidCloseStatusCode,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,11 +24,11 @@ public static async Task InvokeAsync(HttpContext context)

if (context.Request.QueryString.HasValue && context.Request.QueryString.Value.Contains("delay10sec"))
{
Thread.Sleep(10000);
await Task.Delay(10000);
}
else if (context.Request.QueryString.HasValue && context.Request.QueryString.Value.Contains("delay20sec"))
{
Thread.Sleep(20000);
await Task.Delay(20000);
}

try
Expand Down Expand Up @@ -124,14 +124,15 @@ await socket.CloseAsync(
}

bool sendMessage = false;
string receivedMessage = null;
if (receiveResult.MessageType == WebSocketMessageType.Text)
{
string receivedMessage = Encoding.UTF8.GetString(receiveBuffer, 0, offset);
receivedMessage = Encoding.UTF8.GetString(receiveBuffer, 0, offset);
if (receivedMessage == ".close")
{
await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, receivedMessage, CancellationToken.None);
}
if (receivedMessage == ".shutdown")
else if (receivedMessage == ".shutdown")
{
await socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, receivedMessage, CancellationToken.None);
}
Expand Down Expand Up @@ -161,6 +162,14 @@ await socket.SendAsync(
!replyWithPartialMessages,
CancellationToken.None);
}
if (receivedMessage == ".closeafter")
{
await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, receivedMessage, CancellationToken.None);
}
else if (receivedMessage == ".shutdownafter")
{
await socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, receivedMessage, CancellationToken.None);
}
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ internal sealed class BrowserWebSocket : WebSocket
private WebSocketState _state;
private bool _disposed;
private bool _aborted;
private bool _closeReceived;
private bool _closeSent;
private int[] responseStatus = new int[3];
private MemoryHandle? responseStatusHandle;

Expand All @@ -37,7 +39,7 @@ public override WebSocketState State
lock (_thisLock)
{
#endif
if (_innerWebSocket == null || _disposed || (_state != WebSocketState.Connecting && _state != WebSocketState.Open && _state != WebSocketState.CloseSent))
if (_innerWebSocket == null || _disposed || _state == WebSocketState.Aborted || _state == WebSocketState.Closed)
{
return _state;
}
Expand All @@ -46,15 +48,9 @@ public override WebSocketState State
#endif

#if FEATURE_WASM_THREADS
return FastState = _innerWebSocket!.SynchronizationContext.Send(static (BrowserWebSocket self) =>
{
lock (self._thisLock)
{
return GetReadyState(self._innerWebSocket!);
} //lock
}, this);
return _innerWebSocket!.SynchronizationContext.Send(GetReadyState, this);
#else
return FastState = GetReadyState(_innerWebSocket!);
return GetReadyState(this);
#endif
}
}
Expand Down Expand Up @@ -148,7 +144,7 @@ public override Task SendAsync(ArraySegment<byte> buffer, WebSocketMessageType m
ThrowIfDisposed();

// fast check of previous _state instead of GetReadyState(), the readyState would be validated on JS side
if (FastState != WebSocketState.Open)
if (FastState != WebSocketState.Open && FastState != WebSocketState.CloseReceived)
{
throw new InvalidOperationException(SR.net_WebSockets_NotConnected);
}
Expand Down Expand Up @@ -240,7 +236,7 @@ public override Task CloseOutputAsync(WebSocketCloseStatus closeStatus, string?
{
throw new WebSocketException(WebSocketError.InvalidState, SR.Format(SR.net_WebSockets_InvalidState, state, "Connecting, Open, CloseSent, Aborted"));
}
if(state != WebSocketState.Open && state != WebSocketState.Connecting && state != WebSocketState.Aborted)
if (state == WebSocketState.CloseSent)
{
return Task.CompletedTask;
}
Expand Down Expand Up @@ -280,10 +276,6 @@ public override Task CloseAsync(WebSocketCloseStatus closeStatus, string? status
{
throw new WebSocketException(WebSocketError.InvalidState, SR.Format(SR.net_WebSockets_InvalidState, state, "Connecting, Open, CloseSent, Aborted"));
}
if (state != WebSocketState.Open && state != WebSocketState.Connecting && state != WebSocketState.Aborted && state != WebSocketState.CloseSent)
{
return Task.CompletedTask;
}

#if FEATURE_WASM_THREADS
promise = CloseAsyncCore(closeStatus, statusDescription, state != WebSocketState.Aborted, cancellationToken);
Expand Down Expand Up @@ -387,12 +379,13 @@ private void CreateCore(Uri uri, List<string>? requestedSubProtocols)
string[]? subProtocols = requestedSubProtocols?.ToArray();
var onClose = (int code, string reason) =>
{
_closeStatus = (WebSocketCloseStatus)code;
_closeStatusDescription = reason;
#if FEATURE_WASM_THREADS
lock (_thisLock)
{
#endif
_closeStatus = (WebSocketCloseStatus)code;
_closeStatusDescription = reason;
_closeReceived = true;
WebSocketState state = State;
if (state == WebSocketState.Connecting || state == WebSocketState.Open || state == WebSocketState.CloseSent)
{
Expand Down Expand Up @@ -545,15 +538,21 @@ private static WebSocketReceiveResult ConvertResponse(BrowserWebSocket self)
WebSocketMessageType messageType = (WebSocketMessageType)self.responseStatus[typeIndex];
if (messageType == WebSocketMessageType.Close)
{
self._closeReceived = true;
self.FastState = self._closeSent ? WebSocketState.Closed : WebSocketState.CloseReceived;
return new WebSocketReceiveResult(self.responseStatus[countIndex], messageType, self.responseStatus[endIndex] != 0, self.CloseStatus, self.CloseStatusDescription);
}
return new WebSocketReceiveResult(self.responseStatus[countIndex], messageType, self.responseStatus[endIndex] != 0);
}

private async Task CloseAsyncCore(WebSocketCloseStatus closeStatus, string? statusDescription, bool waitForCloseReceived, CancellationToken cancellationToken)
{
_closeStatus = closeStatus;
_closeStatusDescription = statusDescription;
if (!_closeReceived)
{
_closeStatus = closeStatus;
_closeStatusDescription = statusDescription;
}
_closeSent = true;

var closeTask = BrowserInterop.WebSocketClose(_innerWebSocket!, (int)closeStatus, statusDescription, waitForCloseReceived) ?? Task.CompletedTask;
await CancelationHelper(closeTask, cancellationToken, FastState).ConfigureAwait(true);
Expand All @@ -562,6 +561,10 @@ private async Task CloseAsyncCore(WebSocketCloseStatus closeStatus, string? stat
lock (_thisLock)
{
#endif
if (waitForCloseReceived)
{
_closeReceived = true;
}
var state = State;
if (state == WebSocketState.Open || state == WebSocketState.Connecting || state == WebSocketState.CloseSent)
{
Expand Down Expand Up @@ -614,18 +617,42 @@ private async Task CancelationHelper(Task jsTask, CancellationToken cancellation
}
}

private static WebSocketState GetReadyState(JSObject innerWebSocket)
private static WebSocketState GetReadyState(BrowserWebSocket self)
{
var readyState = BrowserInterop.GetReadyState(innerWebSocket);
// https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/readyState
return readyState switch
{
0 => WebSocketState.Connecting, // 0 (CONNECTING)
1 => WebSocketState.Open, // 1 (OPEN)
2 => WebSocketState.CloseSent, // 2 (CLOSING)
3 => WebSocketState.Closed, // 3 (CLOSED)
_ => WebSocketState.None
};
#if FEATURE_WASM_THREADS
lock (self._thisLock)
{
#endif
var readyState = BrowserInterop.GetReadyState(self._innerWebSocket);
// https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/readyState
var st = readyState switch
{
0 => WebSocketState.Connecting, // 0 (CONNECTING)
1 => WebSocketState.Open, // 1 (OPEN)
2 => WebSocketState.CloseSent, // 2 (CLOSING)
3 => WebSocketState.Closed, // 3 (CLOSED)
_ => WebSocketState.None
};
if (st == WebSocketState.Closed || st == WebSocketState.CloseSent)
{
if (self._closeReceived && self._closeSent)
{
st = WebSocketState.Closed;
}
else if (self._closeReceived && !self._closeSent)
{
st = WebSocketState.CloseReceived;
}
else if (!self._closeReceived && self._closeSent)
{
st = WebSocketState.CloseSent;
}
}
self.FastState = st;
return st;
#if FEATURE_WASM_THREADS
} //lock
#endif
}

#endregion
Expand Down
Loading

0 comments on commit 5e9e29f

Please sign in to comment.