<repo_name>siemantic/StratisBitcoinFullNode<file_sep>/src/Stratis.Bitcoin.Features.Consensus.Tests/PeerBanningTest.cs
﻿using System.IO;
using System.Linq;
using System.Net;
using System.Threading;
using System.Threading.Tasks;
using Moq;
using NBitcoin;
using Stratis.Bitcoin.Connection;
using Stratis.Bitcoin.P2P.Peer;
using Xunit;

namespace Stratis.Bitcoin.Features.Consensus.Tests
{
    public class PeerBanningTest
    {
        public PeerBanningTest()
        {
            Block.BlockSignature = false;
            Transaction.TimeStamp = false;
        }

        [Fact]
        public async Task NodeIsSynced_PeerSendsABadBlockAndPeerDiconnected_ThePeerGetsBanned_Async()
        {
            string dataDir = Path.Combine("TestData", nameof(PeerBanningTest), nameof(this.NodeIsSynced_PeerSendsABadBlockAndPeerDiconnected_ThePeerGetsBanned_Async));
            Directory.CreateDirectory(dataDir);

            TestChainContext context = await TestChainFactory.CreateAsync(Network.RegTest, dataDir);
            var peer = new IPEndPoint(IPAddress.Parse("172.16.31.10"), context.Network.DefaultPort);

            context.MockReadOnlyNodesCollection.Setup(s => s.FindByEndpoint(It.IsAny<IPEndPoint>())).Returns((NetworkPeer)null);

            var blocks = await TestChainFactory.MineBlocksAsync(context, 2, new Key().ScriptPubKey);
            var block = blocks.First();
            block.Header.HashPrevBlock = context.Chain.Tip.HashBlock;
            await context.Consensus.AcceptBlockAsync(new BlockValidationContext { Block = block, Peer = peer });

            Assert.True(context.PeerBanning.IsBanned(peer));
        }

        [Fact]
        public async Task NodeIsSynced_PeerSendsABadBlockAndPeerIsConnected_ThePeerGetsBanned_Async()
        {
            string dataDir = Path.Combine("TestData", nameof(PeerBanningTest), nameof(this.NodeIsSynced_PeerSendsABadBlockAndPeerIsConnected_ThePeerGetsBanned_Async));
            Directory.CreateDirectory(dataDir);

            TestChainContext context = await TestChainFactory.CreateAsync(Network.RegTest, dataDir);
            var peer = new IPEndPoint(IPAddress.Parse("1.2.3.4"), context.Network.DefaultPort);

            var connectionManagerBehavior = new ConnectionManagerBehavior(false, context.ConnectionManager, context.LoggerFactory);
            using (var node = new NetworkPeer(context.DateTimeProvider, context.LoggerFactory))
            {
                node.Behaviors.Add(connectionManagerBehavior);
                context.MockReadOnlyNodesCollection.Setup(s => s.FindByEndpoint(It.IsAny<IPEndPoint>())).Returns(node);

                var blocks = await TestChainFactory.MineBlocksAsync(context, 2, new Key().ScriptPubKey);
                // create a new block that breaks consensus.
                var block = blocks.First();
                block.Header.HashPrevBlock = context.Chain.Tip.HashBlock;
                await context.Consensus.AcceptBlockAsync(new BlockValidationContext { Block = block, Peer = peer });
                Assert.True(context.PeerBanning.IsBanned(peer));
            }
        }

