Improvements in Close() (#250)

(cherry picked from commit b766433)

Author: josesimoes

README.md

@@ -115,8 +115,9 @@ A message can be send by calling `SendString` for a text message or `SendBytes`
 
 #### Closing a connection
 
-The connection can be closed by calling `Close`. Calling this method will send a closing message over the line. You can optional specify a `WebSocketCloseStatus` and description on the reason for closing for debugging purposes. 
-Whenever a connection is closed the event `Closed` is fired.  
+The connection can be closed by calling `Close`. Calling this method will send a closing message over the line. You can optional specify a `WebSocketCloseStatus` and description on the reason for closing for debugging purposes (the description is only sent together with a status, it's ignored with the default `WebSocketCloseStatus.Empty`). 
+`Close` doesn't block: the closing message is sent after the messages already queued and the connection is closed when the other end answers it, or after `ServerTimeout`. Closing with `WebSocketCloseStatus.EndpointUnavailable` is synchronous: the call returns after the closing message is sent and the connection is closed, without waiting for an answer. It waits at most `ServerTimeout` for the closing message to be sent; if that expires first the connection is closed anyway and the closing message may not have been sent. Messages still queued may not be sent.
+Whenever a connection is closed the event `ConnectionClosed` is fired.  
 
 ### Server
 

WebSockets/ReceiveAndControllThread.cs

@@ -100,29 +100,27 @@ private void ProcessIncomingMessage()
                             break;
 
                         case OpCode.ConnectionCloseFrame:
-                            _webSocket.CloseStatus = WebSocketCloseStatus.Empty;
-
-                            if (buffer.Length > 1)
+                            if (!TryGetPeerCloseStatus(buffer, out WebSocketCloseStatus peerCloseStatus))
                             {
-                                byte[] closeByteCode = new byte[] { buffer[1], buffer[0] };
-                                UInt16 statusCode = BitConverter.ToUInt16(closeByteCode, 0);
-                                if (statusCode > 999 && statusCode < 1012) 
-                                {
-                                    _webSocket.CloseStatus = (WebSocketCloseStatus)statusCode;
-                                }
-                            }
+                                _webSocket.HasError = true;
 
-                            //connection asked to be closed return answer
-                            if (_webSocket.TryMarkCloseReceived())
-                            {
-                                _webSocket.RawClose(WebSocketCloseStatus.NormalClosure, buffer, true);
+                                Debug.WriteLine($"{_webSocket.RemoteEndPoint} sent an invalid close frame");
+
+                                // no-op if our close message was already queued
+                                _webSocket.RawClose(WebSocketCloseStatus.ProtocolError, Encoding.UTF8.GetBytes("Invalid close frame"), true);
                             }
-                            else
+                            //connection asked to be closed return answer
+                            else if (_webSocket.TryMarkCloseReceived(peerCloseStatus))
                             {
-                                // our close message can still be queued behind pending messages (simultaneous close)
-                                _webSocket.WaitForCloseMessageSent();
+                                // echo the status code received (RFC 6455 section 5.5.1), without the peer's reason
+                                // (no status code is sent if none was received)
+                                _webSocket.RawClose(peerCloseStatus, null, true);
                             }
 
+                            // our close message can still be queued behind pending messages (simultaneous close)
+                            // returns immediately if it was sent or the connection is already closed
+                            _webSocket.WaitForCloseMessageSent();
+
                             // either this is the response to our close, or RawClose lost a race with another close,
                             // so we can shut down the socket (no-op if already closed)
                             _webSocket.HardClose();
@@ -149,7 +147,6 @@ private void ProcessIncomingMessage()
             }
         }
 
-
         // Runs on the timer thread, concurrently with the receive thread.
         // Shared state is accessed through WebSocket helpers that only hold a lock for field access,
         // and nothing here blocks, so a receive thread holding a lock can't stall this check.
@@ -197,6 +194,37 @@ private void CheckTimeouts(object state)
             }
         }
 
