Adding disconnection reason and its use for ensuring better p2p performance

src/neo/Network/P2P/LocalNode.cs

@@ -1,4 +1,5 @@
 using Akka.Actor;
+using Akka.IO;
 using Neo.IO;
 using Neo.Network.P2P.Payloads;
 using System;
@@ -7,6 +8,7 @@
 using System.Linq;
 using System.Net;
 using System.Net.Sockets;
+using System.Net.WebSockets;
 using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
@@ -21,6 +23,7 @@ internal class SendDirectly { public IInventory Inventory; }
 
         public const uint ProtocolVersion = 0;
         private const int MaxCountFromSeedList = 5;
+        private const int PeersToReturnOnDisconnect = 10;
         private readonly IPEndPoint[] SeedList = new IPEndPoint[ProtocolSettings.Default.SeedList.Length];
 
         private static readonly object lockObj = new object();
@@ -103,7 +106,7 @@ private static IPEndPoint GetIPEndpointFromHostPort(string hostNameOrAddress, in
             IPHostEntry entry;
             try
             {
-                entry = Dns.GetHostEntry(hostNameOrAddress);
+                entry = System.Net.Dns.GetHostEntry(hostNameOrAddress);
             }
             catch (SocketException)
             {
@@ -128,6 +131,37 @@ internal static IPEndPoint GetIpEndPoint(string hostAndPort)
             return null;
         }
 
+        /// <summary>
+        /// Check duplicated duplicated Nonce. Usually it occurs when a new remote connection is established, which checks its counterpart's Nonce value. <br/>
+        /// If Nonce is the same we check if the Remote can be added to the known LocalAddresses.<br/>
+        /// If it is equal to the Nonce of other RemoteNode, we just return true, else we'll return false and update the Listener address of the connected remote node.
+        /// </summary>
+        /// <param name="remoteActor">Remote node actor</param>
+        /// <param name="remoteNode">Remote node</param>
+        public bool CheckDuplicateNonce(IActorRef remoteActor, RemoteNode remoteNode)
+        {
+            var version = remoteNode.Version;
+            var remote = remoteNode.Remote;
+
+            if (remote is null) return false;
+            if (version.Nonce == Nonce) return true;
+
+            foreach (var pair in RemoteNodes)
+            {
+                var otherNode = pair.Value;
+                if (otherNode != remoteNode && otherNode.Remote.Address.Equals(remote.Address) && otherNode.Version?.Nonce == version.Nonce)
+                {// filter duplicate connections
+                    return true;
+                }
+            }
+            if (remote.Port != remoteNode.ListenerTcpPort && remoteNode.ListenerTcpPort != 0)
+            {
+                ConnectedPeers.TryUpdate(remoteActor, remoteNode.Listener, remote);
+            }
+
+            return false;
+        }
+
         public IEnumerable<RemoteNode> GetRemoteNodes()
         {
             return RemoteNodes.Values;
@@ -161,6 +195,21 @@ protected override void NeedMorePeers(int count)
             }
         }
 