        [Fact]
        public async Task NodeIsSynced_PeerSendsABadBlockAndPeerIsWhitelisted_ThePeerIsNotBanned_Async()
        {
            string dataDir = Path.Combine("TestData", nameof(PeerBanningTest), nameof(this.NodeIsSynced_PeerSendsABadBlockAndPeerIsWhitelisted_ThePeerIsNotBanned_Async));
            Directory.CreateDirectory(dataDir);

            TestChainContext context = await TestChainFactory.CreateAsync(Network.RegTest, dataDir);
            var peer = new IPEndPoint(IPAddress.Parse("1.2.3.4"), context.Network.DefaultPort);

            var connectionManagerBehavior = new ConnectionManagerBehavior(false, context.ConnectionManager, context.LoggerFactory) { Whitelisted = true };
            using (var node = new NetworkPeer(context.DateTimeProvider, context.LoggerFactory))
            {
                node.Behaviors.Add(connectionManagerBehavior);
                context.MockReadOnlyNodesCollection.Setup(s => s.FindByEndpoint(It.IsAny<IPEndPoint>())).Returns(node);

                var blocks = await TestChainFactory.MineBlocksAsync(context, 2, new Key().ScriptPubKey);
                // create a new block that breaks consensus.
                var block = blocks.First();
                block.Header.HashPrevBlock = context.Chain.Tip.HashBlock;
                await context.Consensus.AcceptBlockAsync(new BlockValidationContext { Block = block, Peer = peer });

                Assert.False(context.PeerBanning.IsBanned(peer));
            }
        }

        [Fact]
        public async Task NodeIsSynced_PeerSendsABadBlockAndErrorIsNotBanError_ThePeerIsNotBanned_Async()
        {
            string dataDir = Path.Combine("TestData", nameof(PeerBanningTest), nameof(this.NodeIsSynced_PeerSendsABadBlockAndErrorIsNotBanError_ThePeerIsNotBanned_Async));
            Directory.CreateDirectory(dataDir);

            TestChainContext context = await TestChainFactory.CreateAsync(Network.RegTest, dataDir);
            var peer = new IPEndPoint(IPAddress.Parse("1.2.3.4"), context.Network.DefaultPort);

            var connectionManagerBehavior = new ConnectionManagerBehavior(false, context.ConnectionManager, context.LoggerFactory) { Whitelisted = true };
            using (var node = new NetworkPeer(context.DateTimeProvider, context.LoggerFactory))
            {
                node.Behaviors.Add(connectionManagerBehavior);
                context.MockReadOnlyNodesCollection.Setup(s => s.FindByEndpoint(It.IsAny<IPEndPoint>())).Returns(node);

                var blocks = await TestChainFactory.MineBlocksAsync(context, 2, new Key().ScriptPubKey);
                // create a new block that breaks consensus.
                var block = blocks.First();
                block.Header.HashPrevBlock = context.Chain.Tip.HashBlock;
                await context.Consensus.AcceptBlockAsync(new BlockValidationContext { Block = block, Peer = peer, BanDurationSeconds = BlockValidationContext.BanDurationNoBan });

                Assert.False(context.PeerBanning.IsBanned(peer));
            }
        }

        [Fact]
        public async Task NodeIsSynced_PeerSendsABadBlockAndPeerIsBandAndBanIsExpired_ThePeerIsNotBanned_Async()
        {
            string dataDir = Path.Combine("TestData", nameof(PeerBanningTest), nameof(this.NodeIsSynced_PeerSendsABadBlockAndPeerIsBandAndBanIsExpired_ThePeerIsNotBanned_Async));
            Directory.CreateDirectory(dataDir);

            TestChainContext context = await TestChainFactory.CreateAsync(Network.RegTest, dataDir);
            var peer = new IPEndPoint(IPAddress.Parse("1.2.3.4"), context.Network.DefaultPort);

            var connectionManagerBehavior = new ConnectionManagerBehavior(false, context.ConnectionManager, context.LoggerFactory) { Whitelisted = true };
            using (var node = new NetworkPeer(context.DateTimeProvider, context.LoggerFactory))
            {
                node.Behaviors.Add(connectionManagerBehavior);
                context.MockReadOnlyNodesCollection.Setup(s => s.FindByEndpoint(It.IsAny<IPEndPoint>())).Returns(node);

                var blocks = await TestChainFactory.MineBlocksAsync(context, 2, new Key().ScriptPubKey);
                // create a new block that breaks consensus.
                var block = blocks.First();
                block.Header.HashPrevBlock = context.Chain.Tip.HashBlock;
                await context.Consensus.AcceptBlockAsync(new BlockValidationContext { Block = block, Peer = peer, BanDurationSeconds = 1 }); // ban for 1 second

                // wait 1 sec for ban to expire.
                Thread.Sleep(1000);

                Assert.False(context.PeerBanning.IsBanned(peer));
            }
        }

    }
}
<file_sep>/src/Stratis.Bitcoin.Tests/P2P/PeerConnectorTests.cs
﻿using System;
using System.IO;
using System.Net;
using NBitcoin;
using NBitcoin.Protocol;
using Stratis.Bitcoin.Configuration;
using Stratis.Bitcoin.Configuration.Logging;
using Stratis.Bitcoin.Configuration.Settings;
using Stratis.Bitcoin.Connection;
using Stratis.Bitcoin.P2P;
using Stratis.Bitcoin.P2P.Peer;
using Stratis.Bitcoin.Utilities;
using Xunit;

