0g-storage-node/node/network/tests/rpc_tests.rs
peilun-conflux 12d0c6b675
Some checks failed
abi-consistent-check / build-and-compare (push) Has been cancelled
code-coverage / unittest-cov (push) Has been cancelled
rust / check (push) Has been cancelled
rust / test (push) Has been cancelled
rust / lints (push) Has been cancelled
functional-test / test (push) Has been cancelled
Check network id in status and ban incompatible peers. (#159)
* Check network id in status and ban incompatible peers.

* Revert debug changes.

* Request chain id from the blockchain server.
2024-08-14 11:35:48 +08:00

398 lines
14 KiB
Rust

#![cfg(test)]
use network::rpc::methods::*;
use network::{BehaviourEvent, Libp2pEvent, ReportSource, Request, Response};
use ssz_types::VariableList;
use std::sync::Arc;
use std::time::Duration;
use tokio::runtime::Runtime;
use tokio::time::sleep;
use tracing::{debug, warn};
use tracing_test::traced_test;
mod common;
// Tests the STATUS RPC message
#[test]
#[traced_test]
#[allow(clippy::single_match)]
fn test_status_rpc() {
let rt = Arc::new(Runtime::new().unwrap());
rt.block_on(async {
// get sender/receiver
let (mut sender, mut receiver) = common::build_node_pair(Arc::downgrade(&rt)).await;
// Dummy STATUS RPC message
let rpc_request = Request::Status(StatusMessage {
data: Default::default(),
});
// Dummy STATUS RPC message
let rpc_response = Response::Status(StatusMessage {
data: Default::default(),
});
// build the sender future
let sender_future = async {
loop {
match sender.next_event().await {
Libp2pEvent::Behaviour(BehaviourEvent::PeerConnectedOutgoing(peer_id)) => {
// Send a STATUS message
debug!("Sending RPC");
sender
.swarm
.behaviour_mut()
.send_request(peer_id, 10, rpc_request.clone());
}
Libp2pEvent::Behaviour(BehaviourEvent::ResponseReceived {
peer_id: _,
id: 10,
response,
}) => {
// Should receive the RPC response
debug!("Sender Received");
assert_eq!(response, rpc_response.clone());
debug!("Sender Completed");
return;
}
_ => {}
}
}
};
// build the receiver future
let receiver_future = async {
loop {
match receiver.next_event().await {
Libp2pEvent::Behaviour(BehaviourEvent::RequestReceived {
peer_id,
id,
request,
}) => {
if request == rpc_request {
// send the response
debug!("Receiver Received");
receiver.swarm.behaviour_mut().send_successful_response(
peer_id,
id,
rpc_response.clone(),
);
}
}
_ => {} // Ignore other events
}
}
};
tokio::select! {
_ = sender_future => {}
_ = receiver_future => {}
_ = sleep(Duration::from_secs(30)) => {
panic!("Future timed out");
}
}
})
}
// Tests a streamed DataByHash RPC Message
#[test]
#[traced_test]
#[allow(clippy::single_match)]
fn test_data_by_hash_chunked_rpc() {
let messages_to_send = 6;
let rt = Arc::new(Runtime::new().unwrap());
rt.block_on(async {
// get sender/receiver
let (mut sender, mut receiver) = common::build_node_pair(Arc::downgrade(&rt)).await;
// DataByHash Request
let rpc_request = Request::DataByHash(DataByHashRequest {
hashes: VariableList::from(vec![
Hash256::from_low_u64_be(0),
Hash256::from_low_u64_be(0),
Hash256::from_low_u64_be(0),
Hash256::from_low_u64_be(0),
Hash256::from_low_u64_be(0),
Hash256::from_low_u64_be(0),
]),
});
// DataByHash Response
let data = ZgsData {
hash: Hash256::from_low_u64_be(0),
};
let rpc_response = Response::DataByHash(Some(Box::new(data)));
// keep count of the number of messages received
let mut messages_received = 0;
let request_id = messages_to_send as usize;
// build the sender future
let sender_future = async {
loop {
match sender.next_event().await {
Libp2pEvent::Behaviour(BehaviourEvent::PeerConnectedOutgoing(peer_id)) => {
// Send a DATA_BY_HASH message
debug!("Sending RPC");
sender.swarm.behaviour_mut().send_request(
peer_id,
request_id,
rpc_request.clone(),
);
}
Libp2pEvent::Behaviour(BehaviourEvent::ResponseReceived {
peer_id: _,
id: _,
response,
}) => {
warn!("Sender received a response: {:?}", &response);
match response {
Response::DataByHash(Some(_)) => {
assert_eq!(response, rpc_response.clone());
messages_received += 1;
warn!("Chunk received");
}
Response::DataByHash(None) => {
// should be exactly `messages_to_send` messages before terminating
assert_eq!(messages_received, messages_to_send);
// end the test
return;
}
_ => panic!("Invalid RPC received"),
}
}
_ => {} // Ignore other behaviour events
}
}
};
// build the receiver future
let receiver_future = async {
loop {
match receiver.next_event().await {
Libp2pEvent::Behaviour(BehaviourEvent::RequestReceived {
peer_id,
id,
request,
}) => {
if request == rpc_request {
// send the response
warn!("Receiver got request");
for _ in 0..messages_to_send {
let rpc_response = rpc_response.clone();
receiver.swarm.behaviour_mut().send_successful_response(
peer_id,
id,
rpc_response.clone(),
);
}
// send the stream termination
receiver.swarm.behaviour_mut().send_successful_response(
peer_id,
id,
Response::DataByHash(None),
);
}
}
_ => {} // Ignore other events
}
}
};
tokio::select! {
_ = sender_future => {}
_ = receiver_future => {}
_ = sleep(Duration::from_secs(30)) => {
panic!("Future timed out");
}
}
})
}
// Tests that a streamed DataByHash RPC Message terminates when all expected chunks were received
#[test]
#[traced_test]
fn test_data_by_hash_chunked_rpc_terminates_correctly() {
let messages_to_send = 10;
let extra_messages_to_send = 10;
let rt = Arc::new(Runtime::new().unwrap());
rt.block_on(async {
// get sender/receiver
let (mut sender, mut receiver) = common::build_node_pair(Arc::downgrade(&rt)).await;
// DataByHash Request
let rpc_request = Request::DataByHash(DataByHashRequest {
hashes: VariableList::from(
(0..messages_to_send)
.map(|_| Hash256::from_low_u64_be(0))
.collect::<Vec<_>>(),
),
});
// DataByHash Response
let data = ZgsData {
hash: Hash256::from_low_u64_be(0),
};
let rpc_response = Response::DataByHash(Some(Box::new(data)));
// keep count of the number of messages received
let mut messages_received: u64 = 0;
let request_id = messages_to_send as usize;
// build the sender future
let sender_future = async {
loop {
match sender.next_event().await {
Libp2pEvent::Behaviour(BehaviourEvent::PeerConnectedOutgoing(peer_id)) => {
// Send a STATUS message
debug!("Sending RPC");
sender.swarm.behaviour_mut().send_request(
peer_id,
request_id,
rpc_request.clone(),
);
}
Libp2pEvent::Behaviour(BehaviourEvent::ResponseReceived {
peer_id: _,
id: _,
response,
}) =>
// Should receive the RPC response
{
debug!("Sender received a response");
match response {
Response::DataByHash(Some(_)) => {
assert_eq!(response, rpc_response.clone());
messages_received += 1;
}
Response::DataByHash(None) => {
// should be exactly 10 messages, as requested
assert_eq!(messages_received, messages_to_send);
}
_ => panic!("Invalid RPC received"),
}
}
_ => {} // Ignore other behaviour events
}
}
};
// determine messages to send (PeerId, RequestId). If some, indicates we still need to send
// messages
let mut message_info = None;
// the number of messages we've sent
let mut messages_sent = 0;
let receiver_future = async {
loop {
// this future either drives the sending/receiving or times out allowing messages to be
// sent in the timeout
match futures::future::select(
Box::pin(receiver.next_event()),
Box::pin(tokio::time::sleep(Duration::from_secs(1))),
)
.await
{
futures::future::Either::Left((
Libp2pEvent::Behaviour(BehaviourEvent::RequestReceived {
peer_id,
id,
request,
}),
_,
)) => {
if request == rpc_request {
// send the response
warn!("Receiver got request");
message_info = Some((peer_id, id));
}
}
futures::future::Either::Right((_, _)) => {} // The timeout hit, send messages if required
_ => continue,
}
// if we need to send messages send them here. This will happen after a delay
if message_info.is_some() {
messages_sent += 1;
let (peer_id, stream_id) = message_info.as_ref().unwrap();
receiver.swarm.behaviour_mut().send_successful_response(
*peer_id,
*stream_id,
rpc_response.clone(),
);
debug!("Sending message {}", messages_sent);
if messages_sent == messages_to_send + extra_messages_to_send {
// stop sending messages
return;
}
}
}
};
tokio::select! {
_ = sender_future => {}
_ = receiver_future => {}
_ = sleep(Duration::from_secs(30)) => {
panic!("Future timed out");
}
}
})
}
// Tests a Goodbye RPC message
#[test]
#[traced_test]
#[allow(clippy::single_match)]
fn test_goodbye_rpc() {
let rt = Arc::new(Runtime::new().unwrap());
rt.block_on(async {
let (mut sender, mut receiver) = common::build_node_pair(Arc::downgrade(&rt)).await;
// build the sender future
let sender_future = async {
loop {
match sender.next_event().await {
Libp2pEvent::Behaviour(BehaviourEvent::PeerConnectedOutgoing(peer_id)) => {
// Send a goodbye and disconnect
debug!("Sending RPC");
sender.swarm.behaviour_mut().goodbye_peer(
&peer_id,
GoodbyeReason::IrrelevantNetwork,
ReportSource::SyncService,
);
}
Libp2pEvent::Behaviour(BehaviourEvent::PeerDisconnected(_)) => {
return;
}
_ => {} // Ignore other RPC messages
}
}
};
// build the receiver future
let receiver_future = async {
loop {
match receiver.next_event().await {
Libp2pEvent::Behaviour(BehaviourEvent::PeerDisconnected(_)) => {
// Should receive sent RPC request
return;
}
_ => {} // Ignore other events
}
}
};
let total_future = futures::future::join(sender_future, receiver_future);
tokio::select! {
_ = total_future => {}
_ = sleep(Duration::from_secs(30)) => {
panic!("Future timed out");
}
}
})
}