Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Test] Add unit test for validator Router connections #3198

Merged
merged 4 commits into from
Apr 3, 2024
Merged

[Test] Add unit test for validator Router connections #3198

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
f654dae
Ensure validator accepts connections from validators
vicsn Apr 1, 2024
9c03d9f
Add unit test asserting validator connection
vicsn Apr 1, 2024
2974612
Only allow validator routers to connect as trusted peers
vicsn Apr 2, 2024
f038b45
Merge branch 'mainnet-staging' into ensure_validator_connects_to_vali…
howardwu Apr 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
7 changes: 4 additions & 3 deletions node/router/src/handshake.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ use crate::{
Peer,
Router,
};
use snarkos_node_router_messages::NodeType;
use snarkos_node_tcp::{ConnectionSide, Tcp, P2P};
use snarkvm::{
ledger::narwhal::Data,
Expand Down Expand Up @@ -202,7 +203,7 @@ impl<N: Network> Router<N> {
let peer_ip = peer_ip.unwrap();

// Knowing the peer's listening address, ensure it is allowed to connect.
if let Err(forbidden_message) = self.ensure_peer_is_allowed(peer_ip) {
if let Err(forbidden_message) = self.ensure_peer_is_allowed(peer_ip, peer_request.node_type) {
return Err(error(format!("{forbidden_message}")));
}
// Verify the challenge request. If a disconnect reason was returned, send the disconnect message and abort.
Expand Down Expand Up @@ -251,7 +252,7 @@ impl<N: Network> Router<N> {
}

/// Ensure the peer is allowed to connect.
fn ensure_peer_is_allowed(&self, peer_ip: SocketAddr) -> Result<()> {
fn ensure_peer_is_allowed(&self, peer_ip: SocketAddr, node_type: NodeType) -> Result<()> {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note that we have to trust the node type as it's given to us by the peer, I don't think we previously used it for anything so we're introducing a new assumption.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right. The easiest immediate resolution is to pass in validator routers explicitly via --peers. I added a unit test for it and updated the PR description. 2974612

// Ensure the peer IP is not this node.
if self.is_local_ip(&peer_ip) {
bail!("Dropping connection request from '{peer_ip}' (attempted to self-connect)")
Expand All @@ -265,7 +266,7 @@ impl<N: Network> Router<N> {
bail!("Dropping connection request from '{peer_ip}' (already connected)")
}
// Only allow trusted peers to connect if allow_external_peers is set
if !self.allow_external_peers() && !self.is_trusted(&peer_ip) {
if !node_type.is_validator() && !self.allow_external_peers() && !self.is_trusted(&peer_ip) {
bail!("Dropping connection request from '{peer_ip}' (untrusted)")
}
// Ensure the peer is not restricted.
Expand Down
2 changes: 1 addition & 1 deletion node/router/tests/cleanups.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ async fn test_connection_cleanups() {
let node = match rng.gen_range(0..3) % 3 {
0 => client(0, 1).await,
1 => prover(0, 1).await,
2 => validator(0, 1).await,
2 => validator(0, 1, true).await,
_ => unreachable!(),
};

Expand Down
4 changes: 2 additions & 2 deletions node/router/tests/common/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,14 +104,14 @@ pub async fn prover(listening_port: u16, max_peers: u16) -> TestRouter<CurrentNe

/// Initializes a validator router. Setting the `listening_port = 0` will result in a random port being assigned.
#[allow(dead_code)]
pub async fn validator(listening_port: u16, max_peers: u16) -> TestRouter<CurrentNetwork> {
pub async fn validator(listening_port: u16, max_peers: u16, allow_external_peers: bool) -> TestRouter<CurrentNetwork> {
Router::new(
SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), listening_port),
NodeType::Validator,
sample_account(),
&[],
max_peers,
true, // Router tests require validators to connect to peers.
allow_external_peers,
true,
)
.await
Expand Down
43 changes: 40 additions & 3 deletions node/router/tests/connect.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ use core::time::Duration;
#[tokio::test]
async fn test_connect_without_handshake() {
// Create 2 routers.
let node0 = validator(0, 2).await;
let node0 = validator(0, 2, true).await;
let node1 = client(0, 2).await;
assert_eq!(node0.number_of_connected_peers(), 0);
assert_eq!(node1.number_of_connected_peers(), 0);
Expand Down Expand Up @@ -78,7 +78,7 @@ async fn test_connect_without_handshake() {
#[tokio::test]
async fn test_connect_with_handshake() {
// Create 2 routers.
let node0 = validator(0, 2).await;
let node0 = validator(0, 2, true).await;
let node1 = client(0, 2).await;
assert_eq!(node0.number_of_connected_peers(), 0);
assert_eq!(node1.number_of_connected_peers(), 0);
Expand Down Expand Up @@ -150,11 +150,48 @@ async fn test_connect_with_handshake() {
}
}

#[tokio::test]
async fn test_validator_connection() {
// Create 2 routers.
let node0 = validator(0, 2, false).await;
let node1 = validator(0, 2, false).await;
assert_eq!(node0.number_of_connected_peers(), 0);
assert_eq!(node1.number_of_connected_peers(), 0);

// Enable handshake protocol.
node0.enable_handshake().await;
node1.enable_handshake().await;

// Start listening.
node0.tcp().enable_listener().await.unwrap();
node1.tcp().enable_listener().await.unwrap();

{
// Connect node0 to node1.
node0.connect(node1.local_ip());
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

aren't we racing against the Heartbeat (that might cause node1 to attempt to connect to node0) here? then again, that direction should fail, so it's probably fine

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

heartbeat is not turned on, only if we call initialize_routing.

// Sleep briefly.
tokio::time::sleep(Duration::from_millis(200)).await;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

deadline could be used here to avoid the sleep.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will keep it in mind for future tests, for now I'll keep the logic similar to the rest of the tests in the file.

Copy link
Collaborator

@ljedrz ljedrz Apr 2, 2024 •
edited

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IIRC connect returns a JoinHandle that can be .awaited (and once it's complete, it means the connection is either ready or that it has failed), so there is no need for a sleep here

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just tried it. Just awaiting node0's connect is not enough, because node1 might not be connected yet. I suggest we leave refactoring this file for later.


print_tcp!(node0);
print_tcp!(node1);

// Check the TCP level.
assert_eq!(node0.tcp().num_connected(), 1);
assert_eq!(node0.tcp().num_connecting(), 0);
assert_eq!(node1.tcp().num_connected(), 1);
assert_eq!(node1.tcp().num_connecting(), 0);

// Check the router level.
assert_eq!(node0.number_of_connected_peers(), 1);
assert_eq!(node1.number_of_connected_peers(), 1);
}
}

#[ignore]
#[tokio::test]
async fn test_connect_simultaneously_with_handshake() {
// Create 2 routers.
let node0 = validator(0, 2).await;
let node0 = validator(0, 2, true).await;
let node1 = client(0, 2).await;
assert_eq!(node0.number_of_connected_peers(), 0);
assert_eq!(node1.number_of_connected_peers(), 0);
Expand Down
4 changes: 2 additions & 2 deletions node/router/tests/disconnect.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ use core::time::Duration;
#[tokio::test]
async fn test_disconnect_without_handshake() {
// Create 2 routers.
let node0 = validator(0, 1).await;
let node0 = validator(0, 1, true).await;
let node1 = client(0, 1).await;
assert_eq!(node0.number_of_connected_peers(), 0);
assert_eq!(node1.number_of_connected_peers(), 0);
Expand Down Expand Up @@ -64,7 +64,7 @@ async fn test_disconnect_without_handshake() {
#[tokio::test]
async fn test_disconnect_with_handshake() {
// Create 2 routers.
let node0 = validator(0, 1).await;
let node0 = validator(0, 1, true).await;
let node1 = client(0, 1).await;
assert_eq!(node0.number_of_connected_peers(), 0);
assert_eq!(node1.number_of_connected_peers(), 0);
Expand Down