diff --git a/src/neo/Network/P2P/LocalNode.cs b/src/neo/Network/P2P/LocalNode.cs index ebdbca98fd..c7f0e472e7 100644 --- a/src/neo/Network/P2P/LocalNode.cs +++ b/src/neo/Network/P2P/LocalNode.cs @@ -1,6 +1,5 @@ using Akka.Actor; using Neo.IO; -using Neo.Ledger; using Neo.Network.P2P.Payloads; using System; using System.Collections.Concurrent; @@ -214,6 +213,11 @@ private void OnRelayDirectly(IInventory inventory) private void OnSendDirectly(IInventory inventory) => SendToRemoteNodes(inventory); + protected override void OnTcpConnected(IActorRef connection) + { + connection.Tell(new RemoteNode.StartProtocol()); + } + public static Props Props(NeoSystem system) { return Akka.Actor.Props.Create(() => new LocalNode(system)); diff --git a/src/neo/Network/P2P/Peer.cs b/src/neo/Network/P2P/Peer.cs index 8e7eb21bd2..dc921b9526 100644 --- a/src/neo/Network/P2P/Peer.cs +++ b/src/neo/Network/P2P/Peer.cs @@ -226,9 +226,14 @@ private void OnTcpConnected(IPEndPoint remote, IPEndPoint local) Context.Watch(connection); Sender.Tell(new Tcp.Register(connection)); ConnectedPeers.TryAdd(connection, remote); + OnTcpConnected(connection); } } + protected virtual void OnTcpConnected(IActorRef connection) + { + } + /// /// Will be triggered when a Tcp.CommandFailed message is received. /// If it's a Tcp.Connect command, remove the related endpoint from ConnectingPeers. diff --git a/src/neo/Network/P2P/RemoteNode.cs b/src/neo/Network/P2P/RemoteNode.cs index 88b58163bf..b1457a8522 100644 --- a/src/neo/Network/P2P/RemoteNode.cs +++ b/src/neo/Network/P2P/RemoteNode.cs @@ -16,6 +16,7 @@ namespace Neo.Network.P2P { public partial class RemoteNode : Connection { + internal class StartProtocol { } internal class Relay { public IInventory Inventory; } private readonly NeoSystem system; @@ -35,16 +36,6 @@ public RemoteNode(NeoSystem system, object connection, IPEndPoint remote, IPEndP { this.system = system; LocalNode.Singleton.RemoteNodes.TryAdd(Self, this); - - var capabilities = new List - { - new FullNodeCapability(Blockchain.Singleton.Height) - }; - - if (LocalNode.Singleton.ListenerTcpPort > 0) capabilities.Add(new ServerCapability(NodeCapabilityType.TcpServer, (ushort)LocalNode.Singleton.ListenerTcpPort)); - if (LocalNode.Singleton.ListenerWsPort > 0) capabilities.Add(new ServerCapability(NodeCapabilityType.WsServer, (ushort)LocalNode.Singleton.ListenerWsPort)); - - SendMessage(Message.Create(MessageCommand.Version, VersionPayload.Create(LocalNode.Nonce, LocalNode.UserAgent, capabilities.ToArray()))); } /// @@ -141,6 +132,9 @@ protected override void OnReceive(object message) case Relay relay: OnRelay(relay.Inventory); break; + case StartProtocol _: + OnStartProtocol(); + break; } } @@ -166,6 +160,19 @@ private void OnSend(IInventory inventory) EnqueueMessage((MessageCommand)inventory.InventoryType, inventory); } + private void OnStartProtocol() + { + var capabilities = new List + { + new FullNodeCapability(Blockchain.Singleton.Height) + }; + + if (LocalNode.Singleton.ListenerTcpPort > 0) capabilities.Add(new ServerCapability(NodeCapabilityType.TcpServer, (ushort)LocalNode.Singleton.ListenerTcpPort)); + if (LocalNode.Singleton.ListenerWsPort > 0) capabilities.Add(new ServerCapability(NodeCapabilityType.WsServer, (ushort)LocalNode.Singleton.ListenerWsPort)); + + SendMessage(Message.Create(MessageCommand.Version, VersionPayload.Create(LocalNode.Nonce, LocalNode.UserAgent, capabilities.ToArray()))); + } + protected override void PostStop() { timer.CancelIfNotNull(); diff --git a/tests/neo.UnitTests/Network/P2P/UT_RemoteNode.cs b/tests/neo.UnitTests/Network/P2P/UT_RemoteNode.cs index 297af38e0c..61702115f3 100644 --- a/tests/neo.UnitTests/Network/P2P/UT_RemoteNode.cs +++ b/tests/neo.UnitTests/Network/P2P/UT_RemoteNode.cs @@ -31,8 +31,6 @@ public void RemoteNode_Test_Abort_DifferentMagic() var connectionTestProbe = CreateTestProbe(); var remoteNodeActor = ActorOfAsTestActorRef(() => new RemoteNode(testBlockchain, connectionTestProbe, null, null)); - connectionTestProbe.ExpectMsg(); - var msg = Message.Create(MessageCommand.Version, new VersionPayload { UserAgent = "".PadLeft(1024, '0'), @@ -58,8 +56,6 @@ public void RemoteNode_Test_Accept_IfSameMagic() var connectionTestProbe = CreateTestProbe(); var remoteNodeActor = ActorOfAsTestActorRef(() => new RemoteNode(testBlockchain, connectionTestProbe, null, null)); - connectionTestProbe.ExpectMsg(); - var msg = Message.Create(MessageCommand.Version, new VersionPayload() { UserAgent = "Unit Test".PadLeft(1024, '0'),