namespace Stratis.Bitcoin.Tests.P2P
{
    public sealed class PeerConnectorTests : TestBase
    {
        private readonly IAsyncLoopFactory asyncLoopFactory;
        private readonly ExtendedLoggerFactory extendedLoggerFactory;
        private readonly INetworkPeerFactory networkPeerFactory;
        private readonly NetworkPeerConnectionParameters networkPeerParameters;
        private readonly NodeLifetime nodeLifetime;
        private readonly Network network;

        public PeerConnectorTests()
        {
            this.extendedLoggerFactory = new ExtendedLoggerFactory();
            this.extendedLoggerFactory.AddConsoleWithFilters();

            this.asyncLoopFactory = new AsyncLoopFactory(this.extendedLoggerFactory);
            this.network = Network.Main;
            this.networkPeerParameters = new NetworkPeerConnectionParameters();
            this.networkPeerFactory = new NetworkPeerFactory(this.network, DateTimeProvider.Default, this.extendedLoggerFactory);
            this.nodeLifetime = new NodeLifetime();
        }

        [Fact]
        public void PeerConnectorAddNode_FindPeerToConnectTo_Returns_AddNodePeers()
        {
            var peerFolder = AssureEmptyDirAsDataFolder(Path.Combine(AppContext.BaseDirectory, "PeerConnectorTests"));
            var peerAddressManager = new PeerAddressManager(DateTimeProvider.Default, peerFolder, this.extendedLoggerFactory);

            var ipAddressOne = IPAddress.Parse("::ffff:192.168.0.1");
            var endpointAddNode = new IPEndPoint(ipAddressOne, 80);

            var ipAddressTwo = IPAddress.Parse("::ffff:192.168.0.2");
            var endpointDiscoveredNode = new IPEndPoint(ipAddressTwo, 80);

            peerAddressManager.AddPeer(endpointAddNode, IPAddress.Loopback);
            peerAddressManager.AddPeer(endpointDiscoveredNode, IPAddress.Loopback);

            var nodeSettings = new NodeSettings();
            nodeSettings.LoadArguments(new string[] { });
            var connectionSettings = new ConnectionManagerSettings();
            connectionSettings.Load(nodeSettings);

            connectionSettings.AddNode.Add(endpointAddNode);

            // TODO: Once we have an interface on NetworkPeer we can test this properly.
            //var connector = this.CreatePeerConnecterAddNode(nodeSettings, peerAddressManager);
            //connector.OnConnectAsync().GetAwaiter().GetResult();
            //Assert.Contains(networkAddressAddNode, connector.ConnectedPeers.Select(p => p.PeerAddress));
        }

        [Fact]
        public void PeerConnectorAddNode_CanAlwaysStart()
        {
            var peerFolder = AssureEmptyDirAsDataFolder(Path.Combine(AppContext.BaseDirectory, "PeerConnectorTests"));
            var peerAddressManager = new PeerAddressManager(DateTimeProvider.Default, peerFolder, this.extendedLoggerFactory);

            var nodeSettings = new NodeSettings();
            nodeSettings.LoadArguments(new string[] { });
            var connectionSettings = new ConnectionManagerSettings();
            connectionSettings.Load(nodeSettings);

            var connector = this.CreatePeerConnecterAddNode(nodeSettings, connectionSettings, peerAddressManager);
            Assert.True(connector.CanStartConnect);
        }

