in src/Microsoft.Azure.Relay/WebSockets/NetStandard20/ManagedWebSocket.cs [618:743]
private async Task<WebSocketReceiveResult> ReceiveAsyncPrivate(ArraySegment<byte> payloadBuffer, CancellationToken cancellationToken)
{
// This is a long method. While splitting it up into pieces would arguably help with readability, doing so would
// also result in more allocations, as each async method that yields ends up with multiple allocations. The impact
// of those allocations is amortized across all of the awaits in the method, and since we generally expect a receive
// operation to require at most a single yield (while waiting for data to arrive), it's more efficient to have
// everything in the one method. We do separate out pieces for handling close and ping/pong messages, as we expect
// those to be much less frequent (e.g. we should only get one close per websocket), and thus we can afford to pay
// a bit more for readability and maintainability.
CancellationTokenRegistration registration = cancellationToken.Register(s => ((ManagedWebSocket)s).Abort(), this);
try
{
while (true) // in case we get control frames that should be ignored from the user's perspective
{
// Get the last received header. If its payload length is non-zero, that means we previously
// received the header but were only able to read a part of the fragment, so we should skip
// reading another header and just proceed to use that same header and read more data associated
// with it. If instead its payload length is zero, then we've completed the processing of
// thta message, and we should read the next header.
MessageHeader header = _lastReceiveHeader;
if (header.PayloadLength == 0)
{
if (_receiveBufferCount < (_isServer ? (MaxMessageHeaderLength - MaskLength) : MaxMessageHeaderLength))
{
// Make sure we have the first two bytes, which includes the start of the payload length.
if (_receiveBufferCount < 2)
{
await EnsureBufferContainsAsync(2, cancellationToken, throwOnPrematureClosure: true).ConfigureAwait(false);
}
// Then make sure we have the full header based on the payload length.
// If this is the server, we also need room for the received mask.
long payloadLength = _receiveBuffer[_receiveBufferOffset + 1] & 0x7F;
if (_isServer || payloadLength > 125)
{
int minNeeded =
2 +
(_isServer ? MaskLength : 0) +
(payloadLength <= 125 ? 0 : payloadLength == 126 ? sizeof(ushort) : sizeof(ulong)); // additional 2 or 8 bytes for 16-bit or 64-bit length
await EnsureBufferContainsAsync(minNeeded, cancellationToken).ConfigureAwait(false);
}
}
if (!TryParseMessageHeaderFromReceiveBuffer(out header))
{
await CloseWithReceiveErrorAndThrowAsync(WebSocketCloseStatus.ProtocolError, WebSocketError.Faulted, cancellationToken).ConfigureAwait(false);
}
_receivedMaskOffsetOffset = 0;
}
// If the header represents a ping or a pong, it's a control message meant
// to be transparent to the user, so handle it and then loop around to read again.
// Alternatively, if it's a close message, handle it and exit.
if (header.Opcode == MessageOpcode.Ping || header.Opcode == MessageOpcode.Pong)
{
await HandleReceivedPingPongAsync(header, cancellationToken).ConfigureAwait(false);
continue;
}
else if (header.Opcode == MessageOpcode.Close)
{
return await HandleReceivedCloseAsync(header, cancellationToken).ConfigureAwait(false);
}
// If this is a continuation, replace the opcode with the one of the message it's continuing
if (header.Opcode == MessageOpcode.Continuation)
{
header.Opcode = _lastReceiveHeader.Opcode;
}
// The message should now be a binary or text message. Handle it by reading the payload and returning the contents.
Debug.Assert(header.Opcode == MessageOpcode.Binary || header.Opcode == MessageOpcode.Text, $"Unexpected opcode {header.Opcode}");
// If there's no data to read, return an appropriate result.
int bytesToRead = (int)Math.Min(payloadBuffer.Count, header.PayloadLength);
if (bytesToRead == 0)
{
_lastReceiveHeader = header;
return new WebSocketReceiveResult(
0,
header.Opcode == MessageOpcode.Text ? WebSocketMessageType.Text : WebSocketMessageType.Binary,
header.PayloadLength == 0 ? header.Fin : false);
}
// Otherwise, read as much of the payload as we can efficiently, and upate the header to reflect how much data
// remains for future reads.
if (_receiveBufferCount == 0)
{
await EnsureBufferContainsAsync(1, cancellationToken, throwOnPrematureClosure: false).ConfigureAwait(false);
}
int bytesToCopy = Math.Min(bytesToRead, _receiveBufferCount);
if (_isServer)
{
_receivedMaskOffsetOffset = ApplyMask(_receiveBuffer, _receiveBufferOffset, header.Mask, _receivedMaskOffsetOffset, bytesToCopy);
}
Buffer.BlockCopy(_receiveBuffer, _receiveBufferOffset, payloadBuffer.Array, payloadBuffer.Offset, bytesToCopy);
ConsumeFromBuffer(bytesToCopy);
header.PayloadLength -= bytesToCopy;
// If this a text message, validate that it contains valid UTF8.
if (header.Opcode == MessageOpcode.Text &&
!TryValidateUtf8(new ArraySegment<byte>(payloadBuffer.Array, payloadBuffer.Offset, bytesToCopy), header.Fin, _utf8TextState))
{
await CloseWithReceiveErrorAndThrowAsync(WebSocketCloseStatus.InvalidPayloadData, WebSocketError.Faulted, cancellationToken).ConfigureAwait(false);
}
_lastReceiveHeader = header;
return new WebSocketReceiveResult(
bytesToCopy,
header.Opcode == MessageOpcode.Text ? WebSocketMessageType.Text : WebSocketMessageType.Binary,
bytesToCopy == 0 || (header.Fin && header.PayloadLength == 0));
}
}
catch (Exception exc)
{
throw _state == WebSocketState.Aborted ?
new WebSocketException(WebSocketError.InvalidState, SR.Format(SR.net_WebSockets_InvalidState_ClosedOrAborted, "System.Net.WebSockets.InternalClientWebSocket", "Aborted"), exc) :
new WebSocketException(WebSocketError.ConnectionClosedPrematurely, exc);
}
finally
{
registration.Dispose();
}
}