+        // Gets the status code from the payload of a close frame received from the remote endpoint.
+        // Returns false if the payload is invalid: 1 byte long, or with a status code that can't be sent in a close frame (RFC 6455 section 7.4).
+        private static bool TryGetPeerCloseStatus(byte[] buffer, out WebSocketCloseStatus closeStatus)
+        {
+            closeStatus = WebSocketCloseStatus.Empty;
+
+            if (buffer.Length == 0)
+            {
+                // no status code
+                return true;
+            }
+
+            if (buffer.Length == 1)
+            {
+                return false;
+            }
+
+            int statusCode = (buffer[0] << 8) | buffer[1];
+
+            if ((statusCode >= 1000 && statusCode <= 1003)
+                || (statusCode >= 1007 && statusCode <= 1014)
+                || (statusCode >= 3000 && statusCode <= 4999))
+            {
+                closeStatus = (WebSocketCloseStatus)statusCode;
+
+                return true;
+            }
+
+            return false;
+        }
+
         private void OnNewMessage(ReceiveMessageFrame message)
         {
             _webSocket.CallbacksMessageReceivedEventHandler?.Invoke(_webSocket, new MessageReceivedEventArgs() { Frame = message });

WebSockets/WebSocket.cs

@@ -218,26 +218,38 @@ public bool SendBytes(byte[] data, int fragmentSize = -1)
         }
 
         /// <summary>
-        /// Will start closing the WebSocket connection using the close handshake defined in the WebSocket protocol specification section 7.
+        /// Will start closing the WebSocket connection using the close handshake defined in the <see cref="WebSocket"/> protocol specification section 7.
         /// </summary>
         /// <param name="closeStatus">Indicates the reason for closing the WebSocket connection.</param>
         /// <param name="statusDescription">Specifies a human readable explanation as to why the connection is closed.</param>
         /// <remarks>
-        /// WebSocketCloseStatus.EndpointUnavailable will close the WebSocket synchronous without awaiting response.
+        /// <para>
+        /// This method doesn't block. The close message is sent after the messages already queued and <see cref="ConnectionClosed"/> is raised
+        /// when the remote endpoint answers the close message, or after <see cref="ServerTimeout"/> if it doesn't.
+        /// </para>
+        /// <para>
+        /// <see cref="WebSocketCloseStatus.EndpointUnavailable"/> will close the <see cref="WebSocket"/> synchronous without awaiting response.
+        /// The call blocks until the close message is sent, waiting at most <see cref="ServerTimeout"/> (5 seconds if it's infinite).
+        /// If that time expires first, the connection is closed anyway and the close message may not have been sent.
+        /// Messages still queued may not be sent.
+        /// </para>
+        /// <para>
+        /// Only has effect if the connection is open. With <see cref="WebSocketCloseStatus.Empty"/> no status code is sent, so <paramref name="statusDescription"/> is ignored.
+        /// </para>
         /// </remarks>
-        public void Close(WebSocketCloseStatus closeStatus = WebSocketCloseStatus.Empty, string statusDescription = null)
+        public void Close(
+            WebSocketCloseStatus closeStatus = WebSocketCloseStatus.Empty,
+            string statusDescription = null)
         {
-            if (State != WebSocketState.Open)
-            {
-                //already closing or closed
-                return;
-            }
-
             bool closeImmediately = closeStatus == WebSocketCloseStatus.EndpointUnavailable;
 
-            // state transition (and ClosingTime) is done atomically by TryBeginClose
-            // a normal close is sent after the messages already queued, so these aren't lost
-            RawClose(closeStatus, statusDescription == null ? null : Encoding.UTF8.GetBytes(statusDescription), closeImmediately, !closeImmediately);
+            // state check and transition (and ClosingTime) are done atomically by TryBeginClose
+            RawClose(
+                closeStatus,
+                statusDescription == null ? null : Encoding.UTF8.GetBytes(statusDescription),
+                closeImmediately,
+                !closeImmediately,
+                true);
         }
 
         /// <summary>
@@ -253,11 +265,19 @@ public void Abort()
             HardClose(); 
         }
 
-        // CloseImediately will Send a close message and not await this message.
+        // CloseImediately will Send a close message and not await the answer from the remote endpoint:
+        // it waits at most ServerTimeout (5 seconds if infinite) for the close message to be sent, then closes the connection,
+        // even if the close message wasn't sent. Messages still queued may not be sent.
         // afterPendingMessages will send the close message after the messages already queued, instead of before them.
