Skip to content

Commit

Permalink
UT for Version/Verack messages
Browse files Browse the repository at this point in the history
  • Loading branch information
lock9 committed Jun 24, 2019
1 parent 90fa9af commit 6e51178
Show file tree
Hide file tree
Showing 4 changed files with 245 additions and 27 deletions.
62 changes: 62 additions & 0 deletions neo.UnitTests/UT_ProtocolHandler.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
using Akka.Actor;
using Akka.TestKit.Xunit2;
using Microsoft.VisualStudio.TestTools.UnitTesting;
using Neo.Network.P2P;
using Neo.Network.P2P.Capabilities;
using Neo.Network.P2P.Payloads;
using System;
using System.Collections.Generic;
using System.Reflection;
using System.Text;

namespace Neo.UnitTests
{
[TestClass]
public class UT_ProtocolHandler : TestKit
{
private static readonly Random TestRandom = new Random(1337); // use fixed seed for guaranteed determinism

private NeoSystem testBlockchain;

[TestCleanup]
public void Cleanup()
{
Shutdown();
}

[TestInitialize]
public void TestSetup()
{
Akka.Actor.ActorSystem system = Sys;
var config = TestKit.DefaultConfig;
var akkaSettings = new Akka.Actor.Settings(system, config);
testBlockchain = TestBlockchain.InitializeMockNeoSystem();
}

[TestMethod]
public void ProtocolHandler_Test_SendVersion_TellParent()
{
var senderProbe = CreateTestProbe();
var parent = CreateTestProbe();
var protocolActor = ActorOfAsTestActorRef<ProtocolHandler>(()=> new ProtocolHandler(testBlockchain, parent));

var payload = new VersionPayload()
{
UserAgent = "".PadLeft(1024, '0'),
Nonce = 1,
Magic = 2,
Timestamp = 5,
Version = 6,
Capabilities = new NodeCapability[]
{
new ServerCapability(NodeCapabilityType.TcpServer, 25)
}
};

senderProbe.Send(protocolActor, Message.Create(MessageCommand.Version, payload));
parent.ExpectMsg<VersionPayload>();
}
}


}
121 changes: 121 additions & 0 deletions neo.UnitTests/UT_RemoteNode.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,121 @@
using Akka.Actor;
using Akka.IO;
using Akka.TestKit.Xunit2;
using FluentAssertions;
using Microsoft.VisualStudio.TestTools.UnitTesting;
using Neo.Network.P2P;
using Neo.Network.P2P.Capabilities;
using Neo.Network.P2P.Payloads;
using System;
using System.Collections.Generic;
using System.Text;

namespace Neo.UnitTests
{
[TestClass]
public class UT_RemoteNode : TestKit
{
private static readonly Random TestRandom = new Random(1337); // use fixed seed for guaranteed determinism

private NeoSystem testBlockchain;

[TestCleanup]
public void Cleanup()
{
Shutdown();
}

[TestInitialize]
public void TestSetup()
{
Akka.Actor.ActorSystem system = Sys;
testBlockchain = TestBlockchain.InitializeMockNeoSystem();
}

[TestMethod]
public void RemoteNode_Test_Abort_DifferentMagic()
{
var testProbe = CreateTestProbe();
var connectionTestProbe = CreateTestProbe();
var protocolActor = ActorOfAsTestActorRef<ProtocolHandler>(() => new ProtocolHandler(testBlockchain));
var protocolFactory = new TestProtocolFactory(protocolActor);
var remoteNodeActor = ActorOfAsTestActorRef<RemoteNode>(() => new RemoteNode(testBlockchain, connectionTestProbe, null, null, protocolFactory));

connectionTestProbe.ExpectMsg<Tcp.Write>();

var payload = new VersionPayload()
{
UserAgent = "".PadLeft(1024, '0'),
Nonce = 1,
Magic = 2,
Timestamp = 5,
Version = 6,
Capabilities = new NodeCapability[]
{
new ServerCapability(NodeCapabilityType.TcpServer, 25)
}
};

testProbe.Send(remoteNodeActor, payload);

connectionTestProbe.ExpectMsg<Tcp.Abort>();
}

[TestMethod]
public void RemoteNode_Test_Accept_IfSameMagic()
{
var testProbe = CreateTestProbe();
var connectionTestProbe = CreateTestProbe();
var protocolActor = ActorOfAsTestActorRef<ProtocolHandler>(() => new ProtocolHandler(testBlockchain));
var protocolFactory = new TestProtocolFactory(protocolActor);
var remoteNodeActor = ActorOfAsTestActorRef<RemoteNode>(() => new RemoteNode(testBlockchain, connectionTestProbe, null, null, protocolFactory));

connectionTestProbe.ExpectMsg<Tcp.Write>();

var payload = new VersionPayload()
{
UserAgent = "Unit Test".PadLeft(1024, '0'),
Nonce = 1,
Magic = ProtocolSettings.Default.Magic,
Timestamp = 5,
Version = 6,
Capabilities = new NodeCapability[]
{
new ServerCapability(NodeCapabilityType.TcpServer, 25)
}
};

testProbe.Send(remoteNodeActor, payload);

var verackMessage = connectionTestProbe.ExpectMsg<Tcp.Write>();

//Verack
verackMessage.Data.Count.Should().Be(3);
}
}

internal class TestProtocolFactory : IActorRefFactory
{
private IActorRef protocolRef;

public TestProtocolFactory(IActorRef protocolRef)
{
this.protocolRef = protocolRef;
}

public IActorRef ActorOf(Props props, string name = null)
{
return protocolRef;
}

public ActorSelection ActorSelection(ActorPath actorPath)
{
throw new NotImplementedException();
}

public ActorSelection ActorSelection(string actorPath)
{
throw new NotImplementedException();
}
}
}
53 changes: 31 additions & 22 deletions neo/Network/P2P/ProtocolHandler.cs
Original file line number Diff line number Diff line change
Expand Up @@ -26,15 +26,24 @@ public class SetFilter { public BloomFilter Filter; }
private VersionPayload version;
private bool verack = false;
private BloomFilter bloom_filter;
private IActorRef parent;