        [Fact]
        public void PeerConnectorConnect_FindPeerToConnectTo_Returns_ConnectNodePeers()
        {
            var peerFolder = AssureEmptyDirAsDataFolder(Path.Combine(AppContext.BaseDirectory, "PeerConnectorTests"));
            var peerAddressManager = new PeerAddressManager(DateTimeProvider.Default, peerFolder, this.extendedLoggerFactory);

            var ipAddressOne = IPAddress.Parse("::ffff:192.168.0.1");
            var endpointAddNode = new IPEndPoint(ipAddressOne, 80);

            var ipAddressTwo = IPAddress.Parse("::ffff:192.168.0.2");
            var endpointDiscoveredNode = new IPEndPoint(ipAddressTwo, 80);

            var ipAddressThree = IPAddress.Parse("::ffff:192.168.0.3");
            var endpointConnectNode = new IPEndPoint(ipAddressThree, 80);

            peerAddressManager.AddPeer(endpointAddNode, IPAddress.Loopback);
            peerAddressManager.AddPeer(endpointConnectNode, IPAddress.Loopback);
            peerAddressManager.AddPeer(endpointDiscoveredNode, IPAddress.Loopback);

            var nodeSettings = new NodeSettings();
            nodeSettings.LoadArguments(new string[] { });

            var connectionSettings = new ConnectionManagerSettings();
            connectionSettings.Load(nodeSettings);

            connectionSettings.Connect.Add(endpointConnectNode);
            var connector = this.CreatePeerConnectorConnectNode(nodeSettings, connectionSettings, peerAddressManager);

            // TODO: Once we have an interface on NetworkPeer we can test this properly.
            //var connector = this.CreatePeerConnecterAddNode(nodeSettings, peerAddressManager);
            //connector.OnConnectAsync().GetAwaiter().GetResult();
            //Assert.Contains(networkAddressConnectNode, connector.ConnectedPeers.Select(p => p.PeerAddress));
        }

        [Fact]
        public void PeerConnectorConnect_WithConnectPeersSpecified_CanStart()
        {
            var peerFolder = AssureEmptyDirAsDataFolder(Path.Combine(AppContext.BaseDirectory, "PeerConnectorTests"));
            var peerAddressManager = new PeerAddressManager(DateTimeProvider.Default, peerFolder, this.extendedLoggerFactory);

            var ipAddressThree = IPAddress.Parse("::ffff:192.168.0.3");
            var endpointConnectNode = new NetworkAddress(ipAddressThree, 80);

            var nodeSettings = new NodeSettings();
            nodeSettings.LoadArguments(new string[] { });

            var connectionSettings = new ConnectionManagerSettings();
            connectionSettings.Load(nodeSettings);

            connectionSettings.Connect.Add(endpointConnectNode.Endpoint);

            var connector = this.CreatePeerConnectorConnectNode(nodeSettings, connectionSettings, peerAddressManager);
            Assert.True(connector.CanStartConnect);
        }

        [Fact]
        public void PeerConnectorConnect_WithNoConnectPeersSpecified_CanNotStart()
        {
            var peerFolder = AssureEmptyDirAsDataFolder(Path.Combine(AppContext.BaseDirectory, "PeerConnectorTests"));
            var peerAddressManager = new PeerAddressManager(DateTimeProvider.Default, peerFolder, this.extendedLoggerFactory);
            var nodeSettings = new NodeSettings();
            nodeSettings.LoadArguments(new string[] { });

            var connectionSettings = new ConnectionManagerSettings();
            connectionSettings.Load(nodeSettings);

            var connector = this.CreatePeerConnectorConnectNode(nodeSettings, connectionSettings, peerAddressManager);
            Assert.False(connector.CanStartConnect);
        }

