diff --git a/src/neo/Network/P2P/LocalNode.cs b/src/neo/Network/P2P/LocalNode.cs index c1e7046725..c8bc937001 100644 --- 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; } + /// + /// Check the new connection
+ /// 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. + ///
+ /// Remote node actor + /// Remote node + 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 GetRemoteNodes() { return RemoteNodes.Values; diff --git a/src/neo/Network/P2P/RemoteNode.ProtocolHandler.cs b/src/neo/Network/P2P/RemoteNode.ProtocolHandler.cs index 3a55d84347..9e73809cc1 100644 --- 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 index 61702115f3..8f5b9c8000 100644 --- 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() {