public ProtocolHandler(NeoSystem system)
{
this.system = system;
this.knownHashes = new FIFOSet<UInt256>(Blockchain.Singleton.MemPool.Capacity * 2);
this.sentHashes = new FIFOSet<UInt256>(Blockchain.Singleton.MemPool.Capacity * 2);
}


public ProtocolHandler(NeoSystem system)
{
this.system = system;
this.knownHashes = new FIFOSet<UInt256>(Blockchain.Singleton.MemPool.Capacity * 2);
this.sentHashes = new FIFOSet<UInt256>(Blockchain.Singleton.MemPool.Capacity * 2);
this.parent = Context.Parent;
}

public ProtocolHandler(NeoSystem system, IActorRef parent) : this(system)
{
this.parent = parent;
}

protected override void OnReceive(object message)
protected override void OnReceive(object message)
{
if (!(message is Message msg)) return;
foreach (IP2PPlugin plugin in Plugin.P2PPlugins)
Expand Down Expand Up @@ -133,13 +142,13 @@ private void OnFilterAddMessageReceived(FilterAddPayload payload)
private void OnFilterClearMessageReceived()
{
bloom_filter = null;
Context.Parent.Tell(new SetFilter { Filter = null });
parent.Tell(new SetFilter { Filter = null });
}

private void OnFilterLoadMessageReceived(FilterLoadPayload payload)
{
bloom_filter = new BloomFilter(payload.Filter.Length * 8, payload.K, payload.Tweak, payload.Filter);
Context.Parent.Tell(new SetFilter { Filter = bloom_filter });
parent.Tell(new SetFilter { Filter = bloom_filter });
}

private void OnGetAddrMessageReceived()
Expand All @@ -152,7 +161,7 @@ private void OnGetAddrMessageReceived()
.Take(AddrPayload.MaxCountToSend);
NetworkAddressWithTime[] networkAddresses = peers.Select(p => NetworkAddressWithTime.Create(p.Listener.Address, p.Version.Timestamp, p.Version.Capabilities)).ToArray();
if (networkAddresses.Length == 0) return;
Context.Parent.Tell(Message.Create(MessageCommand.Addr, AddrPayload.Create(networkAddresses)));
parent.Tell(Message.Create(MessageCommand.Addr, AddrPayload.Create(networkAddresses)));
}

private void OnGetBlocksMessageReceived(GetBlocksPayload payload)
Expand All @@ -172,7 +181,7 @@ private void OnGetBlocksMessageReceived(GetBlocksPayload payload)
hashes.Add(hash);
}
if (hashes.Count == 0) return;
Context.Parent.Tell(Message.Create(MessageCommand.Inv, InvPayload.Create(InventoryType.Block, hashes.ToArray())));
parent.Tell(Message.Create(MessageCommand.Inv, InvPayload.Create(InventoryType.Block, hashes.ToArray())));
}

private void OnGetDataMessageReceived(InvPayload payload)
Expand All @@ -187,7 +196,7 @@ private void OnGetDataMessageReceived(InvPayload payload)
if (inventory == null)
inventory = Blockchain.Singleton.GetTransaction(hash);
if (inventory is Transaction)
Context.Parent.Tell(Message.Create(MessageCommand.Transaction, inventory));
parent.Tell(Message.Create(MessageCommand.Transaction, inventory));
break;
case InventoryType.Block:
if (inventory == null)
Expand All @@ -196,18 +205,18 @@ private void OnGetDataMessageReceived(InvPayload payload)
{
if (bloom_filter == null)
{
Context.Parent.Tell(Message.Create(MessageCommand.Block, inventory));
parent.Tell(Message.Create(MessageCommand.Block, inventory));
}
else
{
BitArray flags = new BitArray(block.Transactions.Select(p => bloom_filter.Test(p)).ToArray());
Context.Parent.Tell(Message.Create(MessageCommand.MerkleBlock, MerkleBlockPayload.Create(block, flags)));
parent.Tell(Message.Create(MessageCommand.MerkleBlock, MerkleBlockPayload.Create(block, flags)));
}
}
break;
case InventoryType.Consensus:
if (inventory != null)
Context.Parent.Tell(Message.Create(MessageCommand.Consensus, inventory));
parent.Tell(Message.Create(MessageCommand.Consensus, inventory));
break;
}
}
Expand All @@ -231,7 +240,7 @@ private void OnGetHeadersMessageReceived(GetBlocksPayload payload)
headers.Add(header);
}
if (headers.Count == 0) return;
Context.Parent.Tell(Message.Create(MessageCommand.Headers, HeadersPayload.Create(headers)));
parent.Tell(Message.Create(MessageCommand.Headers, HeadersPayload.Create(headers)));
}