+        /// <summary>
+        /// Get a maximum number of count random peer currently connected with the node.
+        /// </summary>
+        /// <param name="count"></param>
+        /// <returns></returns>
+        public NetworkAddressWithTime[] GetRandomConnectedPeers(int count)
+        {
+            Random rand = new Random();
+            IEnumerable<RemoteNode> peers = RemoteNodes.Values
+                .Where(p => p.ListenerTcpPort > 0)
+                .OrderBy(p => rand.Next())
+                .Take(count);
+            return peers.Select(p => NetworkAddressWithTime.Create(p.Listener.Address, p.Version.Timestamp, p.Version.Capabilities)).ToArray();
+        }
+
         protected override void OnReceive(object message)
         {
             base.OnReceive(message);
@@ -213,6 +262,46 @@ private void OnRelayDirectly(IInventory inventory)
 
         private void OnSendDirectly(IInventory inventory) => SendToRemoteNodes(inventory);
 
+        /// <summary>
+        /// Disconnect tcp connection with reason message
+        /// </summary>
+        /// <param name="reason"></param>
+        protected override void TcpDisconnect(DisconnectReason reason)
+        {
+            var message = CreateDisconnectMessage(reason);
+            Sender.Ask(Tcp.Write.Create(ByteString.FromBytes(message.ToArray()))).ContinueWith(t => Sender.Tell(Tcp.Abort.Instance));
+        }
+
+        protected override void WsDisconnect(WebSocket ws, DisconnectReason reason)
+        {
+            var message = CreateDisconnectMessage(reason);
+            ws.SendAsync(new ArraySegment<byte>(message.ToArray()), WebSocketMessageType.Binary, true, CancellationToken.None).PipeTo(Self,
+                    failure: ex => new Tcp.ErrorClosed(ex.Message));
+            ws.CloseAsync(WebSocketCloseStatus.NormalClosure, "close ws", CancellationToken.None);
+        }
+
+        /// <summary>
+        /// Create disconnect message with reason and attached with connected peers' addresses.
+        /// </summary>
+        /// <param name="reason">Disconnect reason</param>
+        private Message CreateDisconnectMessage(DisconnectReason reason)
+        {
+            byte[] data;
+            switch (reason)
+            {
+                case DisconnectReason.MaxConnectionReached:
+                case DisconnectReason.MaxConnectionPerAddressReached:
+                    data = GetRandomConnectedPeers(PeersToReturnOnDisconnect).ToByteArray();
+                    break;
+                default:
+                    data = new byte[0];
+                    break;
+            }
+            var payload = DisconnectPayload.Create(reason, data);
+            var message = Message.Create(MessageCommand.Disconnect, payload);
+            return message;
+        }
+
         protected override void OnTcpConnected(IActorRef connection)
         {
             connection.Tell(new RemoteNode.StartProtocol());

src/neo/Network/P2P/MessageCommand.cs

@@ -18,6 +18,8 @@ public enum MessageCommand : byte
         Ping = 0x18,
         [ReflectionCache(typeof(PingPayload))]
         Pong = 0x19,
+        [ReflectionCache(typeof(DisconnectPayload))]
+        Disconnect = 0x1f,
 
         //synchronization
         [ReflectionCache(typeof(GetBlocksPayload))]

src/neo/Network/P2P/Payloads/DisconnectPayload.cs

@@ -0,0 +1,37 @@
+using Neo.IO;
+using System;
+using System.IO;
+
+namespace Neo.Network.P2P.Payloads
+{
+    public class DisconnectPayload : ISerializable
+    {
+        public const int MaxDataSize = 1024;
+
+        public DisconnectReason Reason;
+        public byte[] Data;
+
+        public int Size => sizeof(DisconnectReason) + Data.GetVarSize();
+
+        public static DisconnectPayload Create(DisconnectReason reason, byte[] data = null)
+        {
+            return new DisconnectPayload
+            {
+                Reason = reason,
+                Data = data ?? Array.Empty<byte>()
+            };
+        }
+
+        public void Deserialize(BinaryReader reader)
+        {
+            Reason = (DisconnectReason)reader.ReadByte();
+            Data = reader.ReadVarBytes(MaxDataSize);
+        }
+
+        public void Serialize(BinaryWriter writer)
+        {
+            writer.Write((byte)Reason);
+            writer.WriteVarBytes(Data);
+        }
+    }
+}

src/neo/Network/P2P/Payloads/DisconnectReason.cs

@@ -0,0 +1,14 @@
+namespace Neo.Network.P2P.Payloads
+{
+    public enum DisconnectReason : byte
+    {
+        MaxConnectionReached = 0x01,
+        MaxConnectionPerAddressReached = 0x02,
+        DuplicateNonce = 0x03,
+
+        MagicNumberIncompatible = 0x10,
+        FormatException = 0x11,
+
+        InternalError = 0x20,
+    }
+}

src/neo/Network/P2P/Peer.cs

@@ -4,6 +4,7 @@
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Http;
 using Neo.IO;
+using Neo.Network.P2P.Payloads;
 using System;
 using System.Buffers.Binary;
 using System.Collections.Concurrent;
@@ -14,6 +15,7 @@
 using System.Net.NetworkInformation;
 using System.Net.Sockets;
 using System.Net.WebSockets;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Neo.Network.P2P
@@ -92,6 +94,7 @@ protected void AddPeers(IEnumerable<IPEndPoint> peers)
         protected void ConnectToPeer(IPEndPoint endPoint, bool isTrusted = false)
         {
             endPoint = endPoint.Unmap();
+
             // If the address is the same, the ListenerTcpPort should be different, otherwise, return
             if (endPoint.Port == ListenerTcpPort && localAddresses.Contains(endPoint.Address)) return;
 
@@ -208,26 +211,28 @@ private void OnStart(ChannelsConfig config)
         private void OnTcpConnected(IPEndPoint remote, IPEndPoint local)
         {
             ImmutableInterlocked.Update(ref ConnectingPeers, p => p.Remove(remote));
+
+            // Pre-check includes MaxConnections, MaxConnectionsPerAddress. If the check fails, it'll seed the error message.
             if (MaxConnections != -1 && ConnectedPeers.Count >= MaxConnections && !TrustedIpAddresses.Contains(remote.Address))
             {
-                Sender.Tell(Tcp.Abort.Instance);
+                Sender.Tell(new Tcp.Register(ActorRefs.Nobody));
+                TcpDisconnect(DisconnectReason.MaxConnectionReached);
                 return;
             }
-
             ConnectedAddresses.TryGetValue(remote.Address, out int count);
             if (count >= MaxConnectionsPerAddress)
             {
-                Sender.Tell(Tcp.Abort.Instance);
-            }
-            else
-            {
-                ConnectedAddresses[remote.Address] = count + 1;
-                IActorRef connection = Context.ActorOf(ProtocolProps(Sender, remote, local), $"connection_{Guid.NewGuid()}");
-                Context.Watch(connection);
-                Sender.Tell(new Tcp.Register(connection));
-                ConnectedPeers.TryAdd(connection, remote);
-                OnTcpConnected(connection);
+                Sender.Tell(new Tcp.Register(ActorRefs.Nobody));
+                TcpDisconnect(DisconnectReason.MaxConnectionPerAddressReached);
+                return;
             }
+
+            ConnectedAddresses[remote.Address] = count + 1;
+            IActorRef connection = Context.ActorOf(ProtocolProps(Sender, remote, local), $"connection_{Guid.NewGuid()}");
+            Context.Watch(connection);
+            Sender.Tell(new Tcp.Register(connection));
+            ConnectedPeers.TryAdd(connection, remote);
+            OnTcpConnected(connection);
         }
 
         protected virtual void OnTcpConnected(IActorRef connection)
@@ -262,6 +267,11 @@ private void OnTerminated(IActorRef actorRef)
             }
         }
 
+        protected virtual void TcpDisconnect(DisconnectReason reason)
+        {
+            Sender.Tell(Tcp.Abort.Instance);
+        }
+
         private void OnTimer()
         {
             // Check if the number of desired connections is already enough
@@ -283,13 +293,17 @@ private void OnWsConnected(WebSocket ws, IPEndPoint remote, IPEndPoint local)
             ConnectedAddresses.TryGetValue(remote.Address, out int count);
             if (count >= MaxConnectionsPerAddress)
             {
-                ws.Abort();
-            }
-            else
-            {
-                ConnectedAddresses[remote.Address] = count + 1;
-                Context.ActorOf(ProtocolProps(ws, remote, local), $"connection_{Guid.NewGuid()}");
+                WsDisconnect(ws, DisconnectReason.MaxConnectionPerAddressReached);
+                return;
             }
+
+            ConnectedAddresses[remote.Address] = count + 1;
+            Context.ActorOf(ProtocolProps(ws, remote, local), $"connection_{Guid.NewGuid()}");
+        }
+
+        protected virtual void WsDisconnect(WebSocket ws, DisconnectReason reason)
+        {
+            ws.Abort();
         }
 
         protected override void PostStop()

src/neo/Network/P2P/RemoteNode.ProtocolHandler.cs

@@ -1,5 +1,6 @@
 using Akka.Actor;
 using Neo.Cryptography;
+using Neo.IO;
 using Neo.IO.Caching;
 using Neo.Ledger;
 using Neo.Network.P2P.Capabilities;
@@ -29,7 +30,6 @@ protected override UInt256 GetKeyForItem((UInt256, DateTime) item)
         private readonly PendingKnownHashesCollection pendingKnownHashes = new PendingKnownHashesCollection();
         private readonly HashSetCache<UInt256> knownHashes = new HashSetCache<UInt256>(Blockchain.Singleton.MemPool.Capacity * 2 / 5);
         private readonly HashSetCache<UInt256> sentHashes = new HashSetCache<UInt256>(Blockchain.Singleton.MemPool.Capacity * 2 / 5);
-        private bool verack = false;
         private BloomFilter bloom_filter;
 
         private static readonly TimeSpan TimerInterval = TimeSpan.FromSeconds(30);
@@ -42,6 +42,12 @@ private void OnMessage(Message msg)
             foreach (IP2PPlugin plugin in Plugin.P2PPlugins)
                 if (!plugin.OnP2PMessage(msg))
                     return;
+
+            if (msg.Command == MessageCommand.Disconnect)
+            {
+                OnDisconnectMessageReceived((DisconnectPayload)msg.Payload);
+                return;
+            }
             if (Version == null)
             {
                 if (msg.Command != MessageCommand.Version)
@@ -151,13 +157,7 @@ private void OnFilterLoadMessageReceived(FilterLoadPayload payload)
         /// </summary>
         private void OnGetAddrMessageReceived()
         {
-            Random rand = new Random();
-            IEnumerable<RemoteNode> peers = LocalNode.Singleton.RemoteNodes.Values
-                .Where(p => p.ListenerTcpPort > 0)
-                .GroupBy(p => p.Remote.Address, (k, g) => g.First())
-                .OrderBy(p => rand.Next())
-                .Take(AddrPayload.MaxCountToSend);
-            NetworkAddressWithTime[] networkAddresses = peers.Select(p => NetworkAddressWithTime.Create(p.Listener.Address, p.Version.Timestamp, p.Version.Capabilities)).ToArray();
+            NetworkAddressWithTime[] networkAddresses = LocalNode.Singleton.GetRandomConnectedPeers(AddrPayload.MaxCountToSend);
             if (networkAddresses.Length == 0) return;
             EnqueueMessage(Message.Create(MessageCommand.Addr, AddrPayload.Create(networkAddresses)));
         }
@@ -353,19 +353,44 @@ private void OnVersionMessageReceived(VersionPayload payload)
                         break;
                 }
             }