-        internal void RawClose(WebSocketCloseStatus closeStatus = WebSocketCloseStatus.Empty, byte[] buffer = null, bool CloseImmediately = false, bool afterPendingMessages = false)
+        // requireOpen will only close an open connection (not one answering a close message from the remote endpoint).
+        internal void RawClose(
+            WebSocketCloseStatus closeStatus = WebSocketCloseStatus.Empty,
+            byte[] buffer = null,
+            bool CloseImmediately = false,
+            bool afterPendingMessages = false,
+            bool requireOpen = false)
         {
-            if (!QueueCloseFrame(closeStatus, buffer, afterPendingMessages))
+            if (!QueueCloseFrame(closeStatus, buffer, afterPendingMessages, requireOpen))
             {
                 //already closing or closed
                 return;
@@ -296,9 +316,11 @@ internal void WaitForCloseMessageSent()
 
         // Non blocking close, used by the timeout checker.
         // The connection is hard closed by a later timeout check, once the close message is sent or ServerTimeout expires.
-        internal void BeginClose(WebSocketCloseStatus closeStatus, byte[] buffer)
+        internal void BeginClose(
+            WebSocketCloseStatus closeStatus,
+            byte[] buffer)
         {
-            if (QueueCloseFrame(closeStatus, buffer, false))
+            if (QueueCloseFrame(closeStatus, buffer, false, false))
             {
                 lock (_stateLock)
                 {
@@ -307,9 +329,13 @@ internal void BeginClose(WebSocketCloseStatus closeStatus, byte[] buffer)
             }
         }
 
-        private bool QueueCloseFrame(WebSocketCloseStatus closeStatus, byte[] buffer, bool afterPendingMessages)
+        private bool QueueCloseFrame(
+            WebSocketCloseStatus closeStatus,
+            byte[] buffer,
+            bool afterPendingMessages,
+            bool requireOpen)
         {
-            if (!TryBeginClose(closeStatus))
+            if (!TryBeginClose(closeStatus, requireOpen))
             {
                 return false;
             }
@@ -348,12 +374,14 @@ private bool QueueCloseFrame(WebSocketCloseStatus closeStatus, byte[] buffer, bo
         }
 
         // Atomically moves to CloseSent, making sure only one close message is ever queued.
-        private bool TryBeginClose(WebSocketCloseStatus closeStatus)
+        private bool TryBeginClose(
+            WebSocketCloseStatus closeStatus,
+            bool requireOpen)
         {
             lock (_stateLock)
             {
                 if (_closeFrameQueued
-                    || !(State == WebSocketState.Open || State == WebSocketState.CloseReceived))
+                    || !(State == WebSocketState.Open || (!requireOpen && State == WebSocketState.CloseReceived)))
                 {
                     return false;
                 }
@@ -371,10 +399,12 @@ private bool TryBeginClose(WebSocketCloseStatus closeStatus)
 
         // Called when the peer sent a close message.
         // Returns true if the close message has to be answered, false if the connection is already closing and can be hard closed.
-        internal bool TryMarkCloseReceived()
+        internal bool TryMarkCloseReceived(WebSocketCloseStatus peerCloseStatus)
         {
             lock (_stateLock)
             {
+                CloseStatus = peerCloseStatus;
+
                 if (_closeFrameQueued
                     || !(State == WebSocketState.Open || State == WebSocketState.CloseReceived))
                 {

WebSockets/WebSocketReceiver.cs

@@ -182,6 +182,7 @@ internal byte[] ReadBuffer(int messageLength, byte[] masks = null)
         private ReceiveMessageFrame SetMessageError(ReceiveMessageFrame frame, string errorMsg, WebSocketCloseStatus closeCode)
         {
             frame.ErrorMessage = errorMsg;
+            frame.CloseStatus = closeCode;
             return frame;
         }
 

WebSockets/WebSocketSender.cs

@@ -29,6 +29,8 @@ internal class WebSocketSender
 
         internal bool CloseMessageSent { get; private set; } = false;
 
+        private bool _stopRequested = false;
+
         internal bool ControlMessagesPresent
         {
             get
@@ -72,6 +74,8 @@ internal void QueueMessage(SendMessageFrame sendMessage, bool afterPendingMessag
 
         internal void StopSender()
         {
+            // set before signaling, the event is shared with QueueMessage so a signal alone doesn't mean stop
+            _stopRequested = true;
             CloseMessageSent = true;
 
             _shutdownEvent.Set();
@@ -140,6 +144,12 @@ private bool ProcessMessageFrames()
 
                     for (int i = 0; i < numberOfFrames; i++)
                     {
+                        if (CloseMessageSent)
+                        {
+                            // close message was sent (or sender stopped) while sending fragments, stop sending
+                            break;
+                        }
+
                         //start frame fin = 0 Opcode normal
                         if (i == 0) 
                         {
@@ -225,7 +235,11 @@ private void OnCloseFrameSent()
             CloseMessageSent = true;
 
             // wait until resources can be released.
-            _shutdownEvent.WaitOne();
+            // loop because the event can still be signaled by a message queued before the close
+            while (!_stopRequested)
+            {
+                _shutdownEvent.WaitOne();
+            }
         }