Skip to content

Commit

Permalink
feat: improve handshake test setup
Browse files Browse the repository at this point in the history
  • Loading branch information
@niklaslong
niklaslong committed Nov 24, 2022
1 parent 01dcb0a commit f3c478a
Show file tree
Hide file tree
Showing 2 changed files with 104 additions and 47 deletions.
15 changes: 7 additions & 8 deletions node/tests/common/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@
// You should have received a copy of the GNU General Public License
// along with the snarkOS library. If not, see <https://www.gnu.org/licenses/>.

use pea2pea::{protocols::Handshake, Config, Connection, Node, Pea2Pea};
use snarkos_node_messages::{ChallengeRequest, ChallengeResponse, Data, Message, MessageCodec, NodeType, Status};
use snarkos_node_tcp::{protocols::Handshake, Config, Connection, Tcp, P2P};
use snarkvm::prelude::{Block, FromBytes, Network, Testnet3 as CurrentNetwork};

use futures_util::{sink::SinkExt, TryStreamExt};
Expand All @@ -29,14 +29,13 @@ const ALEO_MAXIMUM_FORK_DEPTH: u32 = 4096;

#[derive(Clone)]
pub struct TestPeer {
// TODO: should be using pea2pea directly (to keep impls separate)
tcp: Tcp,
node: Node,
node_type: NodeType,
}

