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.Ledger;
 using Neo.Network.P2P.Payloads;
@@ -8,6 +9,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;
@@ -22,6 +24,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();
@@ -104,7 +107,7 @@ private static IPEndPoint GetIPEndpointFromHostPort(string hostNameOrAddress, in
             IPHostEntry entry;
             try
             {
-                entry = Dns.GetHostEntry(hostNameOrAddress);
+                entry = System.Net.Dns.GetHostEntry(hostNameOrAddress);
             }
             catch (SocketException)
             {
@@ -129,6 +132,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;
@@ -162,6 +196,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);
@@ -214,6 +263,70 @@ private void OnRelayDirectly(IInventory inventory)
 
         private void OnSendDirectly(IInventory inventory) => SendToRemoteNodes(inventory);
 
+        protected override void OnTcpConnected(IPEndPoint remote, IPEndPoint local)
+        {
+            // Pre-check includes MaxConnections, MaxConnectionsPerAddress. If the check fails, it'll seed the error message.
+            Tcp.Message errorMsg = null;
+            if (MaxConnections != -1 && ConnectedPeers.Count >= MaxConnections && !TrustedIpAddresses.Contains(remote.Address))
+            {
+                Message reason = CreateDisconnectMessage(DisconnectReason.MaxConnectionReached);
+                errorMsg = Tcp.Write.Create(ByteString.FromBytes(reason.ToArray()));
+            }
+            ConnectedAddresses.TryGetValue(remote.Address, out int count);
+            if (count >= MaxConnectionsPerAddress)
+            {
+                Message reason = CreateDisconnectMessage(DisconnectReason.MaxConnectionPerAddressReached);
+                errorMsg = Tcp.Write.Create(ByteString.FromBytes(reason.ToArray()));
+            }
+            if (errorMsg != null)
+            {
+                Sender.Tell(new Tcp.Register(ActorRefs.Nobody));
+                Sender.Ask(errorMsg).ContinueWith(t => Sender.Tell(Tcp.Abort.Instance));
+                return;
+            }
+
+            base.OnTcpConnected(remote, local);
+        }
+
+        protected override void OnWsConnected(WebSocket ws, IPEndPoint remote, IPEndPoint local)
+        {
+            // Pre-check includes MaxConnectionsPerAddress. If the check fails, it'll send the error message.
+            ConnectedAddresses.TryGetValue(remote.Address, out int count);
+            if (count >= MaxConnectionsPerAddress)
+            {
+                var disconnectMessage = CreateDisconnectMessage(DisconnectReason.MaxConnectionPerAddressReached);
+                var errorMsg = new ArraySegment<byte>(disconnectMessage.ToArray());
+                ws.SendAsync(errorMsg, WebSocketMessageType.Binary, true, CancellationToken.None).PipeTo(Self,
+                    failure: ex => new Tcp.ErrorClosed(ex.Message));
+                ws.CloseAsync(WebSocketCloseStatus.NormalClosure, "close ws", CancellationToken.None);
+                return;
+            }
+
+            base.OnWsConnected(ws, remote, local);
+        }
+
+        /// <summary>
+        /// Create disconnect message with reason
+        /// </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;
+        }
+
         public static Props Props(NeoSystem system)
         {
             return Akka.Actor.Props.Create(() => new LocalNode(system));

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
@@ -35,7 +37,7 @@ private class WsConnected { public WebSocket Socket; public IPEndPoint Remote; p
         protected ActorSelection Connections => Context.ActorSelection("connection_*");
 
         private static readonly HashSet<IPAddress> localAddresses = new HashSet<IPAddress>();
-        private readonly Dictionary<IPAddress, int> ConnectedAddresses = new Dictionary<IPAddress, int>();
+        protected readonly Dictionary<IPAddress, int> ConnectedAddresses = new Dictionary<IPAddress, int>();
         /// <summary>
         /// A dictionary that stores the connected nodes.
         /// </summary>
@@ -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;
 
@@ -205,28 +208,15 @@ private void OnStart(ChannelsConfig config)
         /// </summary>
         /// <param name="remote">The remote endpoint of TCP connection.</param>
         /// <param name="local">The local endpoint of TCP connection.</param>
-        private void OnTcpConnected(IPEndPoint remote, IPEndPoint local)
+        protected virtual void OnTcpConnected(IPEndPoint remote, IPEndPoint local)
         {
             ImmutableInterlocked.Update(ref ConnectingPeers, p => p.Remove(remote));
-            if (MaxConnections != -1 && ConnectedPeers.Count >= MaxConnections && !TrustedIpAddresses.Contains(remote.Address))
-            {
-                Sender.Tell(Tcp.Abort.Instance);
-                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);
-            }
+            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);
         }
 
         /// <summary>
@@ -273,18 +263,11 @@ private void OnTimer()
             }
         }
 
-        private void OnWsConnected(WebSocket ws, IPEndPoint remote, IPEndPoint local)
+        protected virtual 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()}");
-            }
+            ConnectedAddresses[remote.Address] = count + 1;
+            Context.ActorOf(ProtocolProps(ws, remote, local), $"connection_{Guid.NewGuid()}");
         }
 
         protected override void PostStop()

src/neo/Network/P2P/ProtocolHandler.cs

@@ -69,6 +69,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)
@@ -181,13 +187,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;
             Context.Parent.Tell(Message.Create(MessageCommand.Addr, AddrPayload.Create(networkAddresses)));
         }
@@ -371,6 +371,31 @@ private void OnVersionMessageReceived(VersionPayload payload)
             Context.Parent.Tell(payload);
         }
 
+        /// <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 OnTimer()
         {
             RefreshPendingKnownHashes();