private void OnHeadersMessageReceived(HeadersPayload payload)
Expand Down Expand Up @@ -268,30 +277,30 @@ private void OnInvMessageReceived(InvPayload payload)
private void OnMemPoolMessageReceived()
{
foreach (InvPayload payload in InvPayload.CreateGroup(InventoryType.TX, Blockchain.Singleton.MemPool.GetVerifiedTransactions().Select(p => p.Hash).ToArray()))
Context.Parent.Tell(Message.Create(MessageCommand.Inv, payload));
parent.Tell(Message.Create(MessageCommand.Inv, payload));
}

private void OnPingMessageReceived(PingPayload payload)
{
Context.Parent.Tell(payload);
Context.Parent.Tell(Message.Create(MessageCommand.Pong, PingPayload.Create(Blockchain.Singleton.Height, payload.Nonce)));
parent.Tell(payload);
parent.Tell(Message.Create(MessageCommand.Pong, PingPayload.Create(Blockchain.Singleton.Height, payload.Nonce)));
}

private void OnPongMessageReceived(PingPayload payload)
{
Context.Parent.Tell(payload);
parent.Tell(payload);
}

private void OnVerackMessageReceived()
{
verack = true;
Context.Parent.Tell(MessageCommand.Verack);
parent.Tell(MessageCommand.Verack);
}

private void OnVersionMessageReceived(VersionPayload payload)
{
version = payload;
Context.Parent.Tell(payload);
parent.Tell(payload);
}

public static Props Props(NeoSystem system)
Expand Down
36 changes: 31 additions & 5 deletions neo/Network/P2P/RemoteNode.cs
Original file line number Diff line number Diff line change
Expand Up @@ -32,12 +32,12 @@ internal class Relay { public IInventory Inventory; }
public uint LastBlockIndex { get; private set; } = 0;
public bool IsFullNode { get; private set; } = false;

public RemoteNode(NeoSystem system, object connection, IPEndPoint remote, IPEndPoint local)
public RemoteNode(NeoSystem system, object connection, IPEndPoint remote, IPEndPoint local, IActorRefFactory protocolFactory)
: base(connection, remote, local)
{
this.system = system;
this.protocol = Context.ActorOf(ProtocolHandler.Props(system));
LocalNode.Singleton.RemoteNodes.TryAdd(Self, this);
this.protocol = protocolFactory.ActorOf(ProtocolHandler.Props(system));
LocalNode.Singleton.RemoteNodes.TryAdd(Self, this);

var capabilities = new List<NodeCapability>
{
Expand All @@ -50,7 +50,7 @@ public RemoteNode(NeoSystem system, object connection, IPEndPoint remote, IPEndP
SendMessage(Message.Create(MessageCommand.Version, VersionPayload.Create(LocalNode.Nonce, LocalNode.UserAgent, capabilities.ToArray())));
}

private void CheckMessageQueue()
private void CheckMessageQueue()
{
if (!verack || !ack) return;
Queue<Message> queue = message_queue_high;
Expand Down Expand Up @@ -224,7 +224,7 @@ protected override void PostStop()

internal static Props Props(NeoSystem system, object connection, IPEndPoint remote, IPEndPoint local)
{
return Akka.Actor.Props.Create(() => new RemoteNode(system, connection, remote, local)).WithMailbox("remote-node-mailbox");
return Akka.Actor.Props.Create(() => new RemoteNode(system, connection, remote, local, new ProtocolActorFactory(Context))).WithMailbox("remote-node-mailbox");
}

private void SendMessage(Message message)
Expand Down Expand Up @@ -269,4 +269,30 @@ internal protected override bool IsHighPriority(object message)
}
}
}

internal class ProtocolActorFactory : IActorRefFactory
{
private IUntypedActorContext context;

public ProtocolActorFactory(IUntypedActorContext context)
{
this.context = context;
}

public IActorRef ActorOf(Props props, string name = null)
{
return context.ActorOf(props);
}

public ActorSelection ActorSelection(ActorPath actorPath)
{
throw new System.NotImplementedException();
}

public ActorSelection ActorSelection(string actorPath)
{
throw new System.NotImplementedException();
}
}

}

0 comments on commit 6e51178

Please sign in to comment.