Fix duplicate connection (#1685)

* fix duplicate connection * fix comments * Rename to AllowNewConnection * fix * Fix comment * fix * Rename * Remove useless check * Update * Update LocalNode.cs Co-authored-by: Shargon <shargon@gmail.com> Co-authored-by: erikzhang <erik@neo.org>
neo-project · Jun 5, 2020 · 44ccf0b · 44ccf0b
1 parent fc67879
commit 44ccf0b
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 7 deletions.
diff --git a/src/neo/Network/P2P/LocalNode.cs b/src/neo/Network/P2P/LocalNode.cs
@@ -127,6 +127,28 @@ internal static IPEndPoint GetIpEndPoint(string hostAndPort)
             return null;
         }
 
+        /// <summary>
+        /// Check the new connection <br/>
+        /// If it is equal to the Nonce of local or any remote node, it'll return false, else we'll return true and update the Listener address of the connected remote node.
+        /// </summary>
+        /// <param name="actor">Remote node actor</param>
+        /// <param name="node">Remote node</param>
+        public bool AllowNewConnection(IActorRef actor, RemoteNode node)
+        {
+            if (node.Version.Magic != ProtocolSettings.Default.Magic) return false;
+            if (node.Version.Nonce == Nonce) return false;
+
+            // filter duplicate connections
+            foreach (var other in RemoteNodes.Values)
+                if (other != node && other.Remote.Address.Equals(node.Remote.Address) && other.Version?.Nonce == node.Version.Nonce)
+                    return false;
+
+            if (node.Remote.Port != node.ListenerTcpPort && node.ListenerTcpPort != 0)
+                ConnectedPeers.TryUpdate(actor, node.Listener, node.Remote);
+
+            return true;
+        }
+
         public IEnumerable<RemoteNode> GetRemoteNodes()
         {
             return RemoteNodes.Values;

diff --git a/src/neo/Network/P2P/RemoteNode.ProtocolHandler.cs b/src/neo/Network/P2P/RemoteNode.ProtocolHandler.cs
@@ -356,12 +356,7 @@ private void OnVersionMessageReceived(VersionPayload payload)
                         break;
                 }
             }
-            if (payload.Nonce == LocalNode.Nonce || payload.Magic != ProtocolSettings.Default.Magic)
-            {
-                Disconnect(true);
-                return;
-            }
-            if (LocalNode.Singleton.RemoteNodes.Values.Where(p => p != this).Any(p => p.Remote.Address.Equals(Remote.Address) && p.Version?.Nonce == payload.Nonce))
+            if (!LocalNode.Singleton.AllowNewConnection(Self, this))
             {
                 Disconnect(true);
                 return;

diff --git a/tests/neo.UnitTests/Network/P2P/UT_RemoteNode.cs b/tests/neo.UnitTests/Network/P2P/UT_RemoteNode.cs
@@ -6,6 +6,7 @@
 using Neo.Network.P2P;
 using Neo.Network.P2P.Capabilities;
 using Neo.Network.P2P.Payloads;
+using System.Net;
 
 namespace Neo.UnitTests.Network.P2P
 {
@@ -54,7 +55,7 @@ public void RemoteNode_Test_Abort_DifferentMagic()
         public void RemoteNode_Test_Accept_IfSameMagic()
         {
             var connectionTestProbe = CreateTestProbe();
-            var remoteNodeActor = ActorOfAsTestActorRef(() => new RemoteNode(testBlockchain, connectionTestProbe, null, null));
+            var remoteNodeActor = ActorOfAsTestActorRef(() => new RemoteNode(testBlockchain, connectionTestProbe, new IPEndPoint(IPAddress.Parse("192.168.1.2"), 8080), new IPEndPoint(IPAddress.Parse("192.168.1.1"), 8080)));
 
             var msg = Message.Create(MessageCommand.Version, new VersionPayload()
             {