Make WebSocketFrame reads fully async

Make WebSocketFrame reads fully async. Process close requests via the message queue, so that anything received prior to the close frame is handled

Author: cswiedler

websocket-sharp/Ext.cs

@@ -618,37 +618,52 @@ internal static byte[] ReadBytes (this Stream stream, long length, int bufferLen
  }
  }
 
- internal static void ReadBytesAsync (
+ internal static void ReadBytesAsync(
  this Stream stream, int length, Action<byte[]> completed, Action<Exception> error)
  {
  var buff = new byte[length];
- stream.BeginRead (
- buff,
- 0,
- length,
- ar => {
- try {
- byte[] bytes = null;
- try {
- var len = stream.EndRead (ar);
- bytes = len < 1
- ? EmptyByteArray
- : len < length
- ? stream.readBytes (buff, len, length - len)
- : buff;
- }
- catch {
- bytes = EmptyByteArray;
- }
-
- if (completed != null)
- completed (bytes);
+ ReadBytesAsync(stream, buff, 0, completed, error);
+ }
+
+ internal static void ReadBytesAsync(
+ this Stream stream, byte[] buff, int currentOffset, Action<byte[]> completed, Action<Exception> error)
+ {
+ AsyncCallback readHandler = (ar) =>
+ {
+ try
+ {
+ var readLength = stream.EndRead(ar);
+ if (readLength == 0)
+ {
+ // EOF/disconnect before reading full length
+ completed(EmptyByteArray);
+ return;
  }
- catch (Exception ex) {
- if (error != null)
- error (ex);
+
+ currentOffset += readLength;
+ if (currentOffset == buff.Length)
+ {
+ // we read everything we needed to read
+ completed(buff);
  }
- },
+ else
+ {
+ // need to read more
+ ReadBytesAsync(stream, buff, currentOffset, completed, error);
+ }
+ }
+ catch (Exception ex)
+ {
+ if (error != null)
+ error(ex);
+ }
+ };
+
+ stream.BeginRead(
+ buff,
+ currentOffset,
+ buff.Length - currentOffset,
+ readHandler,
  null);
  }
 

websocket-sharp/MessageEventArgs.cs

@@ -52,6 +52,7 @@ public class MessageEventArgs : EventArgs
  private string _data;
  private bool _dataSet;
  private Opcode _opcode;
+ private PayloadData _payloadData;
  private byte[] _rawData;
 
  #endregion
@@ -61,6 +62,7 @@ public class MessageEventArgs : EventArgs
  internal MessageEventArgs (WebSocketFrame frame)
  {
  _opcode = frame.Opcode;
+ _payloadData = frame.PayloadData;
  _rawData = frame.PayloadData.ApplicationData;
  }
 
@@ -110,6 +112,14 @@ public byte[] RawData {
  }
  }
 
+ internal PayloadData PayloadData
+ {
+ get
+ {
+ return _payloadData ?? (_payloadData = new PayloadData(_rawData));
+ }
+ }
+
  /// <summary>
  /// Gets the type of the message.
  /// </summary>

websocket-sharp/WebSocket.cs

@@ -1001,8 +1001,7 @@ private void open ()
 
  private bool processCloseFrame (WebSocketFrame frame)
  {
- var payload = frame.PayloadData;
- close (new CloseEventArgs (payload), !payload.IncludesReservedCloseStatusCode, false);
+ enqueueToMessageEventQueue(new MessageEventArgs(frame));
 
  return false;
  }
@@ -1461,43 +1460,80 @@ private void setClientStream ()
  }
  }
 
- private void startReceiving ()
+ private void startReceiving()
  {
  if (_messageEventQueue.Count > 0)
- _messageEventQueue.Clear ();
+ _messageEventQueue.Clear();
 
- _exitReceiving = new AutoResetEvent (false);
- _receivePong = new AutoResetEvent (false);
+ _exitReceiving = new AutoResetEvent(false);
+ _receivePong = new AutoResetEvent(false);
 
  Action receive = null;
- receive = () => WebSocketFrame.ReadAsync (
- _stream,
- false,
- frame => {
- if (processReceivedFrame (frame) && _readyState != WebSocketState.Closed) {
- receive ();
-
- if ((frame.IsControl && !(frame.IsPing && _emitOnPing)) || !frame.IsFinal)
- return;
 
- lock (_forEvent) {
- try {
- var e = dequeueFromMessageEventQueue ();
- if (e != null && _readyState == WebSocketState.Open)
- OnMessage.Emit (this, e);
- }
- catch (Exception ex) {
- processException (ex, "An exception has occurred during an OnMessage event.");
- }
- }
- }
- else if (_exitReceiving != null) {
- _exitReceiving.Set ();
+ Action<WebSocketFrame> messageHandler = (frame) =>
+ {
+ if (processReceivedFrame(frame) && _readyState != WebSocketState.Closed)
+ {
+ receive();
+
+ if ((frame.IsControl && !(frame.IsPing && _emitOnPing)) || !frame.IsFinal)
+ return;
+
+ handleOneMessage();
+ }
+ else
+ {
+ while (_messageEventQueue.Count > 0)
+ {
+ handleOneMessage();
  }
- },
- ex => processException (ex, "An exception has occurred while receiving a message."));
 
- receive ();
+ if (_exitReceiving != null)
+ {
+ _exitReceiving.Set();
+ }
+ }
+ };
+
+ Action<Exception> errorHandler = (ex) =>
+ {
+ processException(ex, "An exception has occurred while receiving a message.");
+ };
+
+ receive = () => WebSocketFrame.Read(
+ _stream,
+ false,
+ messageHandler,
+ errorHandler);
+
+ receive();
+ }
+
+ private void handleOneMessage()
+ {
+ lock (_forEvent)
+ {
+ try
+ {
+ MessageEventArgs e = dequeueFromMessageEventQueue();
+ if (e != null && _readyState == WebSocketState.Open)
+ {
+ if (e.Type == Opcode.Close)
+ {
+ var payload = e.PayloadData;
+ close(new CloseEventArgs(payload), !payload.IncludesReservedCloseStatusCode, false);
+ }
+ else
+ {
+ OnMessage.Emit(this, e);
+ }
+ }
+ }
+ catch (Exception ex)
+ {
+ processException(ex, "An exception has occurred during an OnMessage event.");
+ }
+ }
  }
 
  // As client