        [Fact]
        public void PeerConnectorDiscovery_FindPeerToConnectTo_Returns_DiscoveredPeers()
        {
            var peerFolder = AssureEmptyDirAsDataFolder(Path.Combine(AppContext.BaseDirectory, "PeerConnectorTests"));
            var peerAddressManager = new PeerAddressManager(DateTimeProvider.Default, peerFolder, this.extendedLoggerFactory);

            var ipAddressOne = IPAddress.Parse("::ffff:192.168.0.1");
            var endpointAddNode = new IPEndPoint(ipAddressOne, 80);

            var ipAddressThree = IPAddress.Parse("::ffff:192.168.0.3");
            var endPointConnectNode = new IPEndPoint(ipAddressThree, 80);

            var ipAddressTwo = IPAddress.Parse("::ffff:192.168.0.2");
            var endpointDiscoveredNode = new IPEndPoint(ipAddressTwo, 80);

            peerAddressManager.AddPeer(endpointAddNode, IPAddress.Loopback);
            peerAddressManager.AddPeer(endPointConnectNode, IPAddress.Loopback);
            peerAddressManager.AddPeer(endpointDiscoveredNode, IPAddress.Loopback);

            var nodeSettings = new NodeSettings();
            nodeSettings.LoadArguments(new string[] { });

            var connectionSettings = new ConnectionManagerSettings();
            connectionSettings.Load(nodeSettings);

            connectionSettings.AddNode.Add(endpointAddNode);
            connectionSettings.Connect.Add(endPointConnectNode);
            var connector = this.CreatePeerConnectorDiscovery(nodeSettings, connectionSettings, peerAddressManager);

            // TODO: Once we have an interface on NetworkPeer we can test this properly.
            //var connector = this.CreatePeerConnecterAddNode(nodeSettings, peerAddressManager);
            //connector.OnConnectAsync().GetAwaiter().GetResult();
            //Assert.Contains(networkAddressDiscoverNode, connector.ConnectedPeers.Select(p => p.PeerAddress));
        }

        [Fact]
        public void PeerConnectorDiscover_WithNoConnectPeersSpecified_CanStart()
        {
            var peerFolder = AssureEmptyDirAsDataFolder(Path.Combine(AppContext.BaseDirectory, "PeerConnectorTests"));
            var peerAddressManager = new PeerAddressManager(DateTimeProvider.Default, peerFolder, this.extendedLoggerFactory);

            var nodeSettings = new NodeSettings();
            nodeSettings.LoadArguments(new string[] { });

            var connectionSettings = new ConnectionManagerSettings();
            connectionSettings.Load(nodeSettings);

            var connector = this.CreatePeerConnectorDiscovery(nodeSettings, connectionSettings, peerAddressManager);
            Assert.True(connector.CanStartConnect);
        }

        [Fact]
        public void PeerConnectorDiscover_WithConnectPeersSpecified_CanNotStart()
        {
            var nodeSettings = new NodeSettings();
            nodeSettings.LoadArguments(new string[] { });

            var connectionSettings = new ConnectionManagerSettings();
            connectionSettings.Load(nodeSettings);

            var ipAddressThree = IPAddress.Parse("::ffff:192.168.0.3");
            var networkAddressConnectNode = new NetworkAddress(ipAddressThree, 80);

            connectionSettings.Connect.Add(networkAddressConnectNode.Endpoint);

            var peerFolder = AssureEmptyDirAsDataFolder(Path.Combine(AppContext.BaseDirectory, "PeerConnectorTests"));
            var peerAddressManager = new PeerAddressManager(DateTimeProvider.Default, peerFolder, this.extendedLoggerFactory);

            var connector = this.CreatePeerConnectorDiscovery(nodeSettings, connectionSettings, peerAddressManager);
            Assert.False(connector.CanStartConnect);
        }

        private PeerConnectorAddNode CreatePeerConnecterAddNode(NodeSettings nodeSettings, ConnectionManagerSettings connectionSettings, IPeerAddressManager peerAddressManager)
        {
            var peerConnector = new PeerConnectorAddNode(this.asyncLoopFactory, DateTimeProvider.Default, this.extendedLoggerFactory, this.network, this.networkPeerFactory, this.nodeLifetime, nodeSettings, connectionSettings, peerAddressManager);
            var connectionManager = CreateConnectionManager(nodeSettings, connectionSettings, peerAddressManager, peerConnector);
            peerConnector.Initialize(connectionManager);
            return peerConnector;
        }

