Merge pull request #378 from longfin/bugfix/377

Fix #377
planetarium · Aug 2, 2019 · 21fef6a · 21fef6a
2 parents 6416996 + 6f4e42c
commit 21fef6a
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 8 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -51,6 +51,8 @@ To be released.
     `Address`.  [[#369]]
  -  Removed unnecessary mutex in `Swarm<T>` to avoid continuous delays in peer
     registration in some situations.  [[#375]]
+ -  Fixed a bug that `TurnClient` had thrown `KeyNotFoundException` and
+    `IOException` on startup.  [[#377], [#378]]
 
 [#319]: https://github.com/planetarium/libplanet/issues/319
 [#343]: https://github.com/planetarium/libplanet/pull/343
@@ -61,6 +63,8 @@ To be released.
 [#367]: https://github.com/planetarium/libplanet/pull/367
 [#369]: https://github.com/planetarium/libplanet/pull/369
 [#375]: https://github.com/planetarium/libplanet/pull/375
+[#377]: https://github.com/planetarium/libplanet/issues/377
+[#378]: https://github.com/planetarium/libplanet/pull/378
 
 
 Version 0.4.1

diff --git a/Libplanet.Stun/Stun/TurnClient.cs b/Libplanet.Stun/Stun/TurnClient.cs
@@ -70,7 +70,7 @@ public async Task<IPEndPoint> AllocateRequestAsync(
             {
                 var request = new AllocateRequest((int)lifetime.TotalSeconds);
                 await SendMessageAsync(stream, request, cancellationToken);
-                response = await _responses[request.TransactionId].Task;
+                response = await ReceiveMessage(request.TransactionId);
 
                 if (response is AllocateErrorResponse allocError)
                 {
@@ -101,8 +101,7 @@ public async Task CreatePermissionAsync(
             NetworkStream stream = _control.GetStream();
             var request = new CreatePermissionRequest(peerAddress);
             await SendMessageAsync(stream, request, cancellationToken);
-            StunMessage response =
-                await _responses[request.TransactionId].Task;
+            StunMessage response = await ReceiveMessage(request.TransactionId);
 
             if (response is CreatePermissionErrorResponse)
             {
@@ -154,8 +153,7 @@ public async Task<IPEndPoint> GetMappedAddressAsync(
             NetworkStream stream = _control.GetStream();
             var request = new BindingRequest();
             await SendMessageAsync(stream, request, cancellationToken);
-            StunMessage response =
-                await _responses[request.TransactionId].Task;
+            StunMessage response = await ReceiveMessage(request.TransactionId);
 
             if (response is BindingSuccessResponse success)
             {
@@ -177,8 +175,7 @@ public async Task<TimeSpan> RefreshAllocationAsync(
             var request = new RefreshRequest((int)lifetime.TotalSeconds);
             await SendMessageAsync(stream, request, cancellationToken);
 
-            StunMessage response =
-                await _responses[request.TransactionId].Task;
+            StunMessage response = await ReceiveMessage(request.TransactionId);
             if (response is RefreshSuccessResponse success)
             {
                 return TimeSpan.FromSeconds(success.Lifetime);
@@ -252,7 +249,6 @@ private async Task ProcessMessage()
                         out TaskCompletionSource<StunMessage> tcs))
                     {
                         tcs.SetResult(message);
-                        _responses.Remove(message.TransactionId);
                     }
                 }
                 catch (Exception e)
@@ -274,6 +270,14 @@ private async Task EnsureConnection()
             }
         }
 
+        private async Task<StunMessage> ReceiveMessage(byte[] transactionId)
+        {
+            StunMessage response = await _responses[transactionId].Task;
+            _responses.Remove(transactionId);
+
+            return response;
+        }
+
         private class ByteArrayComparer : IEqualityComparer<byte[]>
         {
             public bool Equals(byte[] left, byte[] right)