impl P2P for TestPeer {
fn tcp(&self) -> &Tcp {
&self.tcp
impl Pea2Pea for TestPeer {
fn node(&self) -> &Node {
&self.node
}
}

Expand All @@ -59,7 +58,7 @@ impl TestPeer {

pub async fn new(node_type: NodeType) -> Self {
let peer = Self {
tcp: Tcp::new(Config {
node: Node::new(Config {
listener_ip: Some(IpAddr::V4(Ipv4Addr::LOCALHOST)),
max_connections: 200,
..Default::default()
Expand All @@ -85,7 +84,7 @@ impl TestPeer {
#[async_trait::async_trait]
impl Handshake for TestPeer {
async fn perform_handshake(&self, mut conn: Connection) -> io::Result<Connection> {
let local_ip = self.tcp().listening_addr().expect("listening address should be present");
let local_ip = self.node().listening_addr().expect("listening address should be present");

let stream = self.borrow_stream(&mut conn);
let mut framed = Framed::new(stream, MessageCodec::<CurrentNetwork>::default());
Expand Down
136 changes: 97 additions & 39 deletions node/tests/new_beacon.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,21 @@
// You should have received a copy of the GNU General Public License
// along with the snarkOS library. If not, see <https://www.gnu.org/licenses/>.

#![recursion_limit = "256"]

mod common;
use common::TestPeer;

use snarkos_account::Account;
use snarkos_node::{Beacon, Validator};
use snarkos_node::{Beacon, Client, Prover, Validator};
use snarkos_node_tcp::P2P;
use snarkvm::prelude::{ConsensusMemory, Testnet3 as CurrentNetwork};

use std::str::FromStr;
use pea2pea::Pea2Pea;

use std::{io, net::SocketAddr, str::FromStr};

/* Node constructors */

async fn beacon() -> Beacon<CurrentNetwork, ConsensusMemory<CurrentNetwork>> {
Beacon::new(
Expand Down Expand Up @@ -50,67 +58,117 @@ async fn validator() -> Validator<CurrentNetwork, ConsensusMemory<CurrentNetwork
.expect("couldn't create beacon instance")
}

macro_rules! test_handshake {
($($node_type:tt -> $peer_type:ident),*) => {
mod handshake_initiator_side {
use snarkos_node_tcp::P2P;
// Trait to unify Pea2Pea and P2P traits.
#[async_trait::async_trait]
trait Connect {
fn listening_addr(&self) -> SocketAddr;

$(
#[tokio::test]
async fn $peer_type () {

let node = $crate::$node_type().await;
async fn connect(&self, target: SocketAddr) -> io::Result<()>;
}

// Spin up a test peer.
let peer = $crate::common::TestPeer::$peer_type().await;
// Implement the `Connect` trait for each node type.
macro_rules! impl_connect {
($($node_type:ident),*) => {
$(
#[async_trait::async_trait]
impl Connect for $node_type<CurrentNetwork, ConsensusMemory<CurrentNetwork>> {
fn listening_addr(&self) -> SocketAddr {
self.tcp().listening_addr().expect("node listener should exist")
}

// Verify the handshake works when the node initiates a connection with the peer.
assert!(
node.tcp().connect(peer.tcp().listening_addr().expect("node listener should exist")).await.is_ok()
);
async fn connect(&self, target: SocketAddr) -> io::Result<()>
where
Self: P2P,
{
self.tcp().connect(target).await
}
}
)*
};
}

)*
}
impl_connect!(Beacon, Client, Prover, Validator);

};
// Implement the `Connect` trait for the test peer.
#[async_trait::async_trait]
impl Connect for TestPeer
where
Self: Pea2Pea,
{
fn listening_addr(&self) -> SocketAddr {
self.node().listening_addr().expect("node listener should exist")
}

($($node_type:tt <- $peer_type:ident),*) => {
mod handshake_responder_side {
use snarkos_node_tcp::P2P;
use snarkos_node_router::Outbound;
async fn connect(&self, target: SocketAddr) -> io::Result<()> {
self.node().connect(target).await
}
}

$(
#[tokio::test]
async fn $peer_type () {
/* Test case */

let node = $crate::$node_type().await;
// Asserts a succesful connection was created from initiator to responder.
async fn assert_connect<T, U>(initiator: T, responder: U)
where
T: Connect,
U: Connect,
{
assert!(initiator.connect(responder.listening_addr()).await.is_ok())
}

// Spin up a test peer.
let peer = $crate::common::TestPeer::$peer_type().await;
// Macro to simply construct handshake cases.
// Syntax:
// - (full_node -> test_peer): full node initiates a handshake to the test peer (synthetic node).
// - (full_node <- test_peer): full node receives a handshake initiated by the test peer.
//
// Test naming: full_node::handshake_<initiator or responder>_side::test_peer.
macro_rules! test_handshake {
($node_type:ident, $peer_type:ident, $is_initiator:expr) => {
#[tokio::test]
async fn $peer_type () {

// Verify the handshake works when the peer initiates a connection with the node.
assert!(
peer.tcp().connect(node.router().tcp().listening_addr().expect("node listener should exist")).await.is_ok()
);
}
// Spin up a full node.
let node = $crate::$node_type().await;

// Spin up a test peer (synthetic node).
let peer = $crate::common::TestPeer::$peer_type().await;

if $is_initiator {
$crate::assert_connect(node, peer).await;
} else {
$crate::assert_connect(peer, node).await;
};
}
};

($($node_type:ident -> $peer_type:ident),*) => {
mod handshake_initiator_side {
$(
test_handshake!($node_type, $peer_type, true);
)*
}

};

($($node_type:ident <- $peer_type:ident),*) => {
mod handshake_responder_side {
$(
test_handshake!($node_type, $peer_type, false);
)*
}

};
}

mod beacon {
// Initiator side.
// Initiator side (full node connects to full node).
test_handshake! {
beacon -> beacon,
beacon -> client,
beacon -> validator,
beacon -> prover
}

// Responder side.
// Responder side (synthetic peer connects to full node).
test_handshake! {
beacon <- beacon,
beacon <- client,
Expand All @@ -120,15 +178,15 @@ mod beacon {
}

mod validator {
// Initiator side.
// Initiator side (full node connects to full node).
test_handshake! {
validator -> beacon,
validator -> client,
validator -> validator,
validator -> prover
}

// Responder side.
// Responder side (synthetic peer connects to full node).
test_handshake! {
validator <- beacon,
validator <- client,
Expand Down

0 comments on commit f3c478a

Please sign in to comment.