-            if (payload.Nonce == LocalNode.Nonce || payload.Magic != ProtocolSettings.Default.Magic)
+            if (Version.Magic != ProtocolSettings.Default.Magic)
             {
-                Disconnect(true);
+                DisconnectWithReason(DisconnectReason.MagicNumberIncompatible, BitConverter.GetBytes(ProtocolSettings.Default.Magic));
                 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.CheckDuplicateNonce(Self, this))
             {
-                Disconnect(true);
+                DisconnectWithReason(DisconnectReason.DuplicateNonce);
                 return;
             }
             SendMessage(Message.Create(MessageCommand.Verack));
         }
 
+        /// <summary>
+        /// Processe received disconnect messages and close the connection
+        /// </summary>
+        /// <param name="payload">Disconnect message</param>
+        private void OnDisconnectMessageReceived(DisconnectPayload payload)
+        {
+            switch (payload.Reason)
+            {
+                case DisconnectReason.MaxConnectionReached:
+                case DisconnectReason.MaxConnectionPerAddressReached:
+                    try
+                    {
+                        var addressList = payload.Data
+                            .AsSerializableArray<NetworkAddressWithTime>(AddrPayload.MaxCountToSend)
+                            .Select(p => p.EndPoint)
+                            .Where(p => p.Port > 0);
+                        system.LocalNode.Tell(new Peer.Peers { EndPoints = addressList });
+                    }
+                    catch { }
+                    break;
+                default: break;
+            }
+            Context.Stop(Self);
+        }
+
         private void RefreshPendingKnownHashes()
         {
             while (pendingKnownHashes.Count > 0)