        private PeerConnectorConnectNode CreatePeerConnectorConnectNode(NodeSettings nodeSettings, ConnectionManagerSettings connectionSettings, IPeerAddressManager peerAddressManager)
        {
            var peerConnector = new PeerConnectorConnectNode(this.asyncLoopFactory, DateTimeProvider.Default, this.extendedLoggerFactory, this.network, this.networkPeerFactory, this.nodeLifetime, nodeSettings, connectionSettings, peerAddressManager);
            var connectionManager = CreateConnectionManager(nodeSettings, connectionSettings, peerAddressManager, peerConnector);
            peerConnector.Initialize(connectionManager);
            return peerConnector;
        }

        private PeerConnectorDiscovery CreatePeerConnectorDiscovery(NodeSettings nodeSettings, ConnectionManagerSettings connectionSettings, IPeerAddressManager peerAddressManager)
        {
            var peerConnector = new PeerConnectorDiscovery(this.asyncLoopFactory, DateTimeProvider.Default, this.extendedLoggerFactory, this.network, this.networkPeerFactory, this.nodeLifetime, nodeSettings, connectionSettings, peerAddressManager);
            var connectionManager = CreateConnectionManager(nodeSettings, connectionSettings, peerAddressManager, peerConnector);
            peerConnector.Initialize(connectionManager);
            return peerConnector;
        }

        private IConnectionManager CreateConnectionManager(NodeSettings nodeSettings, ConnectionManagerSettings connectionSettings, IPeerAddressManager peerAddressManager, IPeerConnector peerConnector)
        {
            var connectionManager = new ConnectionManager(
                DateTimeProvider.Default,
                this.loggerFactory,
                this.network,
                this.networkPeerFactory,
                nodeSettings,
                this.nodeLifetime,
                this.networkPeerParameters,
                peerAddressManager,
                new IPeerConnector[] { peerConnector },
                null,
                connectionSettings);

            return connectionManager;
        }
    }
}<file_sep>/src/Stratis.Bitcoin.IntegrationTests/P2P/PeerConnectionTests.cs
﻿using Microsoft.Extensions.Logging;
using NBitcoin;
using Stratis.Bitcoin.Configuration;
using Stratis.Bitcoin.Configuration.Logging;
using Stratis.Bitcoin.Configuration.Settings;
using Stratis.Bitcoin.Connection;
using Stratis.Bitcoin.IntegrationTests.EnvironmentMockUpHelpers;
using Stratis.Bitcoin.P2P;
using Stratis.Bitcoin.P2P.Peer;
using Stratis.Bitcoin.Utilities;
using Xunit;

namespace Stratis.Bitcoin.IntegrationTests.P2P
{
    public sealed class PeerConnectionTests
    {
        private ILoggerFactory loggerFactory;
        private NetworkPeerConnectionParameters parameters;
        private NetworkPeerFactory networkPeerFactory;
        private INodeLifetime nodeLifetime;
        private NodeSettings nodeSettings;
        private ConnectionManagerSettings connectionSetting;
        private IPeerAddressManager peerAddressManager;
        private readonly Network network;

        public PeerConnectionTests()
        {
            this.network = Network.Main;
        }

