fix: accept after disconnect

MirageNet · Nov 2, 2020 · 3d06e8a · 3d06e8a
1 parent 421ff95
commit 3d06e8a
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 7 deletions.
diff --git a/Assets/Mirror/Runtime/Transport/Kcp/KcpTransport.cs b/Assets/Mirror/Runtime/Transport/Kcp/KcpTransport.cs
@@ -20,7 +20,7 @@ public class KcpTransport : Transport
 
         public KcpDelayMode delayMode = KcpDelayMode.Normal;
         internal readonly Dictionary<IPEndPoint, KcpServerConnection> connectedClients = new Dictionary<IPEndPoint, KcpServerConnection>(new IPEndpointComparer());
-        readonly Channel<KcpServerConnection> acceptedConnections = Cysharp.Threading.Tasks.Channel.CreateSingleConsumerUnbounded<KcpServerConnection>();
+        Channel<KcpServerConnection> acceptedConnections;
 
         public override IEnumerable<string> Scheme => new[] { "kcp" };
 
@@ -41,6 +41,7 @@ public override UniTask ListenAsync()
                 DualMode = true
             };
             socket.Bind(new IPEndPoint(IPAddress.IPv6Any, Port));
+            acceptedConnections = Cysharp.Threading.Tasks.Channel.CreateSingleConsumerUnbounded<KcpServerConnection>();
             return UniTask.CompletedTask;
         }
 
@@ -136,6 +137,7 @@ public override void Disconnect()
         {
             socket?.Close();
             socket = null;
+            acceptedConnections.Writer.TryComplete();
         }
 
         /// <summary>
@@ -146,11 +148,21 @@ public override void Disconnect()
         /// <returns>The connection to a client</returns>
         public override async UniTask<IConnection> AcceptAsync()
         {
-            KcpServerConnection connection = await acceptedConnections.Reader.ReadAsync();
+            if (socket == null)
+                return null;
 
-            await connection.HandshakeAsync();
+            try
+            {
+                KcpServerConnection connection = await acceptedConnections.Reader.ReadAsync();
+
+                await connection.HandshakeAsync();
 
-            return connection;
+                return connection;
+            }
+            catch (ChannelClosedException)
+            {
+                return null;
+            }
         }
 
         /// <summary>

diff --git a/Assets/Tests/Runtime/Transport/TransportTests.cs b/Assets/Tests/Runtime/Transport/TransportTests.cs
@@ -50,14 +50,30 @@ public AsyncTransportTests(string[] scheme, string uri, int port)
         });
 
 
-        [TearDown]
-        public void TearDown()
+        [UnityTearDown]
+        public IEnumerator TearDown() => UniTask.ToCoroutine(async () =>
         {
             clientConnection.Disconnect();
             serverConnection.Disconnect();
             transport.Disconnect();
+
+            try
+            {
+                // make sure we are done accepting,
+                // the transport might take a little bit of time to disconnect
+                while (await transport.AcceptAsync() != null)
+                {
+                    // fine,  just wait until transport stops accepting
+                }
+
+            }
+            catch (Exception)
+            {
+                // fine,  just wait until it is done
+            }
+
             Object.DestroyImmediate(transportObj);
-        }
+        });
 
         #endregion