fixed a problem about close status

src/WebSocket4Net/InternalsVisibleTo.cs

@@ -0,0 +1,3 @@
+using System.Runtime.CompilerServices;
+
+[assembly: InternalsVisibleTo("WebSocket4Net.Tests")]

src/WebSocket4Net/WebSocket.cs

@@ -111,6 +111,12 @@ public void AddSubProtocol(string protocol)
             subProtocols.Add(protocol);
         }
 
+        protected override void SetupConnection(IConnection connection)
+        {
+            this.Closed += OnConnectionClosed;
+            base.SetupConnection(connection);
+        }
+
         public async ValueTask<bool> OpenAsync(CancellationToken cancellationToken = default)
         {
             State = WebSocketState.Connecting;
@@ -197,10 +203,12 @@ private void WriteHandshakeRequest(PipeWriter writer, string secKey)
 
         public new async ValueTask<WebSocketPackage> ReceiveAsync()
         {
-            return await ReceiveAsync(false);
+            return await ReceiveAsync(
+                handleControlPackage: true,
+                returnControlPackage: false);
         }
 
-        public async ValueTask<WebSocketPackage> ReceiveAsync(bool returnControlPackage)
+        internal async ValueTask<WebSocketPackage> ReceiveAsync(bool handleControlPackage, bool returnControlPackage)
         {
             var package = await base.ReceiveAsync();
 
@@ -209,11 +217,14 @@ public async ValueTask<WebSocketPackage> ReceiveAsync(bool returnControlPackage)
 
             if (package.OpCode != OpCode.Binary && package.OpCode != OpCode.Text && package.OpCode != OpCode.Handshake)
             {
-                await HandleControlPackage(package);
+                if (handleControlPackage)
+                {
+                    await HandleControlPackage(package);
+                }
 
                 if (!returnControlPackage)
                 {
-                    return await ReceiveAsync(returnControlPackage);
+                    return await ReceiveAsync(handleControlPackage, returnControlPackage);
                 }
             }
 
@@ -236,7 +247,7 @@ private async ValueTask HandleControlPackage(WebSocketPackage package)
             switch (package.OpCode)
             {
                 case (OpCode.Close):
-                    HandleCloseHandshake(package);
+                    await HandleCloseHandshake(package);
                     break;
 
                 case (OpCode.Ping):
@@ -329,53 +340,68 @@ public async ValueTask CloseAsync(CloseReason closeReason, string message)
 
             State = WebSocketState.CloseSent;
 
-            var closeHandshakeResponse = await ReceiveAsync(true);
+            var closeHandshakeResponse = await ReceiveAsync(
+                handleControlPackage: false,
+                returnControlPackage: true);
 
             if (closeHandshakeResponse.OpCode != OpCode.Close)
             {
                 OnError($"Unexpected close package, OpCode: {closeHandshakeResponse.OpCode}");
             }
             else
             {
-                HandleCloseHandshake(closeHandshakeResponse);
+                await HandleCloseHandshake(closeHandshakeResponse);
             }
 
             await base.CloseAsync();
 
             State = WebSocketState.Closed;
         }
 
-        private void HandleCloseHandshake(WebSocketPackage receivedClosePackage)
+        private CloseStatus DecodeCloseStatus(WebSocketPackage closePackage)
         {
-            var reader = new SequenceReader<byte>(receivedClosePackage.Data);
-
-            reader.TryReadBigEndian(out ushort closeReason);
+            var reader = new SequenceReader<byte>(closePackage.Data);
+            reader.TryReadBigEndian(out ushort closeReason);            
             var reasonText = reader.ReadString(_utf8Encoding);
 
-            var closeStatus = CloseStatus;
-
-            if (closeStatus == null)
-            {
-                CloseStatus = new CloseStatus
+            return new CloseStatus
                 {
                     Reason = (CloseReason)closeReason,
                     ReasonText = reasonText
                 };
+        }
 
-                return;
-            }
+        private async ValueTask HandleCloseHandshake(WebSocketPackage receivedClosePackage)
+        {
+            var closeStatusFromRemote = DecodeCloseStatus(receivedClosePackage);
+
+            var closeStatus = CloseStatus;
 
-            if (closeStatus.Reason != (CloseReason)closeReason)
+            if (closeStatus == null)
             {
-                OnError("Unmatched CloseReason");
-                return;
+                this.State = WebSocketState.CloseReceived;
+                closeStatusFromRemote.RemoteInitiated = true;
+                CloseStatus = closeStatusFromRemote;
+                // Send close pong message to server side.
+                await SendAsync(receivedClosePackage);
             }
-
-            if (closeStatus.ReasonText != reasonText)
+            else
             {
-                OnError("Unmatched CloseReasonText");
-                return;
+                if (closeStatus.Reason != closeStatusFromRemote.Reason)
+                {
+                    OnError("Unmatched CloseReason");
+                    return;
+                }
+
+                // Received the close pong message from server,
+                // so we can close the connection safely now
+                await base.CloseAsync();
             }
         }
+
+        private void OnConnectionClosed(object sender, EventArgs eventArgs)
+        {
+            this.State = WebSocketState.Closed;
+        }
     }
 }

test/WebSocket4Net.Tests/MainTest.cs

@@ -199,7 +199,9 @@ await session.SendAsync(new WebSocketPackage
                     Data = new ReadOnlySequence<byte>(Utf8Encoding.GetBytes("Hello"))
                 });
 
-                var package = await websocket.ReceiveAsync(true);
+                var package = await websocket.ReceiveAsync(
+                    handleControlPackage: true,
+                    returnControlPackage: true);
 
                 Assert.NotNull(package);
 
@@ -363,7 +365,7 @@ public async Task TestCloseWebSocket(Type hostConfiguratorType)
                         {
                             if (p.Message == "QUIT")
                             {
-                                await s.CloseAsync();
+                                await s.CloseAsync(CloseReason.NormalClosure);
                             }
                         });
 
@@ -392,7 +394,9 @@ public async Task TestCloseWebSocket(Type hostConfiguratorType)
 
                 await websocket.SendAsync("QUIT");
 
-                Assert.True(manualResetEvent.WaitOne(TimeSpan.FromSeconds(30)), "The connection failed to close on time");
+                await Task.WhenAny(websocket.ReceiveAsync().AsTask(), Task.Delay(TimeSpan.FromSeconds(30)));
+
+                Assert.True(manualResetEvent.WaitOne(TimeSpan.FromSeconds(5)), "The connection failed to close on time");
                 Assert.Equal(WebSocketState.Closed, websocket.State);
 
                 await server.StopAsync();