        /// <summary>
        /// This tests the fact that we can't find and connect to a
        /// peer that is already connected.
        /// </summary>
        [Fact]
        public void PeerConnectorAddNode_PeerAlreadyConnected_Scenario1()
        {
            this.CreateTestContext("PeerConnectorAddNode_PeerAlreadyConnected_Scenario1");

            var peerConnectorAddNode = new PeerConnectorAddNode(new AsyncLoopFactory(this.loggerFactory), DateTimeProvider.Default, this.loggerFactory, this.network, this.networkPeerFactory, this.nodeLifetime, this.nodeSettings, this.connectionSetting, this.peerAddressManager);
            var peerDiscovery = new PeerDiscovery(new AsyncLoopFactory(this.loggerFactory), this.loggerFactory, this.network, this.networkPeerFactory, this.nodeLifetime, this.nodeSettings, this.peerAddressManager);
            var connectionSettings = new ConnectionManagerSettings();
            connectionSettings.Load(this.nodeSettings);

            IConnectionManager connectionManager = new ConnectionManager(
                DateTimeProvider.Default,
                this.loggerFactory,
                this.network,
                this.networkPeerFactory,
                this.nodeSettings,
                this.nodeLifetime,
                this.parameters,
                this.peerAddressManager,
                new IPeerConnector[] { peerConnectorAddNode },
                peerDiscovery,
                connectionSettings);

            // Create a peer to add to the already connected peers collection.
            using (NodeBuilder builder = NodeBuilder.Create())
            {
                CoreNode coreNode = builder.CreateStratisPowNode();
                builder.StartAll();

                using (NetworkPeer networkPeer = coreNode.CreateNetworkPeerClient())
                {
                    // Add the network peers to the connection manager's
                    // add node collection.
                    connectionManager.AddNodeAddress(networkPeer.PeerEndPoint);

                    // Add the peer to the already connected
                    // peer collection of connection manager.
                    //
                    // This is to simulate that a peer has successfully connected
                    // and that the add node connector's Find method then won't
                    // return the added node.
                    connectionManager.AddConnectedPeer(networkPeer);

                    // Re-initialize the add node peer connector so that it
                    // adds the successful address to the address manager.
                    peerConnectorAddNode.Initialize(connectionManager);

                    // TODO: Once we have an interface on NetworkPeer we can test this properly.
                    // The already connected peer should not be returned.
                    //var peer = peerConnectorAddNode.FindPeerToConnectTo();
                    //Assert.Null(peer);
                }
            }
        }

        private void CreateTestContext(string folder)
        {
            this.loggerFactory = new ExtendedLoggerFactory();
            this.loggerFactory.AddConsoleWithFilters();

            this.parameters = new NetworkPeerConnectionParameters();

            var testFolder = TestDirectory.Create(folder);

            this.nodeSettings = new NodeSettings();
            this.nodeSettings.LoadArguments(new string[] { });
            this.nodeSettings.DataFolder = new DataFolder(this.nodeSettings.DataDir);

            this.connectionSetting = new ConnectionManagerSettings();
            this.connectionSetting.Load(this.nodeSettings);

            this.peerAddressManager = new PeerAddressManager(DateTimeProvider.Default, this.nodeSettings.DataFolder, this.loggerFactory);
            var peerAddressManagerBehaviour = new PeerAddressManagerBehaviour(DateTimeProvider.Default, this.peerAddressManager)
            {
                PeersToDiscover = 10
            };

            this.parameters.TemplateBehaviors.Add(peerAddressManagerBehaviour);

            this.networkPeerFactory = new NetworkPeerFactory(this.network, DateTimeProvider.Default, this.loggerFactory);
            this.nodeLifetime = new NodeLifetime();
        }
    }
}<file_sep>/src/Stratis.Bitcoin/P2P/Peer/NetworkPeerCollection.cs
﻿using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Threading;
using NBitcoin;
using Stratis.Bitcoin.P2P.Protocol;
using Stratis.Bitcoin.Utilities;

namespace Stratis.Bitcoin.P2P.Peer
{
    public class NetworkPeerEventArgs : EventArgs
    {
        public bool Added { get; private set; }
        public NetworkPeer peer { get; private set; }

        public NetworkPeerEventArgs(NetworkPeer peer, bool added)
        {
            this.Added = added;
            this.peer = peer;
        }
    }

    public interface IReadOnlyNetworkPeerCollection : IEnumerable<NetworkPeer>
    {
        event EventHandler<NetworkPeerEventArgs> Added;
        event EventHandler<NetworkPeerEventArgs> Removed;

        NetworkPeer FindByEndpoint(IPEndPoint endpoint);
        NetworkPeer FindByIp(IPAddress ip);
        NetworkPeer FindLocal();
    }

    public class NetworkPeerCollection : IEnumerable<NetworkPeer>, IReadOnlyNetworkPeerCollection
    {
        private ConcurrentDictionary<NetworkPeer, NetworkPeer> networkPeers;

