protected Boolean TryCreateReceiveBuffer(out WebSocketReceiveBuffer buffer)
{
lock ( m_lock )
{
if (State >= WebSocketState.Closed)
{
buffer = default;
return(false);
}
if (IsCompressionEnabled())
{
if (m_inflater == null)
{
m_inflater = new ZLibInflater();
m_inflater.AddRef();
Closed.ContinueWith(x => m_inflater.Release());
}
else
{
m_inflater.AddRef();
}
}
}
buffer = new WebSocketReceiveBuffer(m_inflater);
return(true);
}
private async Task ReceivePayloadAsync(WebSocketReceiveBuffer buffer, Int32 payloadLength)
{
var offset = 0;
while (payloadLength > 0)
{
var memory = buffer.GetMemory(Math.Min(payloadLength, DefaultSegmentSize));
var receiveByteCount = await ReceiveBufferAsync(memory.Slice(0, Math.Min(memory.Length, payloadLength)));
if (m_receiveHeader.Mask != 0) // Apply the mask to the payload we received so far
{
ApplyMask(m_receiveHeader.Mask, memory.Slice(0, receiveByteCount).Span, offset);
}
buffer.Advance(receiveByteCount);
payloadLength -= receiveByteCount;
offset += receiveByteCount;
}
}
protected abstract Task ReceiveAsync(WebSocketReceiveBuffer buffer);
private async Task HandleReceivedCloseAsync(WebSocketReceiveBuffer buffer)
{
var closeStatus = WebSocketCloseStatus.NormalClosure;
if (m_closeReceived)
{
Abort("Close message already has been received.");
return;
}
if (m_receiveHeader.Compressed)
{
Abort("The close message must not be compressed.");
return;
}
if (m_receiveHeader.PayloadLength > MaxMessageSize)
{
Abort($"Close message payload size {m_receiveHeader.PayloadLength} is too big.");
return;
}
if (m_receiveHeader.PayloadLength > 0)
{
await ReceivePayloadAsync(buffer, (Int32)m_receiveHeader.PayloadLength);
buffer.Success = true;
using (var result = buffer.ToResult())
{
ReadPayload(result.Message.Buffer);
}
void ReadPayload(ReadOnlySequence <Byte> sequence)
{
var span = sequence.First.Span;
// The first 2 bytes is the close status
closeStatus = (WebSocketCloseStatus)(span[0] << 8 | span[1]);
if (m_receiveHeader.PayloadLength > 2)
{
m_closeStatusDescription = Encoding.GetString(sequence.Slice(2));
}
}
}
var closeTask = Task.CompletedTask;
lock ( m_lock )
{
m_closeReceived = true;
if (m_state == WebSocketState.Closing)
{
Debug.Assert(m_closeSent);
ChangeState(WebSocketState.Closed);
}
else if (m_state == WebSocketState.Open)
{
closeTask = CloseAsync(closeStatus, m_closeStatusDescription);
}
}
await closeTask;
}
private async Task ReceiveHeaderAsync(WebSocketReceiveBuffer buffer)
{
var memory = buffer.GetMemory(MaxHeaderSize).Slice(0, MaxHeaderSize);
var receivedSize = 0;
var remainingSize = 2;
Receive:
while (remainingSize > 0)
{
var byteCount = await ReceiveBufferAsync(memory.Slice(receivedSize, remainingSize));
remainingSize -= byteCount;
receivedSize += byteCount;
}
var headerSize = 2;
var masked = (memory.Span[1] & 0b1000_0000) != 0;
m_receiveHeader.Mask = 0;
m_receiveHeader.Fin = (memory.Span[0] & 0b1000_0000) != 0;
m_receiveHeader.Compressed = (memory.Span[0] & 0b0100_0000) != 0;
m_receiveHeader.Opcode = (MessageOpcode)(memory.Span[0] & 0b0000_1111);
m_receiveHeader.PayloadLength = memory.Span[1] & 0b0111_1111;
buffer.Compressed = m_receiveHeader.Compressed;
buffer.Type = (WebSocketMessageType)m_receiveHeader.Opcode;
if (masked)
{
headerSize += 4;
}
if (m_receiveHeader.PayloadLength == 126)
{
headerSize += 2;
}
else if (m_receiveHeader.PayloadLength == 127)
{
headerSize += 8;
}
if (receivedSize < headerSize)
{
remainingSize = headerSize - receivedSize;
goto Receive; // More data is needed
}
// Read the remainder of the payload length, if necessary
if (m_receiveHeader.PayloadLength == 126)
{
m_receiveHeader.PayloadLength = (memory.Span[2] << 8) | memory.Span[3];
}
else if (m_receiveHeader.PayloadLength == 127)
{
m_receiveHeader.PayloadLength = 0;
for (var i = 0; i < 8; i++)
{
m_receiveHeader.PayloadLength = (m_receiveHeader.PayloadLength << 8) | memory.Span[2 + i];
}
}
if (masked)
{
m_receiveHeader.Mask = Unsafe.As <Byte, Int32>(ref memory.Span[headerSize - 4]);
}
WebSocketsEventSource.Log.Receive(Id,
m_receiveHeader.Opcode,
m_receiveHeader.PayloadLength,
m_receiveHeader.Compressed,
m_receiveHeader.Fin);
}
protected override async Task ReceiveAsync(WebSocketReceiveBuffer buffer)
{
try
{
while (true)
{
await ReceiveHeaderAsync(buffer);
// Receive header uses the buffer but doesn't advance it, so it's safe to use the 0 length
// check as indication that this is the first receive.
if (buffer.Length == 0)
{
if (m_receiveHeader.Opcode == MessageOpcode.Ping || m_receiveHeader.Opcode == MessageOpcode.Pong)
{
if (m_receiveHeader.Opcode == MessageOpcode.Ping)
{
await ProcessPingAsync();
}
continue;
}
else if (m_receiveHeader.Opcode == MessageOpcode.Close)
{
await HandleReceivedCloseAsync(buffer);
// The receive close procedure sets the buffer status to success in order
// to be able to parse the close reason. Unset it here.
buffer.Success = false;
break;
}
else if (m_receiveHeader.Opcode != MessageOpcode.Text && m_receiveHeader.Opcode != MessageOpcode.Binary)
{
await CloseAsync(WebSocketCloseStatus.InvalidPayloadData, $"Unexpected opcode received {m_receiveHeader.Opcode}.");
break;
}
}
else if (m_receiveHeader.Opcode != MessageOpcode.Continuation)
{
await CloseAsync(WebSocketCloseStatus.InvalidPayloadData, $"Unexpected opcode received {m_receiveHeader.Opcode}.");
break;
}
if (m_receiveHeader.PayloadLength > 0)
{
if (buffer.Length + m_receiveHeader.PayloadLength > MaxMessageSize)
{
await CloseAsync(WebSocketCloseStatus.MessageTooBig, $"Message size must not be greater than {MaxMessageSize}.");
break;
}
await ReceivePayloadAsync(buffer, (Int32)m_receiveHeader.PayloadLength);
}
if (m_receiveHeader.Fin)
{
buffer.Success = true;
break;
}
}
}
catch (Exception ex)
{
Abort(ex);
}
}