private async Task <WebSocketContext> AcceptWebSocketAsync(HttpContext context, CancellationToken token)
{
//if (TD.WebSocketConnectionAcceptStartIsEnabled())
//{
// TD.WebSocketConnectionAcceptStart(this.httpRequestContext.EventTraceActivity);
//}
if (!context.WebSockets.IsWebSocketRequest)
{
context.Response.StatusCode = (int)HttpStatusCode.BadRequest;
context.Features.Get <IHttpResponseFeature>().ReasonPhrase = SR.WebSocketEndpointOnlySupportWebSocketError;
return(null);
}
try
{
using (token.Register(() => { context.Abort(); }))
{
string negotiatedProtocol = null;
// match client protocols vs server protocol
foreach (string protocol in context.WebSockets.WebSocketRequestedProtocols)
{
if (string.Compare(protocol, _httpSettings.WebSocketSettings.SubProtocol, StringComparison.OrdinalIgnoreCase) == 0)
{
negotiatedProtocol = protocol;
break;
}
}
if (negotiatedProtocol == null)
{
string errorMessage = SR.Format(SR.WebSocketInvalidProtocolNotInClientList, _httpSettings.WebSocketSettings.SubProtocol, string.Join(", ", context.WebSockets.WebSocketRequestedProtocols));
Fx.Exception.AsWarning(new WebException(errorMessage));
context.Response.StatusCode = (int)HttpStatusCode.UpgradeRequired;
context.Features.Get <IHttpResponseFeature>().ReasonPhrase = SR.WebSocketEndpointOnlySupportWebSocketError;
return(null);
}
WebSocket webSocket = await context.WebSockets.AcceptWebSocketAsync(negotiatedProtocol);
return(new AspNetCoreWebSocketContext(context, webSocket));
}
}
catch (Exception ex)
{
if (Fx.IsFatal(ex))
{
throw;
}
if (token.IsCancellationRequested)
{
throw Fx.Exception.AsError(new TimeoutException(SR.AcceptWebSocketTimedOutError));
}
WebSocketHelper.ThrowCorrectException(ex);
throw;
}
}
public override async Task <int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
{
Fx.Assert(_messageSource != null, "messageSource should not be null in read case.");
Fx.Assert(cancellationToken.CanBeCanceled, "WebSocketStream should be wrapped by TimeoutStream which should pass a cancellable token");
cancellationToken.ThrowIfCancellationRequested();
if (_endOfMessageReached)
{
return(0);
}
if (_initialReadBuffer.Count != 0)
{
return(GetBytesFromInitialReadBuffer(buffer, offset, count));
}
int receivedBytes = 0;
if (_endofMessageReceived)
{
_endOfMessageReached = true;
}
else
{
//if (TD.WebSocketAsyncReadStartIsEnabled())
//{
// TD.WebSocketAsyncReadStart(this.webSocket.GetHashCode());
//}
WebSocketReceiveResult result = null;
try
{
result = await _webSocket.ReceiveAsync(new ArraySegment <byte>(buffer, offset, count), cancellationToken);
}
catch (Exception ex)
{
if (Fx.IsFatal(ex))
{
throw;
}
WebSocketHelper.ThrowCorrectException(ex, TimeoutHelper.GetOriginalTimeout(cancellationToken), WebSocketHelper.ReceiveOperation);
}
if (result.EndOfMessage)
{
_endofMessageReceived = true;
_endOfMessageReached = true;
}
receivedBytes = result.Count;
CheckResultAndEnsureNotCloseMessage(_messageSource, result);
//if (TD.WebSocketAsyncReadStopIsEnabled())
//{
// TD.WebSocketAsyncReadStop(
// this.webSocket.GetHashCode(),
// receivedBytes,
// this.messageSource != null ? TraceUtility.GetRemoteEndpointAddressPort(this.messageSource.RemoteEndpointMessageProperty) : string.Empty);
//}
}
if (_endOfMessageReached)
{
await CleanupAsync(cancellationToken);
}
return(receivedBytes);
}
private async Task ReadBufferedMessageAsync()
{
byte[] internalBuffer = null;
try
{
internalBuffer = _bufferManager.TakeBuffer(_receiveBufferSize);
int receivedByteCount = 0;
bool endOfMessage = false;
WebSocketReceiveResult result = null;
do
{
try
{
//if (TD.WebSocketAsyncReadStartIsEnabled())
//{
// TD.WebSocketAsyncReadStart(this.webSocket.GetHashCode());
//}
Task <WebSocketReceiveResult> receiveTask = _webSocket.ReceiveAsync(
new ArraySegment <byte>(internalBuffer, receivedByteCount, internalBuffer.Length - receivedByteCount),
CancellationToken.None);
await receiveTask.ConfigureAwait(false);
result = receiveTask.Result;
CheckCloseStatus(result);
endOfMessage = result.EndOfMessage;
receivedByteCount += result.Count;
if (receivedByteCount >= internalBuffer.Length && !result.EndOfMessage)
{
if (internalBuffer.Length >= _maxBufferSize)
{
_pendingException = Fx.Exception.AsError(new QuotaExceededException(SR.Format(SR.MaxReceivedMessageSizeExceeded, _maxBufferSize)));
return;
}
int newSize = (int)Math.Min(((double)internalBuffer.Length) * 2, _maxBufferSize);
Fx.Assert(newSize > 0, "buffer size should be larger than zero.");
byte[] newBuffer = _bufferManager.TakeBuffer(newSize);
Buffer.BlockCopy(internalBuffer, 0, newBuffer, 0, receivedByteCount);
_bufferManager.ReturnBuffer(internalBuffer);
internalBuffer = newBuffer;
}
//if (TD.WebSocketAsyncReadStopIsEnabled())
//{
// TD.WebSocketAsyncReadStop(
// this.webSocket.GetHashCode(),
// receivedByteCount,
// TraceUtility.GetRemoteEndpointAddressPort(this.RemoteEndpointMessageProperty));
//}
}
catch (AggregateException ex)
{
WebSocketHelper.ThrowCorrectException(ex, TimeSpan.MaxValue, WebSocketHelper.ReceiveOperation);
}
}while (!endOfMessage && !_closureReceived);
byte[] buffer = null;
bool success = false;
try
{
buffer = _bufferManager.TakeBuffer(receivedByteCount);
Buffer.BlockCopy(internalBuffer, 0, buffer, 0, receivedByteCount);
Fx.Assert(result != null, "Result should not be null");
_pendingMessage = await PrepareMessageAsync(result, buffer, receivedByteCount);
success = true;
}
finally
{
if (buffer != null && (!success || _pendingMessage == null))
{
_bufferManager.ReturnBuffer(buffer);
}
}
}
catch (Exception ex)
{
if (Fx.IsFatal(ex))
{
throw;
}
_pendingException = WebSocketHelper.ConvertAndTraceException(ex, TimeSpan.MaxValue, WebSocketHelper.ReceiveOperation);
}
finally
{
if (internalBuffer != null)
{
_bufferManager.ReturnBuffer(internalBuffer);
}
}
}
private async void StartNextReceiveAsync()
{
Fx.Assert(_receiveTask == null || _receiveTask.Task.IsCompleted, "this.receiveTask is not completed.");
_receiveTask = new TaskCompletionSource <object>();
int currentState = Interlocked.CompareExchange(ref _asyncReceiveState, AsyncReceiveState.Started, AsyncReceiveState.Finished);
Fx.Assert(currentState == AsyncReceiveState.Finished, "currentState is not AsyncReceiveState.Finished: " + currentState);
if (currentState != AsyncReceiveState.Finished)
{
throw Fx.Exception.AsError(new InvalidOperationException());
}
try
{
if (_useStreaming)
{
if (_streamWaitTask != null)
{
//// Wait until the previous stream message finished.
await _streamWaitTask.Task.ConfigureAwait(false);
}
_streamWaitTask = new TaskCompletionSource <object>();
}
if (_pendingException == null)
{
if (!_useStreaming)
{
await ReadBufferedMessageAsync().ConfigureAwait(false);
}
else
{
byte[] buffer = _bufferManager.TakeBuffer(_receiveBufferSize);
bool success = false;
try
{
//if (TD.WebSocketAsyncReadStartIsEnabled())
//{
// TD.WebSocketAsyncReadStart(this.webSocket.GetHashCode());
//}
try
{
Task <WebSocketReceiveResult> receiveTask = _webSocket.ReceiveAsync(
new ArraySegment <byte>(buffer, 0, _receiveBufferSize),
CancellationToken.None);
await receiveTask.ConfigureAwait(false);
WebSocketReceiveResult result = receiveTask.Result;
CheckCloseStatus(result);
_pendingMessage = await PrepareMessageAsync(result, buffer, result.Count);
//if (TD.WebSocketAsyncReadStopIsEnabled())
//{
// TD.WebSocketAsyncReadStop(
// this.webSocket.GetHashCode(),
// result.Count,
// TraceUtility.GetRemoteEndpointAddressPort(this.remoteEndpointMessageProperty));
//}
}
catch (AggregateException ex)
{
WebSocketHelper.ThrowCorrectException(ex, _asyncReceiveTimeout, WebSocketHelper.ReceiveOperation);
}
success = true;
}
catch (Exception ex)
{
if (Fx.IsFatal(ex))
{
throw;
}
_pendingException = WebSocketHelper.ConvertAndTraceException(ex, _asyncReceiveTimeout, WebSocketHelper.ReceiveOperation);
}
finally
{
if (!success)
{
_bufferManager.ReturnBuffer(buffer);
}
}
}
}
}
finally
{
if (Interlocked.CompareExchange(ref _asyncReceiveState, AsyncReceiveState.Finished, AsyncReceiveState.Started) == AsyncReceiveState.Started)
{
_receiveTask.SetResult(null);
}
}
}
protected override async Task OnSendCoreAsync(Message message, CancellationToken token)
{
Fx.Assert(message != null, "message should not be null.");
WebSocketMessageType outgoingMessageType = GetWebSocketMessageType(message);
if (IsStreamedOutput)
{
WebSocketStream webSocketStream = new WebSocketStream(WebSocket, outgoingMessageType, token);
TimeoutStream timeoutStream = new TimeoutStream(webSocketStream, token);
await MessageEncoder.WriteMessageAsync(message, timeoutStream);
await webSocketStream.WriteEndOfMessageAsync(token);
}
else
{
ArraySegment <byte> messageData = EncodeMessage(message);
bool success = false;
try
{
//if (TD.WebSocketAsyncWriteStartIsEnabled())
//{
// TD.WebSocketAsyncWriteStart(
// this.WebSocket.GetHashCode(),
// messageData.Count,
// this.RemoteAddress != null ? this.RemoteAddress.ToString() : string.Empty);
//}
try
{
await WebSocket.SendAsync(messageData, outgoingMessageType, true, token);
}
catch (Exception ex)
{
if (Fx.IsFatal(ex))
{
throw;
}
WebSocketHelper.ThrowCorrectException(ex, TimeoutHelper.GetOriginalTimeout(token), WebSocketHelper.SendOperation);
}
//if (TD.WebSocketAsyncWriteStopIsEnabled())
//{
// TD.WebSocketAsyncWriteStop(this.webSocket.GetHashCode());
//}
success = true;
}
finally
{
try
{
BufferManager.ReturnBuffer(messageData.Array);
}
catch (Exception ex)
{
if (Fx.IsFatal(ex) || success)
{
throw;
}
Fx.Exception.TraceUnhandledException(ex);
}
}
}
}