        public int Count
        {
            get
            {
                return this.networkPeers.Count;
            }
        }

        public event EventHandler<NetworkPeerEventArgs> Added;
        public event EventHandler<NetworkPeerEventArgs> Removed;

        /// <summary>
        /// Provides a comparer to specify how peers are compared for equality.
        /// </summary>
        public class NetworkPeerComparer : IEqualityComparer<NetworkPeer>
        {
            public bool Equals(NetworkPeer peerA, NetworkPeer peerB)
            {
                if ((peerA == null) || (peerB == null))
                    return (peerA == null) && (peerB == null);

                return (peerA.RemoteSocketAddress.MapToIPv6().ToString() == peerB.RemoteSocketAddress.MapToIPv6().ToString()) && (peerA.RemoteSocketPort == peerB.RemoteSocketPort);
            }

            public int GetHashCode(NetworkPeer peer)
            {
                if (peer == null)
                    return 0;

                return peer.RemoteSocketPort.GetHashCode() ^ peer.RemoteSocketAddress.MapToIPv6().ToString().GetHashCode();
            }
        }

        public NetworkPeerCollection()
        {
            this.networkPeers = new ConcurrentDictionary<NetworkPeer, NetworkPeer>(new NetworkPeerComparer());
        }

        public bool Add(NetworkPeer peer)
        {
            Guard.NotNull(peer, nameof(peer));

            if (this.networkPeers.TryAdd(peer, peer))
            {
                this.OnPeerAdded(peer);
                return true;
            }

            return false;
        }

        public bool Remove(NetworkPeer peer, string reason)
        {
            NetworkPeer old;
            if (this.networkPeers.TryRemove(peer, out old))
            {
                this.OnPeerRemoved(old);
                peer.Dispose(reason);
                return true;
            }

            return false;
        }

        private void OnPeerAdded(NetworkPeer peer)
        {
            this.Added?.Invoke(this, new NetworkPeerEventArgs(peer, true));
        }

        public void OnPeerRemoved(NetworkPeer peer)
        {
            this.Removed?.Invoke(this, new NetworkPeerEventArgs(peer, false));
        }

        public NetworkPeer FindLocal()
        {
            return this.FindByIp(IPAddress.Loopback);
        }

        public NetworkPeer FindByIp(IPAddress ip)
        {
            ip = ip.EnsureIPv6();
            return this.networkPeers.Where(n => Match(ip, null, n.Key)).Select(s => s.Key).FirstOrDefault();
        }

        public NetworkPeer FindByEndpoint(IPEndPoint endpoint)
        {
            IPAddress ip = endpoint.Address.EnsureIPv6();
            int port = endpoint.Port;
            return this.networkPeers.Select(n => n.Key).FirstOrDefault(n => Match(ip, port, n));
        }

        private static bool Match(IPAddress ip, int? port, NetworkPeer peer)
        {
            if (port.HasValue)
            {
                return ((peer.State > NetworkPeerState.Disconnecting) && peer.RemoteSocketAddress.Equals(ip) && (peer.RemoteSocketPort == port.Value))
                    || (peer.PeerVersion.AddressFrom.Address.Equals(ip) && (peer.PeerVersion.AddressFrom.Port == port.Value));
            }
            else
            {
                return ((peer.State > NetworkPeerState.Disconnecting) && peer.RemoteSocketAddress.Equals(ip))
                    || peer.PeerVersion.AddressFrom.Address.Equals(ip);
            }
        }

        public IEnumerator<NetworkPeer> GetEnumerator()
        {
            return this.networkPeers.Select(n => n.Key).AsEnumerable().GetEnumerator();
        }

        System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
        {
            return this.GetEnumerator();
        }

        public void DisconnectAll(string reason, CancellationToken cancellation = default(CancellationToken))
        {
            foreach (KeyValuePair<NetworkPeer, NetworkPeer> peer in this.networkPeers)
                peer.Key.Dispose(reason);
        }

        public void Clear()
        {
            this.networkPeers.Clear();
        }
    }
}