0g-storage-node/node/router/src/libp2p_event_handler.rs
2024-06-20 17:33:32 +08:00

1015 lines
34 KiB
Rust

use std::{ops::Neg, sync::Arc};
use chunk_pool::ChunkPoolMessage;
use file_location_cache::FileLocationCache;
use network::types::{AnnounceShardConfig, SignedAnnounceShardConfig};
use network::{
rpc::StatusMessage,
types::{
AnnounceChunks, AnnounceFile, FindChunks, FindFile, HasSignature, SignedAnnounceChunks,
SignedAnnounceFile, SignedMessage,
},
Keypair, MessageAcceptance, MessageId, NetworkGlobals, NetworkMessage, PeerId, PeerRequestId,
PublicKey, PubsubMessage, Request, RequestId, Response,
};
use shared_types::{bytes_to_chunks, timestamp_now, TxID};
use storage::config::ShardConfig;
use storage_async::Store;
use sync::{SyncMessage, SyncSender};
use tokio::sync::mpsc::UnboundedSender;
use tokio::sync::{mpsc, RwLock};
use crate::peer_manager::PeerManager;
lazy_static::lazy_static! {
pub static ref FIND_FILE_TIMEOUT: chrono::Duration = chrono::Duration::minutes(2);
pub static ref ANNOUNCE_FILE_TIMEOUT: chrono::Duration = chrono::Duration::minutes(2);
pub static ref ANNOUNCE_SHARD_CONFIG_TIMEOUT: chrono::Duration = chrono::Duration::minutes(2);
pub static ref TOLERABLE_DRIFT: chrono::Duration = chrono::Duration::seconds(5);
}
#[allow(deprecated)]
fn duration_since(timestamp: u32) -> chrono::Duration {
let timestamp = i64::from(timestamp);
let timestamp = chrono::NaiveDateTime::from_timestamp_opt(timestamp, 0).expect("should fit");
let now = chrono::Utc::now().naive_utc();
now.signed_duration_since(timestamp)
}
fn peer_id_to_public_key(peer_id: &PeerId) -> Result<PublicKey, String> {
// A libp2p peer id byte representation should be 2 length bytes + 4 protobuf bytes + compressed pk bytes
// if generated from a PublicKey with Identity multihash.
let pk_bytes = &peer_id.to_bytes()[2..];
PublicKey::from_protobuf_encoding(pk_bytes).map_err(|e| {
format!(
" Cannot parse libp2p public key public key from peer id: {}",
e
)
})
}
fn verify_signature(msg: &dyn HasSignature, peer_id: &PeerId, propagation_source: PeerId) -> bool {
match peer_id_to_public_key(peer_id) {
Ok(pub_key) => msg.verify_signature(&pub_key),
Err(err) => {
error!(
?err,
?peer_id,
?propagation_source,
"Failed to verify signature"
);
false
}
}
}
pub struct Libp2pEventHandler {
/// A collection of global variables, accessible outside of the network service.
network_globals: Arc<NetworkGlobals>,
/// A channel to the router service.
network_send: mpsc::UnboundedSender<NetworkMessage>,
/// A channel to the syncing service.
sync_send: SyncSender,
/// A channel to the RPC chunk pool service.
chunk_pool_send: mpsc::UnboundedSender<ChunkPoolMessage>,
/// Node keypair for signing messages.
local_keypair: Keypair,
/// Log and transaction storage.
store: Store,
/// Cache for storing and serving gossip messages.
file_location_cache: Arc<FileLocationCache>,
/// All connected peers.
peers: Arc<RwLock<PeerManager>>,
}
impl Libp2pEventHandler {
#[allow(clippy::too_many_arguments)]
pub fn new(
network_globals: Arc<NetworkGlobals>,
network_send: mpsc::UnboundedSender<NetworkMessage>,
sync_send: SyncSender,
chunk_pool_send: UnboundedSender<ChunkPoolMessage>,
local_keypair: Keypair,
store: Store,
file_location_cache: Arc<FileLocationCache>,
peers: Arc<RwLock<PeerManager>>,
) -> Self {
Self {
network_globals,
network_send,
sync_send,
chunk_pool_send,
local_keypair,
store,
file_location_cache,
peers,
}
}
fn send_to_network(&self, message: NetworkMessage) {
self.network_send.send(message).unwrap_or_else(|err| {
warn!(%err, "Could not send message to the network service");
});
}
pub fn send_to_sync(&self, message: SyncMessage) {
self.sync_send.notify(message).unwrap_or_else(|err| {
warn!(%err, "Could not send message to the sync service");
});
}
pub fn send_to_chunk_pool(&self, message: ChunkPoolMessage) {
self.chunk_pool_send.send(message).unwrap_or_else(|err| {
warn!(%err, "Could not send message to the chunk pool service");
});
}
pub fn publish(&self, msg: PubsubMessage) {
self.send_to_network(NetworkMessage::Publish {
messages: vec![msg],
});
}
pub fn send_status(&self, peer_id: PeerId) {
let status_message = StatusMessage { data: 123 }; // dummy status message
debug!(%peer_id, ?status_message, "Sending Status request");
self.send_to_network(NetworkMessage::SendRequest {
peer_id,
request_id: RequestId::Router,
request: Request::Status(status_message),
});
}
pub async fn on_peer_connected(&self, peer_id: PeerId, outgoing: bool) {
self.peers.write().await.add(peer_id, outgoing);
if outgoing {
self.send_status(peer_id);
self.send_to_sync(SyncMessage::PeerConnected { peer_id });
}
}
pub async fn on_peer_disconnected(&self, peer_id: PeerId) {
self.peers.write().await.remove(&peer_id);
self.send_to_sync(SyncMessage::PeerDisconnected { peer_id });
}
pub async fn on_rpc_request(
&self,
peer_id: PeerId,
request_id: PeerRequestId,
request: Request,
) {
self.peers.write().await.update(&peer_id);
match request {
Request::Status(status) => {
self.on_status_request(peer_id, request_id, status);
}
Request::GetChunks(request) => {
self.send_to_sync(SyncMessage::RequestChunks {
peer_id,
request_id,
request,
});
}
Request::DataByHash(_) => {
// ignore
}
}
}
fn on_status_request(&self, peer_id: PeerId, request_id: PeerRequestId, status: StatusMessage) {
debug!(%peer_id, ?status, "Received Status request");
let status_message = StatusMessage { data: 456 }; // dummy status message
debug!(%peer_id, ?status_message, "Sending Status response");
self.send_to_network(NetworkMessage::SendResponse {
peer_id,
id: request_id,
response: Response::Status(status_message),
});
}
pub async fn on_rpc_response(
&self,
peer_id: PeerId,
request_id: RequestId,
response: Response,
) {
self.peers.write().await.update(&peer_id);
match response {
Response::Status(status_message) => {
debug!(%peer_id, ?status_message, "Received Status response");
}
Response::Chunks(response) => {
let request_id = match request_id {
RequestId::Sync(sync_id) => sync_id,
_ => unreachable!("All Chunks responses belong to sync"),
};
self.send_to_sync(SyncMessage::ChunksResponse {
peer_id,
request_id,
response,
});
}
Response::DataByHash(_) => {
// ignore
}
}
}
pub async fn on_rpc_error(&self, peer_id: PeerId, request_id: RequestId) {
self.peers.write().await.update(&peer_id);
// Check if the failed RPC belongs to sync
if let RequestId::Sync(request_id) = request_id {
self.send_to_sync(SyncMessage::RpcError {
peer_id,
request_id,
});
}
}
pub async fn on_pubsub_message(
&self,
propagation_source: PeerId,
source: PeerId,
id: &MessageId,
message: PubsubMessage,
) -> MessageAcceptance {
info!(?message, %propagation_source, %source, %id, "Received pubsub message");
match message {
PubsubMessage::ExampleMessage(_) => MessageAcceptance::Ignore,
PubsubMessage::FindFile(msg) => self.on_find_file(msg).await,
PubsubMessage::FindChunks(msg) => self.on_find_chunks(msg).await,
PubsubMessage::AnnounceFile(msg) => self.on_announce_file(propagation_source, msg),
PubsubMessage::AnnounceChunks(msg) => self.on_announce_chunks(propagation_source, msg),
PubsubMessage::AnnounceShardConfig(msg) => {
self.on_announce_shard_config(propagation_source, msg)
}
}
}
pub async fn construct_announce_file_message(&self, tx_id: TxID) -> Option<PubsubMessage> {
let peer_id = *self.network_globals.peer_id.read();
let addr = match self.network_globals.listen_multiaddrs.read().first() {
Some(addr) => addr.clone(),
None => {
error!("No listen address available");
return None;
}
};
let timestamp = timestamp_now();
let shard_config = self.store.get_store().flow().get_shard_config();
let msg = AnnounceFile {
tx_id,
num_shard: shard_config.num_shard,
shard_id: shard_config.shard_id,
peer_id: peer_id.into(),
at: addr.into(),
timestamp,
};
let mut signed = match SignedMessage::sign_message(msg, &self.local_keypair) {
Ok(signed) => signed,
Err(e) => {
error!(%tx_id.seq, %e, "Failed to sign AnnounceFile message");
return None;
}
};
signed.resend_timestamp = timestamp;
Some(PubsubMessage::AnnounceFile(signed))
}
pub async fn construct_announce_shard_config_message(
&self,
shard_config: ShardConfig,
) -> Option<PubsubMessage> {
let peer_id = *self.network_globals.peer_id.read();
let addr = match self.network_globals.listen_multiaddrs.read().first() {
Some(addr) => addr.clone(),
None => {
error!("No listen address available");
return None;
}
};
let timestamp = timestamp_now();
let msg = AnnounceShardConfig {
num_shard: shard_config.num_shard,
shard_id: shard_config.shard_id,
peer_id: peer_id.into(),
at: addr.into(),
timestamp,
};
let mut signed = match SignedMessage::sign_message(msg, &self.local_keypair) {
Ok(signed) => signed,
Err(e) => {
error!(%e, "Failed to sign AnnounceShardConfig message");
return None;
}
};
signed.resend_timestamp = timestamp;
Some(PubsubMessage::AnnounceShardConfig(signed))
}
async fn on_find_file(&self, msg: FindFile) -> MessageAcceptance {
let FindFile { tx_id, timestamp } = msg;
// verify timestamp
let d = duration_since(timestamp);
if d < TOLERABLE_DRIFT.neg() || d > *FIND_FILE_TIMEOUT {
debug!(%timestamp, "Invalid timestamp, ignoring FindFile message");
return MessageAcceptance::Ignore;
}
// check if we have it
if matches!(self.store.check_tx_completed(tx_id.seq).await, Ok(true)) {
if let Ok(Some(tx)) = self.store.get_tx_by_seq_number(tx_id.seq).await {
if tx.id() == tx_id {
debug!(?tx_id, "Found file locally, responding to FindFile query");
return match self.construct_announce_file_message(tx_id).await {
Some(msg) => {
self.publish(msg);
MessageAcceptance::Ignore
}
// propagate FindFile query to other nodes
None => MessageAcceptance::Accept,
};
}
}
}
// try from cache
if let Some(mut msg) = self.file_location_cache.get_one(tx_id) {
debug!(?tx_id, "Found file in cache, responding to FindFile query");
msg.resend_timestamp = timestamp_now();
self.publish(PubsubMessage::AnnounceFile(msg));
return MessageAcceptance::Ignore;
}
// propagate FindFile query to other nodes
MessageAcceptance::Accept
}
pub fn construct_announce_chunks_message(
&self,
tx_id: TxID,
index_start: u64,
index_end: u64,
) -> Option<PubsubMessage> {
let peer_id = *self.network_globals.peer_id.read();
let addr = self
.network_globals
.listen_multiaddrs
.read()
.first()?
.clone();
let timestamp = timestamp_now();
let msg = AnnounceChunks {
tx_id,
index_start,
index_end,
peer_id: peer_id.into(),
at: addr.into(),
timestamp,
};
let mut signed = match SignedMessage::sign_message(msg, &self.local_keypair) {
Ok(signed) => signed,
Err(e) => {
error!(%tx_id.seq, %e, "Failed to sign AnnounceChunks message");
return None;
}
};
signed.resend_timestamp = timestamp;
Some(PubsubMessage::AnnounceChunks(signed))
}
async fn on_find_chunks(&self, msg: FindChunks) -> MessageAcceptance {
// validate message
if msg.index_start >= msg.index_end {
debug!(?msg, "Invalid chunk index range");
return MessageAcceptance::Reject;
}
// verify timestamp
let d = duration_since(msg.timestamp);
if d < TOLERABLE_DRIFT.neg() || d > *FIND_FILE_TIMEOUT {
debug!(%msg.timestamp, "Invalid timestamp, ignoring FindFile message");
return MessageAcceptance::Ignore;
}
// check if we have specified chunks even file not finalized yet
// validate end index
let tx = match self.store.get_tx_by_seq_number(msg.tx_id.seq).await {
Ok(Some(tx)) if tx.id() == msg.tx_id => tx,
_ => return MessageAcceptance::Accept,
};
// validate index range
if let Ok(size) = usize::try_from(tx.size) {
let num_chunks = bytes_to_chunks(size);
if msg.index_end > num_chunks as u64 {
debug!(?msg, "Invalid chunk end index for FindChunks message");
return MessageAcceptance::Reject;
}
}
// TODO(qhz): check if there is better way to check existence of requested chunks.
match self
.store
.get_chunks_by_tx_and_index_range(
msg.tx_id.seq,
msg.index_start as usize,
msg.index_end as usize,
)
.await
{
Ok(Some(_)) => (),
_ => return MessageAcceptance::Accept,
};
debug!(?msg, "Found chunks to respond FindChunks message");
match self.construct_announce_chunks_message(msg.tx_id, msg.index_start, msg.index_end) {
Some(msg) => {
self.publish(msg);
MessageAcceptance::Ignore
}
// propagate FindFile query to other nodes
None => MessageAcceptance::Accept,
}
}
fn on_announce_file(
&self,
propagation_source: PeerId,
msg: SignedAnnounceFile,
) -> MessageAcceptance {
// verify message signature
if !verify_signature(&msg, &msg.peer_id, propagation_source) {
return MessageAcceptance::Reject;
}
// propagate gossip to peers
let d = duration_since(msg.resend_timestamp);
if d < TOLERABLE_DRIFT.neg() || d > *ANNOUNCE_FILE_TIMEOUT {
debug!(%msg.resend_timestamp, "Invalid resend timestamp, ignoring AnnounceFile message");
return MessageAcceptance::Ignore;
}
// notify sync layer
self.send_to_sync(SyncMessage::AnnounceFileGossip {
tx_id: msg.tx_id,
peer_id: msg.peer_id.clone().into(),
addr: msg.at.clone().into(),
});
// insert message to cache
self.file_location_cache.insert(msg);
MessageAcceptance::Accept
}
fn on_announce_shard_config(
&self,
propagation_source: PeerId,
msg: SignedAnnounceShardConfig,
) -> MessageAcceptance {
// verify message signature
if !verify_signature(&msg, &msg.peer_id, propagation_source) {
return MessageAcceptance::Reject;
}
// propagate gossip to peers
let d = duration_since(msg.resend_timestamp);
if d < TOLERABLE_DRIFT.neg() || d > *ANNOUNCE_SHARD_CONFIG_TIMEOUT {
debug!(%msg.resend_timestamp, "Invalid resend timestamp, ignoring AnnounceShardConfig message");
return MessageAcceptance::Ignore;
}
let shard_config = ShardConfig {
shard_id: msg.shard_id,
num_shard: msg.num_shard,
};
// notify sync layer
self.send_to_sync(SyncMessage::AnnounceShardConfig {
shard_config,
peer_id: msg.peer_id.clone().into(),
addr: msg.at.clone().into(),
});
// insert message to cache
self.file_location_cache
.insert_peer_config(msg.peer_id.clone().into(), shard_config);
MessageAcceptance::Accept
}
fn on_announce_chunks(
&self,
propagation_source: PeerId,
msg: SignedAnnounceChunks,
) -> MessageAcceptance {
// verify message signature
if !verify_signature(&msg, &msg.peer_id, propagation_source) {
return MessageAcceptance::Reject;
}
// propagate gossip to peers
let d = duration_since(msg.resend_timestamp);
if d < TOLERABLE_DRIFT.neg() || d > *ANNOUNCE_FILE_TIMEOUT {
debug!(%msg.resend_timestamp, "Invalid resend timestamp, ignoring AnnounceChunks message");
return MessageAcceptance::Ignore;
}
// notify sync layer
self.send_to_sync(SyncMessage::AnnounceChunksGossip { msg: msg.inner });
MessageAcceptance::Accept
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use channel::Message::*;
use file_location_cache::{test_util::AnnounceFileBuilder, FileLocationCache};
use network::{
discovery::{CombinedKey, ConnectionId},
discv5::enr::EnrBuilder,
rpc::{GetChunksRequest, StatusMessage, SubstreamId},
types::FindFile,
CombinedKeyExt, Keypair, MessageAcceptance, MessageId, Multiaddr, NetworkGlobals,
NetworkMessage, PeerId, PubsubMessage, Request, RequestId, Response, SyncId,
};
use shared_types::{timestamp_now, ChunkArray, ChunkArrayWithProof, FlowRangeProof, TxID};
use storage::{
log_store::{log_manager::LogConfig, Store},
LogManager,
};
use sync::{test_util::create_2_store, SyncMessage, SyncReceiver, SyncSender};
use task_executor::test_utils::TestRuntime;
use tokio::sync::{
mpsc::{self, error::TryRecvError},
RwLock,
};
use crate::{peer_manager::PeerManager, Config};
use super::*;
struct Context {
runtime: TestRuntime,
network_globals: Arc<NetworkGlobals>,
keypair: Keypair,
network_send: mpsc::UnboundedSender<NetworkMessage>,
network_recv: mpsc::UnboundedReceiver<NetworkMessage>,
sync_send: SyncSender,
sync_recv: SyncReceiver,
chunk_pool_send: mpsc::UnboundedSender<ChunkPoolMessage>,
chunk_pool_recv: mpsc::UnboundedReceiver<ChunkPoolMessage>,
store: Arc<dyn Store>,
file_location_cache: Arc<FileLocationCache>,
peers: Arc<RwLock<PeerManager>>,
}
impl Default for Context {
fn default() -> Self {
let runtime = TestRuntime::default();
let (network_globals, keypair) = Context::new_network_globals();
let (network_send, network_recv) = mpsc::unbounded_channel();
let (sync_send, sync_recv) = channel::Channel::unbounded();
let (chunk_pool_send, chunk_pool_recv) = mpsc::unbounded_channel();
let store = LogManager::memorydb(LogConfig::default()).unwrap();
Self {
runtime,
network_globals: Arc::new(network_globals),
keypair,
network_send,
network_recv,
sync_send,
sync_recv,
chunk_pool_send,
chunk_pool_recv,
store: Arc::new(store),
file_location_cache: Arc::new(FileLocationCache::default()),
peers: Arc::new(RwLock::new(PeerManager::new(Config::default()))),
}
}
}
impl Context {
fn new_handler(&self) -> Libp2pEventHandler {
Libp2pEventHandler::new(
self.network_globals.clone(),
self.network_send.clone(),
self.sync_send.clone(),
self.chunk_pool_send.clone(),
self.keypair.clone(),
storage_async::Store::new(self.store.clone(), self.runtime.task_executor.clone()),
self.file_location_cache.clone(),
self.peers.clone(),
)
}
fn new_network_globals() -> (NetworkGlobals, Keypair) {
let keypair = Keypair::generate_secp256k1();
let enr_key = CombinedKey::from_libp2p(&keypair).unwrap();
let enr = EnrBuilder::new("v4").build(&enr_key).unwrap();
let network_globals = NetworkGlobals::new(enr, 30000, 30000, vec![]);
let listen_addr: Multiaddr = "/ip4/127.0.0.1/tcp/30000".parse().unwrap();
network_globals.listen_multiaddrs.write().push(listen_addr);
(network_globals, keypair)
}
fn assert_status_request(&mut self, expected_peer_id: PeerId) {
match self.network_recv.try_recv() {
Ok(NetworkMessage::SendRequest {
peer_id,
request,
request_id,
}) => {
assert_eq!(peer_id, expected_peer_id);
assert!(matches!(request, Request::Status(..)));
assert!(matches!(request_id, RequestId::Router))
}
Ok(_) => panic!("Unexpected network message type received"),
Err(e) => panic!("No network message received: {:?}", e),
}
}
fn assert_file_announcement_published(&mut self, expected_tx_id: TxID) {
match self.network_recv.try_recv() {
Ok(NetworkMessage::Publish { messages }) => {
assert_eq!(messages.len(), 1);
assert!(
matches!(&messages[0], PubsubMessage::AnnounceFile(file) if file.tx_id == expected_tx_id)
);
}
Ok(_) => panic!("Unexpected network message type received"),
Err(e) => panic!("No network message received: {:?}", e),
}
}
}
#[test]
fn test_send_status() {
let mut ctx = Context::default();
let handler = ctx.new_handler();
assert!(matches!(
ctx.network_recv.try_recv(),
Err(TryRecvError::Empty)
));
let alice = PeerId::random();
handler.send_status(alice);
ctx.assert_status_request(alice);
}
#[tokio::test]
async fn test_on_peer_connected_incoming() {
let mut ctx = Context::default();
let handler = ctx.new_handler();
assert_eq!(handler.peers.read().await.size(), 0);
let alice = PeerId::random();
handler.on_peer_connected(alice, false).await;
assert_eq!(handler.peers.read().await.size(), 1);
assert!(matches!(
ctx.network_recv.try_recv(),
Err(TryRecvError::Empty)
));
assert!(matches!(ctx.sync_recv.try_recv(), Err(TryRecvError::Empty)));
}
#[tokio::test]
async fn test_on_peer_connected_outgoing() {
let mut ctx = Context::default();
let handler = ctx.new_handler();
assert_eq!(handler.peers.read().await.size(), 0);
let alice = PeerId::random();
handler.on_peer_connected(alice, true).await;
assert_eq!(handler.peers.read().await.size(), 1);
ctx.assert_status_request(alice);
assert!(matches!(
ctx.sync_recv.try_recv(),
Ok(Notification(SyncMessage::PeerConnected {peer_id})) if peer_id == alice
));
}
#[tokio::test]
async fn test_on_peer_disconnected() {
let mut ctx = Context::default();
let handler = ctx.new_handler();
let alice = PeerId::random();
handler.on_peer_connected(alice, false).await;
assert_eq!(handler.peers.read().await.size(), 1);
handler.on_peer_disconnected(alice).await;
assert_eq!(handler.peers.read().await.size(), 0);
assert!(matches!(
ctx.sync_recv.try_recv(),
Ok(Notification(SyncMessage::PeerDisconnected {peer_id})) if peer_id == alice
));
}
#[tokio::test]
async fn test_on_rpc_request_status() {
let mut ctx = Context::default();
let handler = ctx.new_handler();
let alice = PeerId::random();
let req_id = (ConnectionId::new(4), SubstreamId(12));
let request = Request::Status(StatusMessage { data: 412 });
handler.on_rpc_request(alice, req_id, request).await;
match ctx.network_recv.try_recv() {
Ok(NetworkMessage::SendResponse {
peer_id,
response,
id,
}) => {
assert_eq!(peer_id, alice);
assert!(matches!(response, Response::Status(..)));
assert_eq!(id, req_id);
}
Ok(_) => panic!("Unexpected network message type received"),
Err(e) => panic!("No network message received: {:?}", e),
}
}
#[tokio::test]
async fn test_on_rpc_request_get_chunks() {
let mut ctx = Context::default();
let handler = ctx.new_handler();
let alice = PeerId::random();
let id = (ConnectionId::new(4), SubstreamId(12));
let raw_request = GetChunksRequest {
tx_id: TxID::random_hash(7),
index_start: 66,
index_end: 99,
};
handler
.on_rpc_request(alice, id, Request::GetChunks(raw_request.clone()))
.await;
match ctx.sync_recv.try_recv() {
Ok(Notification(SyncMessage::RequestChunks {
peer_id,
request_id,
request,
})) => {
assert_eq!(peer_id, alice);
assert_eq!(request_id, id);
assert_eq!(request, raw_request);
}
Ok(_) => panic!("Unexpected sync message type received"),
Err(e) => panic!("No sync message received: {:?}", e),
}
}
#[tokio::test]
async fn test_on_rpc_response() {
let mut ctx = Context::default();
let handler = ctx.new_handler();
let alice = PeerId::random();
let id = TxID::random_hash(555);
let data = ChunkArrayWithProof {
chunks: ChunkArray {
data: vec![1, 2, 3, 4],
start_index: 16,
},
proof: FlowRangeProof::new_empty(),
};
handler
.on_rpc_response(
alice,
RequestId::Sync(SyncId::SerialSync { tx_id: id }),
Response::Chunks(data.clone()),
)
.await;
match ctx.sync_recv.try_recv() {
Ok(Notification(SyncMessage::ChunksResponse {
peer_id,
request_id,
response,
})) => {
assert_eq!(peer_id, alice);
assert!(matches!(request_id, SyncId::SerialSync { tx_id } if tx_id == id ));
assert_eq!(response, data);
}
Ok(_) => panic!("Unexpected sync message type received"),
Err(e) => panic!("No sync message received: {:?}", e),
}
}
#[tokio::test]
async fn test_on_rpc_error() {
let mut ctx = Context::default();
let handler = ctx.new_handler();
let alice = PeerId::random();
let id = TxID::random_hash(555);
handler
.on_rpc_error(alice, RequestId::Sync(SyncId::SerialSync { tx_id: id }))
.await;
match ctx.sync_recv.try_recv() {
Ok(Notification(SyncMessage::RpcError {
peer_id,
request_id,
})) => {
assert_eq!(peer_id, alice);
assert!(matches!(request_id, SyncId::SerialSync { tx_id } if tx_id == id ));
}
Ok(_) => panic!("Unexpected sync message type received"),
Err(e) => panic!("No sync message received: {:?}", e),
}
}
async fn handle_find_file_msg(
handler: &Libp2pEventHandler,
tx_id: TxID,
timestamp: u32,
) -> MessageAcceptance {
let (alice, bob) = (PeerId::random(), PeerId::random());
let id = MessageId::new(b"dummy message");
let message = PubsubMessage::FindFile(FindFile { tx_id, timestamp });
handler.on_pubsub_message(alice, bob, &id, message).await
}
#[tokio::test]
async fn test_on_pubsub_find_file_invalid_timestamp() {
let ctx = Context::default();
let handler = ctx.new_handler();
// message too future
let result = handle_find_file_msg(
&handler,
TxID::random_hash(412),
timestamp_now() + 10 + TOLERABLE_DRIFT.num_seconds() as u32,
)
.await;
assert!(matches!(result, MessageAcceptance::Ignore));
// message too old
let result = handle_find_file_msg(
&handler,
TxID::random_hash(412),
timestamp_now() - 10 - FIND_FILE_TIMEOUT.num_seconds() as u32,
)
.await;
assert!(matches!(result, MessageAcceptance::Ignore));
}
#[tokio::test]
async fn test_on_pubsub_find_file_not_found() {
let ctx = Context::default();
let handler = ctx.new_handler();
let result = handle_find_file_msg(&handler, TxID::random_hash(412), timestamp_now()).await;
assert!(matches!(result, MessageAcceptance::Accept));
}
#[tokio::test]
async fn test_on_pubsub_find_file_in_store() {
let mut ctx = Context::default();
// prepare store with txs
let (_, store, txs, _) = create_2_store(vec![1314]);
ctx.store = store;
let handler = ctx.new_handler();
// receive find file request
let result = handle_find_file_msg(&handler, txs[0].id(), timestamp_now()).await;
assert!(matches!(result, MessageAcceptance::Ignore));
ctx.assert_file_announcement_published(txs[0].id());
}
#[tokio::test]
async fn test_on_pubsub_find_file_in_cache() {
let mut ctx = Context::default();
let handler = ctx.new_handler();
// prepare tx in cache
let tx_id = TxID::random_hash(412);
let signed = AnnounceFileBuilder::default()
.with_tx_id(tx_id)
.with_timestamp(timestamp_now() - 5)
.build();
ctx.file_location_cache.insert(signed);
// receive find file request
let result = handle_find_file_msg(&handler, tx_id, timestamp_now()).await;
assert!(matches!(result, MessageAcceptance::Ignore));
ctx.assert_file_announcement_published(tx_id);
}
#[tokio::test]
async fn test_on_pubsub_announce_file_invalid_sig() {
let ctx = Context::default();
let handler = ctx.new_handler();
let (alice, bob) = (PeerId::random(), PeerId::random());
let id = MessageId::new(b"dummy message");
let tx_id = TxID::random_hash(412);
// change signed message
let message = match handler
.construct_announce_file_message(tx_id)
.await
.unwrap()
{
PubsubMessage::AnnounceFile(mut file) => {
let malicious_addr: Multiaddr = "/ip4/127.0.0.38/tcp/30000".parse().unwrap();
file.inner.at = malicious_addr.into();
PubsubMessage::AnnounceFile(file)
}
_ => panic!("Unexpected pubsub message type"),
};
// failed to verify signature
let result = handler.on_pubsub_message(alice, bob, &id, message).await;
assert!(matches!(result, MessageAcceptance::Reject));
}
#[tokio::test]
async fn test_on_pubsub_announce_file() {
let mut ctx = Context::default();
let handler = ctx.new_handler();
// prepare message
let (alice, bob) = (PeerId::random(), PeerId::random());
let id = MessageId::new(b"dummy message");
let tx = TxID::random_hash(412);
let message = handler.construct_announce_file_message(tx).await.unwrap();
// succeeded to handle
let result = handler.on_pubsub_message(alice, bob, &id, message).await;
assert!(matches!(result, MessageAcceptance::Accept));
// ensure notify to sync layer
match ctx.sync_recv.try_recv() {
Ok(Notification(SyncMessage::AnnounceFileGossip {
tx_id,
peer_id,
addr,
})) => {
assert_eq!(tx_id, tx);
assert_eq!(peer_id, *ctx.network_globals.peer_id.read());
assert_eq!(
addr,
*ctx.network_globals
.listen_multiaddrs
.read()
.first()
.unwrap()
);
}
Ok(_) => panic!("Unexpected sync message type received"),
Err(e) => panic!("No sync message received: {:?}", e),
}
// ensure cache updated
assert_eq!(ctx.file_location_cache.get_all(tx).len(), 1);
}
}