From ebe098ffe8e2d616b16681428e1eceb78dd0bfd0 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Tue, 18 Feb 2025 02:23:28 -0700 Subject: [PATCH 01/52] feat: overhaul networking, removing setup/networking/gossip/round_based --- .cursor/rules/p2p.mdc | 46 ++ Cargo.lock | 2 + Cargo.toml | 6 +- crates/networking/Cargo.toml | 3 + crates/networking/src/behaviours.rs | 150 ++-- .../src/blueprint_protocol/behaviour.rs | 329 ++++++++ .../networking/src/blueprint_protocol/mod.rs | 56 ++ crates/networking/src/discovery/behaviour.rs | 352 ++++++++ crates/networking/src/discovery/config.rs | 173 ++++ crates/networking/src/discovery/mod.rs | 32 + crates/networking/src/discovery/peers.rs | 254 ++++++ crates/networking/src/error.rs | 10 + crates/networking/src/gossip.rs | 356 -------- crates/networking/src/handlers/connections.rs | 4 +- crates/networking/src/handlers/mdns.rs | 32 - crates/networking/src/handlers/mod.rs | 1 - crates/networking/src/handlers/p2p.rs | 24 +- crates/networking/src/lib.rs | 32 +- crates/networking/src/networking.rs | 552 ------------- crates/networking/src/round_based_compat.rs | 253 ------ crates/networking/src/service.rs | 376 +++++++++ crates/networking/src/setup.rs | 345 -------- crates/networking/src/tests/mod.rs | 773 ++++++++++++++++++ crates/networking/src/types.rs | 128 +-- 24 files changed, 2623 insertions(+), 1666 deletions(-) create mode 100644 .cursor/rules/p2p.mdc create mode 100644 crates/networking/src/blueprint_protocol/behaviour.rs create mode 100644 crates/networking/src/blueprint_protocol/mod.rs create mode 100644 crates/networking/src/discovery/behaviour.rs create mode 100644 crates/networking/src/discovery/config.rs create mode 100644 crates/networking/src/discovery/mod.rs create mode 100644 crates/networking/src/discovery/peers.rs delete mode 100644 crates/networking/src/gossip.rs delete mode 100644 crates/networking/src/handlers/mdns.rs delete mode 100644 crates/networking/src/networking.rs delete mode 100644 crates/networking/src/round_based_compat.rs create mode 100644 crates/networking/src/service.rs delete mode 100644 crates/networking/src/setup.rs create mode 100644 crates/networking/src/tests/mod.rs diff --git a/.cursor/rules/p2p.mdc b/.cursor/rules/p2p.mdc new file mode 100644 index 000000000..f1302339c --- /dev/null +++ b/.cursor/rules/p2p.mdc @@ -0,0 +1,46 @@ +--- +description: Peer to peer networking expert & Rust engineer +globs: **/*.rs +--- +# Peer to peer networking expert + +## P2P Development Cursor (Rust + libp2p) + +1. **Initial Analysis** + - Review input requirements and resources + - Identify minimal viable protocol components + - Review code that exists. + - Never build new when it already exists. Improve instead, remove, and optimize. + - Move tests to a `src/tests` directory, create it if it doesn't exist. + +2. **Implementation Flow** + - Start with concise implementations and build from there. + - Suggest improvements with sound reasoning. + - Leverage libp2p's modular design patterns + - Leverage existing code and improve it, optimize it. + - When a code snippet is provided, understand it, and adapt it with existing code. These are meant as resources not as copy/pastes. + +3. **Code Standards** + - Implement proper error handling with custom error types + - Prioritize concise and efficient code. + - Add relevant and detailed documentation + - Always put tests inside a `src/tests` directory + +4. **Efficiency Guidelines** + - Prefer bounded channels for peer message handling + - Implement connection pooling where appropriate + - Leverage existing libp2p protocols before custom ones + +5. **Review & Integration** + - Verify protocol compatibility + - Test network behaviour under various conditions + - Test in `src/tests` directory + - Ensure proper resource cleanup + - Document failure modes and recovery + +6. **Core Principles** + - Start minimal, expand as needed + - Test thoroughly between iterations + - Maintain clear protocol boundaries + - Document network assumptions and requirements + - \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 51ad6242a..30e3ecbc9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7147,6 +7147,7 @@ dependencies = [ name = "gadget-networking" version = "0.1.0" dependencies = [ + "anyhow", "async-trait", "auto_impl", "bincode", @@ -7168,6 +7169,7 @@ dependencies = [ "serde_json", "thiserror 2.0.11", "tokio", + "tokio-stream", "tracing", "tracing-subscriber 0.3.19", ] diff --git a/Cargo.toml b/Cargo.toml index 520139240..2aa2e85ea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -104,9 +104,12 @@ gadget-runner-symbiotic = { version = "0.1.0", path = "./crates/runners/symbioti gadget-config = { version = "0.1.0", path = "./crates/config", default-features = false } gadget-keystore = { version = "0.1.0", path = "./crates/keystore", default-features = false } gadget-logging = { version = "0.1.0", path = "./crates/logging", default-features = false } -gadget-networking = { version = "0.1.0", path = "./crates/networking", default-features = false } gadget-std = { version = "0.1.0", path = "./crates/std", default-features = false } +# P2P +gadget-networking = { version = "0.1.0", path = "./crates/networking", default-features = false } +gadget-networking-behaviours = { version = "0.1.0", path = "./crates/networking/behaviours", default-features = false } + # Utilities gadget-utils = { version = "0.1.0", path = "./crates/utils", default-features = false } gadget-utils-evm = { version = "0.1.0", path = "./crates/utils/evm", default-features = false } @@ -145,6 +148,7 @@ futures-util = { version = "0.3.31", default-features = false } tokio = { version = "1.40", default-features = false } tokio-util = { version = "0.7.12", default-features = false } tokio-cron-scheduler = "0.13.0" +tokio-stream = { version = "0.1.17", default-features = false } # CLI & Configuration cargo-generate = { version = "0.22.1", default-features = false } diff --git a/crates/networking/Cargo.toml b/crates/networking/Cargo.toml index 1f17f1047..26e5e4993 100644 --- a/crates/networking/Cargo.toml +++ b/crates/networking/Cargo.toml @@ -19,6 +19,7 @@ blake3 = { workspace = true } dashmap = { workspace = true } libp2p = { workspace = true } tokio = { workspace = true, features = ["macros"] } +tokio-stream = { workspace = true } futures = { workspace = true } tracing = { workspace = true } bincode = { workspace = true } @@ -30,6 +31,7 @@ hex = { workspace = true } itertools = { workspace = true, features = ["use_alloc"] } parking_lot = { workspace = true } thiserror = { workspace = true } +anyhow = { workspace = true } # Crypto dependencies gadget-crypto = { workspace = true, features = ["k256", "hashing"] } @@ -58,6 +60,7 @@ features = [ "ping", "dns", "autonat", + "upnp", ] [dev-dependencies] diff --git a/crates/networking/src/behaviours.rs b/crates/networking/src/behaviours.rs index 32257bd93..6f3b4aa6c 100644 --- a/crates/networking/src/behaviours.rs +++ b/crates/networking/src/behaviours.rs @@ -1,55 +1,109 @@ -use crate::key_types::{GossipMsgPublicKey, GossipSignedMsgSignature}; -use libp2p::{gossipsub, kad::store::MemoryStore, mdns, request_response, swarm::NetworkBehaviour}; -use serde::{Deserialize, Serialize}; - -#[non_exhaustive] -#[derive(Serialize, Deserialize, Debug)] -// TODO: Needs better name -pub enum GossipOrRequestResponse { - Gossip(GossipMessage), - Request(MyBehaviourRequest), - Response(MyBehaviourResponse), -} +use crate::{ + blueprint_protocol::{BlueprintProtocolBehaviour, BlueprintProtocolEvent}, + discovery::{ + behaviour::{DiscoveryBehaviour, DiscoveryEvent}, + config::DiscoveryConfig, + PeerInfo, PeerManager, + }, +}; +use libp2p::{ + connection_limits::{self, ConnectionLimits}, + identity::Keypair, + kad::QueryId, + ping, + swarm::NetworkBehaviour, + Multiaddr, PeerId, +}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, + time::Duration, +}; -#[derive(Serialize, Deserialize, Debug)] -pub struct GossipMessage { - pub topic: String, - pub raw_payload: Vec, -} +const MAX_ESTABLISHED_PER_PEER: u32 = 4; -#[non_exhaustive] -#[derive(Serialize, Deserialize, Debug)] -pub enum MyBehaviourRequest { - Handshake { - public_key: GossipMsgPublicKey, - signature: GossipSignedMsgSignature, - }, - Message { - topic: String, - raw_payload: Vec, - }, +/// Events that can be emitted by the GadgetBehavior +#[derive(Debug)] +pub enum GadgetEvent { + /// Discovery-related events + Discovery(DiscoveryEvent), + /// Ping events for connection liveness + Ping(ping::Event), + /// Blueprint protocol events + Blueprint(BlueprintProtocolEvent), } -#[non_exhaustive] -#[derive(Serialize, Deserialize, Debug)] -pub enum MyBehaviourResponse { - Handshaked { - public_key: GossipMsgPublicKey, - signature: GossipSignedMsgSignature, - }, - MessageHandled, +#[derive(NetworkBehaviour)] +pub struct GadgetBehaviour { + /// Connection limits to prevent DoS + connection_limits: connection_limits::Behaviour, + /// Discovery mechanisms (Kademlia, mDNS, etc) + discovery: DiscoveryBehaviour, + /// Direct P2P messaging + blueprint_protocol: BlueprintProtocolBehaviour, + /// Connection liveness checks + ping: ping::Behaviour, } -// We create a custom network behaviour that combines Gossipsub and Mdns. -#[derive(NetworkBehaviour)] -pub struct MyBehaviour { - pub gossipsub: gossipsub::Behaviour, - pub mdns: mdns::tokio::Behaviour, - pub p2p: request_response::cbor::Behaviour, - pub identify: libp2p::identify::Behaviour, - pub kadmelia: libp2p::kad::Behaviour, - pub dcutr: libp2p::dcutr::Behaviour, - pub relay: libp2p::relay::Behaviour, - pub ping: libp2p::ping::Behaviour, - pub autonat: libp2p::autonat::Behaviour, +impl GadgetBehaviour { + pub fn new( + network_name: &str, + local_key: &Keypair, + target_peer_count: u64, + peer_manager: Arc, + ) -> Self { + let connection_limits = connection_limits::Behaviour::new( + ConnectionLimits::default() + .with_max_pending_incoming(Some( + target_peer_count as u32 * MAX_ESTABLISHED_PER_PEER, + )) + .with_max_pending_outgoing(Some( + target_peer_count as u32 * MAX_ESTABLISHED_PER_PEER, + )) + .with_max_established_incoming(Some( + target_peer_count as u32 * MAX_ESTABLISHED_PER_PEER, + )) + .with_max_established_outgoing(Some( + target_peer_count as u32 * MAX_ESTABLISHED_PER_PEER, + )) + .with_max_established_per_peer(Some(MAX_ESTABLISHED_PER_PEER)), + ); + + let ping = ping::Behaviour::new(ping::Config::new().with_interval(Duration::from_secs(30))); + + let discovery = DiscoveryConfig::new(local_key.public(), network_name) + .with_mdns(true) + .with_kademlia(true) + .with_target_peer_count(target_peer_count) + .build() + .unwrap(); + + let blueprint_protocol = BlueprintProtocolBehaviour::new(local_key, peer_manager); + + Self { + connection_limits, + discovery, + blueprint_protocol, + ping, + } + } + + /// Bootstrap Kademlia network + pub fn bootstrap(&mut self) -> Result { + self.discovery.bootstrap() + } + + /// Returns a set of peer ids + pub fn peers(&self) -> &HashSet { + self.discovery.get_peers() + } + + /// Returns a map of peer ids and their multi-addresses + pub fn peer_addresses(&self) -> HashMap> { + self.discovery.get_peer_addresses() + } + + pub fn peer_info(&self, peer_id: &PeerId) -> Option<&PeerInfo> { + self.discovery.get_peer_info(peer_id) + } } diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs new file mode 100644 index 000000000..133e939ab --- /dev/null +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -0,0 +1,329 @@ +use crate::{Curve, InstanceMsgPublicKey, InstanceSignedMsgSignature}; +use dashmap::DashMap; +use gadget_crypto::KeyType; +use gadget_logging::{debug, trace, warn}; +use libp2p::{ + core::transport::PortUse, + gossipsub::{self, IdentTopic, MessageAuthenticity, MessageId, Sha256Topic}, + identity::Keypair, + request_response::{self, OutboundRequestId, ResponseChannel}, + swarm::{ + ConnectionDenied, ConnectionId, FromSwarm, NetworkBehaviour, THandler, THandlerInEvent, + THandlerOutEvent, ToSwarm, + }, + Multiaddr, PeerId, StreamProtocol, +}; +use std::{ + sync::Arc, + task::Poll, + time::{Duration, Instant}, +}; + +use crate::discovery::PeerManager; + +use super::{InstanceMessageRequest, InstanceMessageResponse}; + +/// Events emitted by the BlueprintProtocolBehaviour +#[derive(Debug)] +pub enum BlueprintProtocolEvent { + /// Response received from a peer + Response { + peer: PeerId, + request_id: OutboundRequestId, + response: InstanceMessageResponse, + }, + /// Request received from a peer + Request { + peer: PeerId, + request: InstanceMessageRequest, + channel: ResponseChannel, + }, + /// Gossip message received + GossipMessage { + source: PeerId, + message: Vec, + topic: IdentTopic, + }, +} + +/// Behaviour that handles the blueprint protocol request/response and gossip +pub struct BlueprintProtocolBehaviour { + /// Request/response protocol for direct messaging + request_response: + request_response::cbor::Behaviour, + /// Gossipsub for broadcast messaging + gossipsub: gossipsub::Behaviour, + /// Peer manager for tracking peer states + peer_manager: Arc, + /// Peers with outstanding handshake requests + pending_handshake_peers: DashMap, + /// Active response channels + response_channels: DashMap>, +} + +impl BlueprintProtocolBehaviour { + /// Create a new blueprint protocol behaviour + pub fn new(local_key: &Keypair, peer_manager: Arc) -> Self { + let protocols = vec![( + StreamProtocol::new("/gadget/blueprint_protocol/1.0.0"), + request_response::ProtocolSupport::Full, + )]; + + // Initialize gossipsub with message signing + let gossipsub_config = gossipsub::ConfigBuilder::default() + .heartbeat_interval(Duration::from_secs(1)) + .validation_mode(gossipsub::ValidationMode::Strict) + .build() + .expect("Valid gossipsub config"); + + let gossipsub = gossipsub::Behaviour::new( + MessageAuthenticity::Signed(local_key.clone()), + gossipsub_config, + ) + .expect("Valid gossipsub behaviour"); + + let config = libp2p::request_response::Config::default() + .with_request_timeout(Duration::from_secs(30)) + .with_max_concurrent_streams(50); + + Self { + request_response: request_response::cbor::Behaviour::new(protocols, config), + gossipsub, + peer_manager, + pending_handshake_peers: DashMap::new(), + response_channels: DashMap::new(), + } + } + + /// Send a request to a peer + pub fn send_request( + &mut self, + peer: &PeerId, + request: InstanceMessageRequest, + ) -> OutboundRequestId { + debug!(%peer, ?request, "sending request"); + self.request_response.send_request(peer, request) + } + + /// Send a response through a response channel + pub fn send_response( + &mut self, + channel: ResponseChannel, + response: InstanceMessageResponse, + ) -> Result<(), InstanceMessageResponse> { + debug!(?response, "sending response"); + self.request_response.send_response(channel, response) + } + + /// Subscribe to a gossip topic + pub fn subscribe(&mut self, topic: &str) -> Result { + let topic = Sha256Topic::new(topic); + self.gossipsub.subscribe(&topic) + } + + /// Publish a message to a gossip topic + pub fn publish( + &mut self, + topic: &str, + data: impl Into>, + ) -> Result { + let topic = Sha256Topic::new(topic); + self.gossipsub.publish(topic, data) + } + + /// Verify and handle a handshake with a peer + pub fn verify_handshake( + &self, + peer: &PeerId, + public_key: &InstanceMsgPublicKey, + signature: &InstanceSignedMsgSignature, + ) -> Result<(), InstanceMessageResponse> { + let msg = peer.to_bytes(); + let valid = ::verify(public_key, &msg, signature); + if !valid { + warn!("Invalid initial handshake signature from peer: {peer}"); + return Err(InstanceMessageResponse::Error { + code: 400, + message: "Invalid handshake signature".to_string(), + }); + } + + trace!("Received valid handshake from peer: {peer}"); + + Ok(()) + } + + pub fn handle_handshake( + &self, + peer: &PeerId, + public_key: &InstanceMsgPublicKey, + signature: &InstanceSignedMsgSignature, + ) -> Result<(), InstanceMessageResponse> { + self.verify_handshake(peer, public_key, signature)?; + self.peer_manager + .add_peer_id_to_public_key(peer, public_key); + + Ok(()) + } + /// Handle a failed handshake with a peer + pub fn handle_handshake_failure(&self, peer: &PeerId, reason: &str) { + // Update peer info and potentially ban peer + if let Some(mut peer_info) = self.peer_manager.get_peer_info(peer) { + peer_info.failures += 1; + self.peer_manager.update_peer(*peer, peer_info.clone()); + + // Ban peer if too many failures + if peer_info.failures >= 3 { + self.peer_manager + .ban_peer(*peer, reason, Some(Duration::from_secs(300))); + } + } + } +} + +impl NetworkBehaviour for BlueprintProtocolBehaviour { + type ConnectionHandler = as NetworkBehaviour>::ConnectionHandler; + + type ToSwarm = as NetworkBehaviour>::ToSwarm; + + fn handle_established_inbound_connection( + &mut self, + connection_id: ConnectionId, + peer: PeerId, + local_addr: &libp2p::Multiaddr, + remote_addr: &libp2p::Multiaddr, + ) -> Result, ConnectionDenied> { + self.request_response.handle_established_inbound_connection( + connection_id, + peer, + local_addr, + remote_addr, + ) + } + + fn handle_established_outbound_connection( + &mut self, + connection_id: ConnectionId, + peer: PeerId, + addr: &Multiaddr, + role_override: libp2p::core::Endpoint, + port_use: PortUse, + ) -> Result, ConnectionDenied> { + self.request_response + .handle_established_outbound_connection( + connection_id, + peer, + addr, + role_override, + port_use, + ) + } + + fn handle_pending_inbound_connection( + &mut self, + connection_id: ConnectionId, + local_addr: &libp2p::Multiaddr, + remote_addr: &libp2p::Multiaddr, + ) -> Result<(), ConnectionDenied> { + self.request_response.handle_pending_inbound_connection( + connection_id, + local_addr, + remote_addr, + ) + } + + fn handle_pending_outbound_connection( + &mut self, + connection_id: ConnectionId, + maybe_peer: Option, + addresses: &[libp2p::Multiaddr], + effective_role: libp2p::core::Endpoint, + ) -> Result, ConnectionDenied> { + self.request_response.handle_pending_outbound_connection( + connection_id, + maybe_peer, + addresses, + effective_role, + ) + } + + fn on_connection_handler_event( + &mut self, + peer_id: PeerId, + connection_id: ConnectionId, + event: THandlerOutEvent, + ) { + self.request_response + .on_connection_handler_event(peer_id, connection_id, event) + } + + fn on_swarm_event(&mut self, event: FromSwarm<'_>) { + if let FromSwarm::ConnectionEstablished(e) = &event { + if e.other_established == 0 { + self.pending_handshake_peers + .insert(e.peer_id, Instant::now()); + } + } + + self.request_response.on_swarm_event(event) + } + + fn poll( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll>> { + if let Poll::Ready(ev) = self.request_response.poll(cx) { + // Remove a peer from `pending_handshake_peers` when its handshake request is received. + if let ToSwarm::GenerateEvent(request_response::Event::Message { + peer, + message: + request_response::Message::Request { + request: + InstanceMessageRequest::Handshake { + public_key, + signature, + }, + .. + }, + .. + }) = &ev + { + match self.handle_handshake(peer, public_key, signature) { + Ok(_) => { + self.pending_handshake_peers.remove(&peer); + } + Err(e) => { + self.handle_handshake_failure(peer, "Handshake request failed"); + } + } + } + + return Poll::Ready(ev); + } + + // Track peers whose handshake requests have timed out + const INBOUND_HANDSHAKE_WAIT_TIMEOUT: Duration = Duration::from_secs(30); + let now = Instant::now(); + if let Some((&peer_id, _)) = self + .pending_handshake_peers + .clone() + .into_read_only() + .iter() + .find(|(_, &connected_instant)| { + now.duration_since(connected_instant) > INBOUND_HANDSHAKE_WAIT_TIMEOUT + }) + { + self.pending_handshake_peers.remove(&peer_id); + self.handle_handshake_failure(&peer_id, "Handshake request timeout"); + debug!(peer=%peer_id, "Handshake request timed out, marking failure"); + } + + Poll::Pending + } +} diff --git a/crates/networking/src/blueprint_protocol/mod.rs b/crates/networking/src/blueprint_protocol/mod.rs new file mode 100644 index 000000000..4de601a41 --- /dev/null +++ b/crates/networking/src/blueprint_protocol/mod.rs @@ -0,0 +1,56 @@ +mod behaviour; + +pub use behaviour::{BlueprintProtocolBehaviour, BlueprintProtocolEvent}; + +use crate::key_types::{InstanceMsgPublicKey, InstanceSignedMsgSignature}; +use serde::{Deserialize, Serialize}; + +/// A message sent to a specific instance or broadcast to all instances +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum InstanceMessageRequest { + /// Handshake request with authentication + Handshake { + /// Public key for authentication + public_key: InstanceMsgPublicKey, + /// Signature for verification + signature: InstanceSignedMsgSignature, + }, + /// Protocol-specific message with custom payload + Protocol { + /// Protocol identifier (e.g., "consensus/1.0.0", "sync/1.0.0") + protocol: String, + /// Protocol-specific message payload + payload: Vec, + /// Optional metadata for the protocol handler + metadata: Option>, + }, +} + +/// Response to an instance message +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum InstanceMessageResponse { + /// Handshake response with authentication + Handshake { + /// Public key for authentication + public_key: InstanceMsgPublicKey, + /// Signature for verification + signature: InstanceSignedMsgSignature, + }, + /// Success response with optional data + Success { + /// Response data specific to the protocol + data: Option>, + }, + /// Error response with details + Error { + /// Error code + code: u16, + /// Error message + message: String, + }, + /// Protocol-specific response + Protocol { + /// Protocol-specific response data + data: Vec, + }, +} diff --git a/crates/networking/src/discovery/behaviour.rs b/crates/networking/src/discovery/behaviour.rs new file mode 100644 index 000000000..af521d080 --- /dev/null +++ b/crates/networking/src/discovery/behaviour.rs @@ -0,0 +1,352 @@ +use std::{ + cmp, + collections::{HashMap, HashSet, VecDeque}, + task::{Context, Poll}, + time::Duration, +}; + +use gadget_logging::trace; +use libp2p::{ + autonat, + core::Multiaddr, + identify, + identity::PeerId, + kad::{self, store::MemoryStore}, + mdns::{tokio::Behaviour as Mdns, Event as MdnsEvent}, + relay, + swarm::{ + behaviour::toggle::Toggle, derive_prelude::*, dial_opts::DialOpts, NetworkBehaviour, + ToSwarm, + }, + upnp, +}; +use tokio::time::Interval; +use tracing::{debug, info}; + +use super::PeerInfo; + +#[derive(NetworkBehaviour)] +pub struct DerivedDiscoveryBehaviour { + /// Kademlia discovery + pub kademlia: Toggle>, + /// Local network discovery via mDNS + pub mdns: Toggle, + /// Identify protocol for peer information exchange + pub identify: identify::Behaviour, + /// NAT traversal + pub autonat: autonat::Behaviour, + /// UPnP port mapping + pub upnp: Toggle, + /// Circuit relay for NAT traversal + pub relay: Toggle, +} + +/// Event generated by the `DiscoveryBehaviour`. +#[derive(Debug)] +pub enum DiscoveryEvent { + /// Event that notifies that we connected to the node with the given peer + /// id. + PeerConnected(PeerId), + + /// Event that notifies that we disconnected with the node with the given + /// peer id. + PeerDisconnected(PeerId), + + /// Discovery event + Discovery(Box), +} + +pub struct DiscoveryBehaviour { + /// Discovery behaviour + pub discovery: DerivedDiscoveryBehaviour, + /// Stream that fires when we need to perform the next random Kademlia + /// query. + pub next_kad_random_query: Interval, + /// After `next_kad_random_query` triggers, the next one triggers after this + /// duration. + pub duration_to_next_kad: Duration, + /// Events to return in priority when polled. + pub pending_events: VecDeque, + /// Number of nodes we're currently connected to. + pub n_node_connected: u64, + /// Peers + pub peers: HashSet, + /// Peer info + pub peer_info: HashMap, + /// Target peer count + pub target_peer_count: u64, + /// Options to configure dials to known peers. + pub pending_dial_opts: VecDeque, +} + +impl DiscoveryBehaviour { + pub fn bootstrap(&mut self) -> Result { + if let Some(kademlia) = self.discovery.kademlia.as_mut() { + kademlia.bootstrap().map_err(|e| e.to_string()) + } else { + Err("Kademlia is not enabled".to_string()) + } + } + + pub fn get_peers(&self) -> &HashSet { + &self.peers + } + + pub fn get_peer_info(&self, peer_id: &PeerId) -> Option<&PeerInfo> { + self.peer_info.get(peer_id) + } + + pub fn nat_status(&self) -> autonat::NatStatus { + self.discovery.autonat.nat_status() + } + + pub fn get_peer_addresses(&self) -> HashMap> { + self.peer_info + .iter() + .map(|(peer_id, info)| (*peer_id, info.addresses.clone())) + .collect() + } +} + +impl NetworkBehaviour for DiscoveryBehaviour { + type ConnectionHandler = ::ConnectionHandler; + type ToSwarm = DiscoveryEvent; + + fn handle_established_inbound_connection( + &mut self, + connection_id: ConnectionId, + peer: PeerId, + local_addr: &Multiaddr, + remote_addr: &Multiaddr, + ) -> Result, ConnectionDenied> { + self.discovery.handle_established_inbound_connection( + connection_id, + peer, + local_addr, + remote_addr, + ) + } + + fn handle_established_outbound_connection( + &mut self, + connection_id: ConnectionId, + peer: PeerId, + addr: &libp2p::Multiaddr, + role_override: libp2p::core::Endpoint, + port_use: PortUse, + ) -> Result, ConnectionDenied> { + self.peer_info + .entry(peer) + .or_insert_with(|| PeerInfo { + addresses: HashSet::new(), + identify_info: None, + last_seen: std::time::SystemTime::now(), + ping_latency: None, + successes: 0, + failures: 0, + average_response_time: None, + }) + .addresses + .insert(addr.clone()); + self.discovery.handle_established_outbound_connection( + connection_id, + peer, + addr, + role_override, + port_use, + ) + } + + fn on_connection_handler_event( + &mut self, + peer_id: PeerId, + connection: ConnectionId, + event: THandlerOutEvent, + ) { + self.discovery + .on_connection_handler_event(peer_id, connection, event); + } + + fn on_swarm_event(&mut self, event: FromSwarm<'_>) { + match &event { + FromSwarm::ConnectionEstablished(e) => { + if e.other_established == 0 { + self.n_node_connected += 1; + self.peers.insert(e.peer_id); + self.pending_events + .push_back(DiscoveryEvent::PeerConnected(e.peer_id)); + } + } + FromSwarm::ConnectionClosed(e) => { + if e.remaining_established == 0 { + self.n_node_connected -= 1; + self.peers.remove(&e.peer_id); + self.peer_info.remove(&e.peer_id); + self.pending_events + .push_back(DiscoveryEvent::PeerDisconnected(e.peer_id)); + } + } + _ => {} + }; + self.discovery.on_swarm_event(event) + } + + #[allow(clippy::type_complexity)] + fn poll( + &mut self, + cx: &mut Context<'_>, + ) -> Poll>> { + // Immediately process the content of `discovered`. + if let Some(ev) = self.pending_events.pop_front() { + return Poll::Ready(ToSwarm::GenerateEvent(ev)); + } + + // Dial to peers + if let Some(opts) = self.pending_dial_opts.pop_front() { + return Poll::Ready(ToSwarm::Dial { opts }); + } + + // Poll the stream that fires when we need to start a random Kademlia query. + while self.next_kad_random_query.poll_tick(cx).is_ready() { + if self.n_node_connected < self.target_peer_count { + // We still have not hit the discovery max, send random request for peers. + let random_peer_id = PeerId::random(); + debug!( + "Libp2p <= Starting random Kademlia request for {:?}", + random_peer_id + ); + if let Some(kademlia) = self.discovery.kademlia.as_mut() { + kademlia.get_closest_peers(random_peer_id); + } + } + + // Schedule the next random query with exponentially increasing delay, + // capped at 60 seconds. + self.next_kad_random_query = tokio::time::interval(self.duration_to_next_kad); + // we need to reset the interval, otherwise the next tick completes immediately. + self.next_kad_random_query.reset(); + + self.duration_to_next_kad = + cmp::min(self.duration_to_next_kad * 2, Duration::from_secs(60)); + } + + // Poll discovery events. + while let Poll::Ready(ev) = self.discovery.poll(cx) { + match ev { + ToSwarm::GenerateEvent(ev) => { + match &ev { + DerivedDiscoveryBehaviourEvent::Identify(ev) => { + if let identify::Event::Received { peer_id, info, .. } = ev { + self.peer_info.entry(*peer_id).or_default().identify_info = + Some(info.clone()); + if let Some(kademlia) = self.discovery.kademlia.as_mut() { + for address in &info.listen_addrs { + kademlia.add_address(peer_id, address.clone()); + } + } + } + } + DerivedDiscoveryBehaviourEvent::Autonat(_) => {} + DerivedDiscoveryBehaviourEvent::Upnp(ev) => match ev { + upnp::Event::NewExternalAddr(addr) => { + info!("UPnP NewExternalAddr: {addr}"); + } + upnp::Event::ExpiredExternalAddr(addr) => { + info!("UPnP ExpiredExternalAddr: {addr}"); + } + upnp::Event::GatewayNotFound => { + info!("UPnP GatewayNotFound"); + } + upnp::Event::NonRoutableGateway => { + info!("UPnP NonRoutableGateway"); + } + }, + DerivedDiscoveryBehaviourEvent::Kademlia(ev) => match ev { + // Adding to Kademlia buckets is automatic with our config, + // no need to do manually. + kad::Event::RoutingUpdated { .. } => {} + kad::Event::RoutablePeer { .. } => {} + kad::Event::PendingRoutablePeer { .. } => { + // Intentionally ignore + } + other => { + trace!("Libp2p => Unhandled Kademlia event: {:?}", other) + } + }, + DerivedDiscoveryBehaviourEvent::Mdns(ev) => match ev { + MdnsEvent::Discovered(list) => { + if self.n_node_connected >= self.target_peer_count { + // Already over discovery max, don't add discovered peers. + // We could potentially buffer these addresses to be added later, + // but mdns is not an important use case and may be removed in future. + continue; + } + + // Add any discovered peers to Kademlia + for (peer_id, multiaddr) in list { + if let Some(kad) = self.discovery.kademlia.as_mut() { + kad.add_address(peer_id, multiaddr.clone()); + } + } + } + MdnsEvent::Expired(_) => {} + }, + DerivedDiscoveryBehaviourEvent::Relay(relay_event) => match relay_event { + relay::Event::ReservationReqAccepted { src_peer_id, .. } => { + debug!("Relay accepted reservation request from: {src_peer_id:#?}"); + } + relay::Event::ReservationReqDenied { src_peer_id } => { + debug!("Reservation request was denied for: {src_peer_id:#?}"); + } + relay::Event::ReservationTimedOut { src_peer_id } => { + debug!("Reservation timed out for: {src_peer_id:#?}"); + } + _ => {} + }, + } + self.pending_events + .push_back(DiscoveryEvent::Discovery(Box::new(ev))); + } + ToSwarm::Dial { opts } => { + return Poll::Ready(ToSwarm::Dial { opts }); + } + ToSwarm::NotifyHandler { + peer_id, + handler, + event, + } => { + return Poll::Ready(ToSwarm::NotifyHandler { + peer_id, + handler, + event, + }) + } + ToSwarm::CloseConnection { + peer_id, + connection, + } => { + return Poll::Ready(ToSwarm::CloseConnection { + peer_id, + connection, + }) + } + ToSwarm::ListenOn { opts } => return Poll::Ready(ToSwarm::ListenOn { opts }), + ToSwarm::RemoveListener { id } => { + return Poll::Ready(ToSwarm::RemoveListener { id }) + } + ToSwarm::NewExternalAddrCandidate(addr) => { + return Poll::Ready(ToSwarm::NewExternalAddrCandidate(addr)) + } + ToSwarm::ExternalAddrConfirmed(addr) => { + return Poll::Ready(ToSwarm::ExternalAddrConfirmed(addr)) + } + ToSwarm::ExternalAddrExpired(addr) => { + return Poll::Ready(ToSwarm::ExternalAddrExpired(addr)) + } + _ => {} + } + } + + Poll::Pending + } +} diff --git a/crates/networking/src/discovery/config.rs b/crates/networking/src/discovery/config.rs new file mode 100644 index 000000000..7eb869135 --- /dev/null +++ b/crates/networking/src/discovery/config.rs @@ -0,0 +1,173 @@ +use super::{ + behaviour::{DerivedDiscoveryBehaviour, DiscoveryBehaviour}, + new_kademlia, +}; +use gadget_logging::warn; +use libp2p::{ + autonat, identify, identity::PublicKey, mdns, relay, upnp, Multiaddr, PeerId, StreamProtocol, +}; +use std::{ + collections::{HashMap, HashSet, VecDeque}, + time::Duration, +}; + +pub struct DiscoveryConfig { + /// The local peer ID. + local_peer_id: PeerId, + /// The local public key. + local_public_key: PublicKey, + /// The bootstrap peers. + bootstrap_peers: Vec<(PeerId, Multiaddr)>, + /// The relay nodes. + relay_nodes: Vec<(PeerId, Multiaddr)>, + /// The number of peers to connect to. + target_peer_count: u64, + /// Enable mDNS discovery. + enable_mdns: bool, + /// Enable Kademlia discovery. + enable_kademlia: bool, + /// Enable UPnP discovery. + enable_upnp: bool, + /// Enable relay nodes. + enable_relay: bool, + /// The name of the network. + network_name: String, + /// Protocol version string that uniquely identifies your P2P service. + /// This should be unique to your application to avoid conflicts with other P2P networks. + /// Format recommendation: "/" + /// Example: "my-blockchain/1.0.0" or "my-chat-app/0.1.0" + protocol_version: String, +} + +impl DiscoveryConfig { + pub fn new(local_public_key: PublicKey, network_name: impl Into) -> Self { + Self { + local_peer_id: local_public_key.to_peer_id(), + local_public_key, + bootstrap_peers: Vec::new(), + relay_nodes: Vec::new(), + target_peer_count: 25, // Reasonable default + enable_mdns: true, // Enable by default for local development + enable_kademlia: true, // Enable by default for production + enable_upnp: true, // Enable by default for better connectivity + enable_relay: true, // Enable by default for relay functionality + network_name: network_name.into(), + protocol_version: String::from("gadget/1.0.0"), // Default version + } + } + + /// Set the protocol version that uniquely identifies your P2P service. + /// This should be unique to your application to avoid conflicts with other P2P networks. + /// Format recommendation: "/" + pub fn with_protocol_version(mut self, version: impl Into) -> Self { + self.protocol_version = version.into(); + self + } + + pub fn with_bootstrap_peers(mut self, peers: Vec<(PeerId, Multiaddr)>) -> Self { + self.bootstrap_peers = peers; + self + } + + pub fn with_relay_nodes(mut self, nodes: Vec<(PeerId, Multiaddr)>) -> Self { + self.relay_nodes = nodes; + self + } + + pub fn with_target_peer_count(mut self, count: u64) -> Self { + self.target_peer_count = count; + self + } + + pub fn with_mdns(mut self, enable: bool) -> Self { + self.enable_mdns = enable; + self + } + + pub fn with_kademlia(mut self, enable: bool) -> Self { + self.enable_kademlia = enable; + self + } + + pub fn with_upnp(mut self, enable: bool) -> Self { + self.enable_upnp = enable; + self + } + + pub fn with_relay(mut self, enable: bool) -> Self { + self.enable_relay = enable; + self + } + + pub fn build(self) -> anyhow::Result { + let kademlia_opt = if self.enable_kademlia { + let protocol = StreamProtocol::try_from_owned(format!( + "/gadget/kad/{}/kad/1.0.0", + self.network_name + ))?; + + let mut kademlia = new_kademlia(self.local_peer_id, protocol); + + // Add bootstrap peers + for (peer_id, addr) in &self.bootstrap_peers { + kademlia.add_address(peer_id, addr.clone()); + } + + // Start bootstrap process + if let Err(e) = kademlia.bootstrap() { + warn!("Kademlia bootstrap failed: {}", e); + } + + Some(kademlia) + } else { + None + }; + + let mdns_opt = if self.enable_mdns { + Some(mdns::Behaviour::new( + Default::default(), + self.local_peer_id, + )?) + } else { + None + }; + + let upnp_opt = if self.enable_upnp { + Some(upnp::tokio::Behaviour::default()) + } else { + None + }; + + let relay_opt = if self.enable_relay { + let relay = relay::Behaviour::new(self.local_peer_id, Default::default()); + Some(relay) + } else { + None + }; + + let behaviour = DerivedDiscoveryBehaviour { + kademlia: kademlia_opt.into(), + mdns: mdns_opt.into(), + identify: identify::Behaviour::new( + identify::Config::new(self.protocol_version, self.local_public_key) + .with_agent_version(format!("gadget-{}", env!("CARGO_PKG_VERSION"))) + .with_push_listen_addr_updates(true), + ), + autonat: autonat::Behaviour::new(self.local_peer_id, Default::default()), + upnp: upnp_opt.into(), + relay: relay_opt.into(), + }; + + Ok(DiscoveryBehaviour { + discovery: behaviour, + peers: HashSet::new(), + peer_info: HashMap::new(), + target_peer_count: self.target_peer_count, + next_kad_random_query: tokio::time::interval(Duration::from_secs(1)), + duration_to_next_kad: Duration::from_secs(1), + pending_events: VecDeque::new(), + n_node_connected: 0, + pending_dial_opts: VecDeque::new(), + }) + } +} diff --git a/crates/networking/src/discovery/mod.rs b/crates/networking/src/discovery/mod.rs new file mode 100644 index 000000000..c9aa05957 --- /dev/null +++ b/crates/networking/src/discovery/mod.rs @@ -0,0 +1,32 @@ +use std::{num::NonZero, time::Duration}; + +use libp2p::{ + kad::{self, store::MemoryStore}, + PeerId, StreamProtocol, +}; + +pub mod behaviour; +pub mod config; +pub mod peers; + +pub use peers::{PeerEvent, PeerInfo, PeerManager}; + +const MAX_ESTABLISHED_PER_PEER: u32 = 4; + +pub fn new_kademlia(peer_id: PeerId, protocol: StreamProtocol) -> kad::Behaviour { + let store = kad::store::MemoryStore::new(peer_id); + let mut config = kad::Config::new(protocol); + + // Optimize Kademlia configuration + config + .set_query_timeout(Duration::from_secs(60)) + .set_replication_factor(NonZero::new(3).unwrap()) + .set_publication_interval(Some(Duration::from_secs(120))) + .set_provider_record_ttl(Some(Duration::from_secs(24 * 60 * 60))) + .set_record_ttl(Some(Duration::from_secs(24 * 60 * 60))) + .set_parallelism(NonZero::new(5).unwrap()); + + let mut kademlia = kad::Behaviour::with_config(peer_id, store, config); + kademlia.set_mode(Some(kad::Mode::Server)); + kademlia +} diff --git a/crates/networking/src/discovery/peers.rs b/crates/networking/src/discovery/peers.rs new file mode 100644 index 000000000..01e427127 --- /dev/null +++ b/crates/networking/src/discovery/peers.rs @@ -0,0 +1,254 @@ +use std::{ + collections::{BTreeMap, HashSet}, + sync::Arc, + time::{Duration, Instant, SystemTime}, +}; + +use crate::InstanceMsgPublicKey; +use dashmap::{DashMap, DashSet}; +use libp2p::{core::Multiaddr, identify, PeerId}; +use tokio::sync::{broadcast, RwLock}; +use tracing::debug; + +/// Information about a peer's connection and behavior +#[derive(Clone, Debug)] +pub struct PeerInfo { + /// Known addresses for the peer + pub addresses: HashSet, + /// Information from the identify protocol + pub identify_info: Option, + /// When the peer was last seen + pub last_seen: SystemTime, + /// Latest ping latency + pub ping_latency: Option, + /// Number of successful protocol interactions + pub successes: u32, + /// Number of failed protocol interactions + pub failures: u32, + /// Average response time for protocol requests + pub average_response_time: Option, +} + +impl Default for PeerInfo { + fn default() -> Self { + Self { + addresses: HashSet::new(), + identify_info: None, + last_seen: SystemTime::now(), + ping_latency: None, + successes: 0, + failures: 0, + average_response_time: None, + } + } +} + +#[derive(Debug, Clone)] +pub enum PeerEvent { + /// A peer was added or updated + PeerUpdated { peer_id: PeerId, info: PeerInfo }, + /// A peer was removed + PeerRemoved { peer_id: PeerId, reason: String }, + /// A peer was banned + PeerBanned { + peer_id: PeerId, + reason: String, + expires_at: Option, + }, + /// A peer was unbanned + PeerUnbanned { peer_id: PeerId }, +} + +pub struct PeerManager { + /// Active peers and their information + peers: DashMap, + /// Completed handshakes + completed_handshakes: DashSet, + /// Handshake keys to peer ids + public_keys_to_peer_ids: Arc>>, + /// Banned peers with optional expiration time + banned_peers: DashMap>, + /// Protected peers that won't be banned + protected_peers: DashSet, + /// Event sender for peer updates + event_tx: broadcast::Sender, +} + +impl Default for PeerManager { + fn default() -> Self { + let (event_tx, _) = broadcast::channel(100); + Self { + peers: Default::default(), + banned_peers: Default::default(), + protected_peers: Default::default(), + completed_handshakes: Default::default(), + public_keys_to_peer_ids: Arc::new(RwLock::new(BTreeMap::new())), + event_tx, + } + } +} + +impl PeerManager { + /// Get a subscription to peer events + pub fn subscribe(&self) -> broadcast::Receiver { + self.event_tx.subscribe() + } + + /// Update or add peer information + pub fn update_peer(&self, peer_id: PeerId, mut info: PeerInfo) { + // Update last seen time + info.last_seen = SystemTime::now(); + + // Insert or update peer info + self.peers.insert(peer_id, info.clone()); + + // Emit event + let _ = self.event_tx.send(PeerEvent::PeerUpdated { peer_id, info }); + } + + /// Remove a peer + pub fn remove_peer(&self, peer_id: &PeerId, reason: impl Into) { + if self.peers.remove(peer_id).is_some() { + let reason = reason.into(); + debug!(%peer_id, %reason, "removed peer"); + let _ = self.event_tx.send(PeerEvent::PeerRemoved { + peer_id: *peer_id, + reason, + }); + } + } + + /// Ban a peer with optional expiration + pub fn ban_peer(&self, peer_id: PeerId, reason: impl Into, duration: Option) { + // Don't ban protected peers + if self.protected_peers.contains(&peer_id) { + return; + } + + let expires_at = duration.map(|d| Instant::now() + d); + + // Remove from active peers + self.remove_peer(&peer_id, "banned"); + + // Add to banned peers + self.banned_peers.insert(peer_id, expires_at); + + let reason = reason.into(); + debug!(%peer_id, %reason, "banned peer"); + let _ = self.event_tx.send(PeerEvent::PeerBanned { + peer_id, + reason, + expires_at, + }); + } + + /// Unban a peer + pub fn unban_peer(&self, peer_id: &PeerId) { + if self.banned_peers.remove(peer_id).is_some() { + debug!(%peer_id, "unbanned peer"); + let _ = self + .event_tx + .send(PeerEvent::PeerUnbanned { peer_id: *peer_id }); + } + } + + /// Check if a peer is banned + pub fn is_banned(&self, peer_id: &PeerId) -> bool { + self.banned_peers.contains_key(peer_id) + } + + /// Log a successful interaction with a peer + pub fn log_success(&self, peer_id: &PeerId, duration: Duration) { + if let Some(mut info) = self.peers.get_mut(peer_id) { + info.successes += 1; + update_average_time(&mut info, duration); + self.update_peer(*peer_id, info.clone()); + } + } + + /// Log a failed interaction with a peer + pub fn log_failure(&self, peer_id: &PeerId, duration: Duration) { + if let Some(mut info) = self.peers.get_mut(peer_id) { + info.failures += 1; + update_average_time(&mut info, duration); + self.update_peer(*peer_id, info.clone()); + } + } + + /// Protect a peer from being banned + pub fn protect_peer(&self, peer_id: PeerId) { + self.protected_peers.insert(peer_id); + } + + /// Remove protection from a peer + pub fn unprotect_peer(&self, peer_id: &PeerId) { + self.protected_peers.remove(peer_id); + } + + /// Get peer information + pub fn get_peer_info(&self, peer_id: &PeerId) -> Option { + self.peers.get(peer_id).map(|info| info.value().clone()) + } + + /// Get all active peers + pub fn get_peers(&self) -> DashMap { + self.peers.clone() + } + + /// Get number of active peers + pub fn peer_count(&self) -> usize { + self.peers.len() + } + + /// Start the background task to clean up expired bans + pub async fn run_ban_cleanup(self: Arc) { + loop { + let now = Instant::now(); + let mut to_unban = Vec::new(); + + // Find expired bans + let banned_peers = self.banned_peers.clone().into_read_only(); + for (peer_id, expires_at) in banned_peers.iter() { + if let Some(expiry) = expires_at { + if now >= *expiry { + to_unban.push(*peer_id); + } + } + } + + // Unban expired peers + for peer_id in to_unban { + self.unban_peer(&peer_id); + } + + tokio::time::sleep(Duration::from_secs(60)).await; + } + } + + /// Add a peer id to the public key to peer id map after verifying handshake + pub async fn add_peer_id_to_public_key( + &self, + peer_id: &PeerId, + public_key: &InstanceMsgPublicKey, + ) { + self.public_keys_to_peer_ids + .write() + .await + .insert(public_key.clone(), *peer_id); + } +} + +/// Update the average response time for a peer +fn update_average_time(info: &mut PeerInfo, duration: Duration) { + const ALPHA: u32 = 5; // Smoothing factor for the moving average + + if info.average_response_time.is_none() { + info.average_response_time = Some(duration); + } else if duration < info.average_response_time.unwrap() { + let delta = (info.average_response_time.unwrap() - duration) / ALPHA; + info.average_response_time = Some(info.average_response_time.unwrap() - delta); + } else { + let delta = (duration - info.average_response_time.unwrap()) / ALPHA; + info.average_response_time = Some(info.average_response_time.unwrap() + delta); + } +} diff --git a/crates/networking/src/error.rs b/crates/networking/src/error.rs index 9c4d13ead..3dc77e3ed 100644 --- a/crates/networking/src/error.rs +++ b/crates/networking/src/error.rs @@ -1,3 +1,5 @@ +use crate::{service::NetworkMessage, NetworkEvent}; + #[derive(Debug, thiserror::Error)] pub enum Error { #[error("Network error: {0}")] @@ -42,14 +44,22 @@ pub enum Error { // libp2p compat #[error(transparent)] Dial(#[from] libp2p::swarm::DialError), + #[error(transparent)] Noise(#[from] libp2p::noise::Error), + #[error(transparent)] Behaviour(#[from] libp2p::BehaviourBuilderError), + #[error(transparent)] Subscription(#[from] libp2p::gossipsub::SubscriptionError), + #[error(transparent)] TransportIo(#[from] libp2p::TransportError), + #[error(transparent)] Multiaddr(#[from] libp2p::multiaddr::Error), + + #[error(transparent)] + TokioSendError(#[from] tokio::sync::mpsc::error::SendError), } diff --git a/crates/networking/src/gossip.rs b/crates/networking/src/gossip.rs deleted file mode 100644 index 2f97475cb..000000000 --- a/crates/networking/src/gossip.rs +++ /dev/null @@ -1,356 +0,0 @@ -#![allow( - missing_debug_implementations, - unused_results, - clippy::module_name_repetitions, - clippy::exhaustive_enums -)] - -use crate::behaviours::{ - GossipMessage, GossipOrRequestResponse, MyBehaviour, MyBehaviourEvent, MyBehaviourRequest, -}; -use crate::error::Error; -use crate::key_types::{GossipMsgKeyPair, GossipMsgPublicKey}; -use crate::types::{IntraNodePayload, MessageType, ParticipantInfo, ProtocolMessage}; -use async_trait::async_trait; -use gadget_crypto::hashing::blake3_256; -use gadget_std::collections::BTreeMap; -use gadget_std::string::ToString; -use gadget_std::sync::atomic::AtomicUsize; -use gadget_std::sync::Arc; -use libp2p::gossipsub::IdentTopic; -use libp2p::{swarm::SwarmEvent, PeerId}; -use lru_mem::LruCache; -use tokio::sync::mpsc::UnboundedSender; -use tokio::sync::{Mutex, RwLock}; - -use crate::networking::Network; -use gadget_std::{boxed::Box, format, vec::Vec}; - -pub type InboundMapping = (IdentTopic, UnboundedSender>, Arc); - -pub struct NetworkServiceWithoutSwarm<'a> { - pub inbound_mapping: &'a [InboundMapping], - pub public_key_to_libp2p_id: Arc>>, - pub secret_key: &'a GossipMsgKeyPair, - pub connected_peers: Arc, - pub span: tracing::Span, - pub my_id: PeerId, -} - -impl<'a> NetworkServiceWithoutSwarm<'a> { - pub(crate) fn with_swarm( - &'a self, - swarm: &'a mut libp2p::Swarm, - ) -> NetworkService<'a> { - NetworkService { - swarm, - inbound_mapping: self.inbound_mapping, - public_key_to_libp2p_id: &self.public_key_to_libp2p_id, - secret_key: self.secret_key, - connected_peers: self.connected_peers.clone(), - span: &self.span, - my_id: self.my_id, - } - } -} - -pub struct NetworkService<'a> { - pub swarm: &'a mut libp2p::Swarm, - pub inbound_mapping: &'a [InboundMapping], - pub public_key_to_libp2p_id: &'a Arc>>, - pub connected_peers: Arc, - pub secret_key: &'a GossipMsgKeyPair, - pub span: &'a tracing::Span, - pub my_id: PeerId, -} - -impl NetworkService<'_> { - /// Handle local requests that are meant to be sent to the network. - pub(crate) fn handle_intra_node_payload(&mut self, msg: IntraNodePayload) { - let _enter = self.span.enter(); - match (msg.message_type, msg.payload) { - (MessageType::Broadcast, GossipOrRequestResponse::Gossip(payload)) => { - let gossip_message = bincode::serialize(&payload).expect("Should serialize"); - if let Err(e) = self - .swarm - .behaviour_mut() - .gossipsub - .publish(msg.topic, gossip_message) - { - gadget_logging::error!("Publish error: {e:?}"); - } - } - - (MessageType::P2P(peer_id), GossipOrRequestResponse::Request(req)) => { - // Send the outer payload in order to attach the topic to it - // "Requests are sent using Behaviour::send_request and the responses - // received as Message::Response via Event::Message." - self.swarm.behaviour_mut().p2p.send_request(&peer_id, req); - } - (MessageType::Broadcast, GossipOrRequestResponse::Request(_)) => { - gadget_logging::error!("Broadcasting a request is not supported"); - } - (MessageType::Broadcast, GossipOrRequestResponse::Response(_)) => { - gadget_logging::error!("Broadcasting a response is not supported"); - } - (MessageType::P2P(_), GossipOrRequestResponse::Gossip(_)) => { - gadget_logging::error!("P2P message should be a request or response"); - } - (MessageType::P2P(_), GossipOrRequestResponse::Response(_)) => { - // TODO: Send the response to the peer. - } - } - } - - /// Handle inbound events from the networking layer - #[allow(clippy::too_many_lines)] - pub(crate) async fn handle_swarm_event(&mut self, event: SwarmEvent) { - use MyBehaviourEvent::{Dcutr, Gossipsub, Identify, Kadmelia, Mdns, P2p, Ping, Relay}; - use SwarmEvent::{ - Behaviour, ConnectionClosed, ConnectionEstablished, Dialing, ExpiredListenAddr, - ExternalAddrConfirmed, ExternalAddrExpired, IncomingConnection, - IncomingConnectionError, ListenerClosed, ListenerError, NewExternalAddrCandidate, - NewExternalAddrOfPeer, NewListenAddr, OutgoingConnectionError, - }; - let _enter = self.span.enter(); - match event { - Behaviour(P2p(event)) => { - self.handle_p2p(event).await; - } - Behaviour(Gossipsub(event)) => { - self.handle_gossip(event).await; - } - Behaviour(Mdns(event)) => { - self.handle_mdns_event(event).await; - } - Behaviour(Identify(event)) => { - self.handle_identify_event(event).await; - } - Behaviour(Kadmelia(event)) => { - gadget_logging::trace!("Kadmelia event: {event:?}"); - } - Behaviour(Dcutr(event)) => { - self.handle_dcutr_event(event).await; - } - Behaviour(Relay(event)) => { - self.handle_relay_event(event).await; - } - Behaviour(Ping(event)) => { - self.handle_ping_event(event).await; - } - - NewListenAddr { - address, - listener_id, - } => { - gadget_logging::trace!("{listener_id} has a new address: {address}"); - } - ConnectionEstablished { - peer_id, - num_established, - .. - } => { - self.handle_connection_established(peer_id, num_established.get()) - .await; - } - ConnectionClosed { - peer_id, - num_established, - cause, - .. - } => { - self.handle_connection_closed(peer_id, num_established, cause) - .await; - } - IncomingConnection { - connection_id, - local_addr, - send_back_addr, - } => { - self.handle_incoming_connection(connection_id, local_addr, send_back_addr) - .await; - } - IncomingConnectionError { - connection_id, - local_addr, - send_back_addr, - error, - } => { - self.handle_incoming_connection_error( - connection_id, - local_addr, - send_back_addr, - error, - ) - .await; - } - OutgoingConnectionError { - connection_id, - peer_id, - error, - } => { - self.handle_outgoing_connection_error(connection_id, peer_id, error) - .await; - } - ExpiredListenAddr { - listener_id, - address, - } => { - gadget_logging::trace!("{listener_id} has an expired address: {address}"); - } - ListenerClosed { - listener_id, - addresses, - reason, - } => { - gadget_logging::trace!( - "{listener_id} on {addresses:?} has been closed: {reason:?}" - ); - } - ListenerError { listener_id, error } => { - gadget_logging::error!("{listener_id} has an error: {error}"); - } - Dialing { - peer_id, - connection_id, - } => { - gadget_logging::trace!( - "Dialing peer: {peer_id:?} with connection_id: {connection_id}" - ); - } - NewExternalAddrCandidate { address } => { - gadget_logging::trace!("New external address candidate: {address}"); - } - ExternalAddrConfirmed { address } => { - gadget_logging::trace!("External address confirmed: {address}"); - } - ExternalAddrExpired { address } => { - gadget_logging::trace!("External address expired: {address}"); - } - NewExternalAddrOfPeer { peer_id, address } => { - gadget_logging::trace!( - "New external address of peer: {peer_id} with address: {address}" - ); - } - unknown => { - gadget_logging::warn!("Unknown swarm event: {unknown:?}"); - } - } - } -} - -pub struct GossipHandle { - pub topic: IdentTopic, - pub tx_to_outbound: UnboundedSender, - pub rx_from_inbound: Arc>>>, - pub connected_peers: Arc, - pub public_key_to_libp2p_id: Arc>>, - pub recent_messages: parking_lot::Mutex>, - pub my_id: GossipMsgPublicKey, -} - -impl GossipHandle { - #[must_use] - pub fn connected_peers(&self) -> usize { - self.connected_peers - .load(gadget_std::sync::atomic::Ordering::Relaxed) - } - - #[must_use] - pub fn topic(&self) -> IdentTopic { - self.topic.clone() - } - - /// Returns an ordered vector of public keys of the peers that are connected to the gossipsub topic. - pub async fn peers(&self) -> Vec { - self.public_key_to_libp2p_id - .read() - .await - .keys() - .copied() - .collect() - } -} - -#[async_trait] -impl Network for GossipHandle { - async fn next_message(&self) -> Option { - loop { - let mut lock = self - .rx_from_inbound - .try_lock() - .expect("There should be only a single caller for `next_message`"); - - let message_bytes = lock.recv().await?; - drop(lock); - match bincode::deserialize::(&message_bytes) { - Ok(message) => { - let hash = blake3_256(&message_bytes); - let mut map = self.recent_messages.lock(); - if map - .insert(hash, ()) - .expect("Should not exceed memory limit (rx)") - .is_none() - { - return Some(message); - } - } - Err(e) => { - gadget_logging::error!("Failed to deserialize message (gossip): {e}"); - } - } - } - } - - async fn send_message(&self, mut message: ProtocolMessage) -> Result<(), Error> { - message.sender.public_key = Some(self.my_id); - let message_type = if let Some(ParticipantInfo { - public_key: Some(to), - .. - }) = message.recipient - { - let pub_key_to_libp2p_id = self.public_key_to_libp2p_id.read().await; - gadget_logging::trace!("Handshake count: {}", pub_key_to_libp2p_id.len()); - let libp2p_id = pub_key_to_libp2p_id - .get(&to) - .copied() - .ok_or_else(|| { - Error::NetworkError(format!( - "No libp2p ID found for crypto public key: {:?}. No handshake happened? Total handshakes: {}", - to, pub_key_to_libp2p_id.len(), - )) - })?; - - MessageType::P2P(libp2p_id) - } else { - MessageType::Broadcast - }; - - let raw_payload = - bincode::serialize(&message).map_err(|err| Error::MessagingError(err.to_string()))?; - let payload_inner = match message_type { - MessageType::Broadcast => GossipOrRequestResponse::Gossip(GossipMessage { - topic: self.topic.to_string(), - raw_payload, - }), - MessageType::P2P(_) => GossipOrRequestResponse::Request(MyBehaviourRequest::Message { - topic: self.topic.to_string(), - raw_payload, - }), - }; - - let payload = IntraNodePayload { - topic: self.topic.clone(), - payload: payload_inner, - message_type, - }; - - self.tx_to_outbound - .send(payload) - .map_err(|e| Error::NetworkError(format!("Failed to send intra-node payload: {e}"))) - } - - fn public_id(&self) -> GossipMsgPublicKey { - self.my_id - } -} diff --git a/crates/networking/src/handlers/connections.rs b/crates/networking/src/handlers/connections.rs index ea8b812bf..a511ba066 100644 --- a/crates/networking/src/handlers/connections.rs +++ b/crates/networking/src/handlers/connections.rs @@ -1,6 +1,6 @@ #![allow(unused_results, clippy::used_underscore_binding)] -use crate::behaviours::MyBehaviourRequest; +use crate::behaviours::Peer2PeerRequest; use crate::gossip::NetworkService; use crate::key_types::Curve; use gadget_crypto::KeyType; @@ -27,7 +27,7 @@ impl NetworkService<'_> { let msg = my_peer_id.to_bytes(); match ::sign_with_secret(&mut self.secret_key.clone(), &msg) { Ok(signature) => { - let handshake = MyBehaviourRequest::Handshake { + let handshake = Peer2PeerRequest::Handshake { public_key: self.secret_key.public(), signature, }; diff --git a/crates/networking/src/handlers/mdns.rs b/crates/networking/src/handlers/mdns.rs deleted file mode 100644 index 9bdd1ac8d..000000000 --- a/crates/networking/src/handlers/mdns.rs +++ /dev/null @@ -1,32 +0,0 @@ -use crate::gossip::NetworkService; -use libp2p::mdns; - -impl NetworkService<'_> { - #[tracing::instrument(skip(self, event))] - pub(crate) async fn handle_mdns_event(&mut self, event: mdns::Event) { - use mdns::Event::{Discovered, Expired}; - match event { - Discovered(list) => { - for (peer_id, multiaddr) in list { - gadget_logging::trace!("discovered a new peer: {peer_id} on {multiaddr}"); - self.swarm - .behaviour_mut() - .gossipsub - .add_explicit_peer(&peer_id); - if let Err(err) = self.swarm.dial(multiaddr) { - gadget_logging::error!("Failed to dial peer: {err}"); - } - } - } - Expired(list) => { - for (peer_id, multiaddr) in list { - gadget_logging::trace!("discover peer has expired: {peer_id} with {multiaddr}"); - self.swarm - .behaviour_mut() - .gossipsub - .remove_explicit_peer(&peer_id); - } - } - } - } -} diff --git a/crates/networking/src/handlers/mod.rs b/crates/networking/src/handlers/mod.rs index 63b39123e..c9e63fff3 100644 --- a/crates/networking/src/handlers/mod.rs +++ b/crates/networking/src/handlers/mod.rs @@ -4,7 +4,6 @@ pub mod dcutr; pub mod gossip; pub mod identify; pub mod kadmelia; -pub mod mdns; pub mod p2p; pub mod ping; pub mod relay; diff --git a/crates/networking/src/handlers/p2p.rs b/crates/networking/src/handlers/p2p.rs index be9b279ff..ee27af024 100644 --- a/crates/networking/src/handlers/p2p.rs +++ b/crates/networking/src/handlers/p2p.rs @@ -1,6 +1,6 @@ #![allow(unused_results)] -use crate::behaviours::{MyBehaviourRequest, MyBehaviourResponse}; +use crate::behaviours::{Peer2PeerRequest, Peer2PeerResponse}; use crate::gossip::NetworkService; use crate::key_types::Curve; use gadget_crypto::KeyType; @@ -13,7 +13,7 @@ impl NetworkService<'_> { #[tracing::instrument(skip(self, event))] pub(crate) async fn handle_p2p( &mut self, - event: request_response::Event, + event: request_response::Event, ) { use request_response::Event::{InboundFailure, Message, OutboundFailure, ResponseSent}; match event { @@ -57,7 +57,7 @@ impl NetworkService<'_> { async fn handle_p2p_message( &mut self, peer: PeerId, - message: request_response::Message, + message: request_response::Message, ) { use request_response::Message::{Request, Response}; match message { @@ -89,11 +89,11 @@ impl NetworkService<'_> { &mut self, peer: PeerId, request_id: request_response::InboundRequestId, - req: MyBehaviourRequest, - channel: request_response::ResponseChannel, + req: Peer2PeerRequest, + channel: request_response::ResponseChannel, ) { let result = match req { - MyBehaviourRequest::Handshake { + Peer2PeerRequest::Handshake { public_key, signature, } => { @@ -121,7 +121,7 @@ impl NetworkService<'_> { match ::sign_with_secret(&mut self.secret_key.clone(), &msg) { Ok(signature) => self.swarm.behaviour_mut().p2p.send_response( channel, - MyBehaviourResponse::Handshaked { + Peer2PeerResponse::Handshaked { public_key: self.secret_key.public(), signature, }, @@ -132,7 +132,7 @@ impl NetworkService<'_> { } } } - MyBehaviourRequest::Message { topic, raw_payload } => { + Peer2PeerRequest::Message { topic, raw_payload } => { // Reject messages from self if peer == self.my_id { return; @@ -153,7 +153,7 @@ impl NetworkService<'_> { self.swarm .behaviour_mut() .p2p - .send_response(channel, MyBehaviourResponse::MessageHandled) + .send_response(channel, Peer2PeerResponse::MessageHandled) } }; if result.is_err() { @@ -166,10 +166,10 @@ impl NetworkService<'_> { &mut self, peer: PeerId, request_id: request_response::OutboundRequestId, - message: MyBehaviourResponse, + message: Peer2PeerResponse, ) { match message { - MyBehaviourResponse::Handshaked { + Peer2PeerResponse::Handshaked { public_key, signature, } => { @@ -198,7 +198,7 @@ impl NetworkService<'_> { let _ = self.connected_peers.fetch_add(1, Ordering::Relaxed); } } - MyBehaviourResponse::MessageHandled => {} + Peer2PeerResponse::MessageHandled => {} } } } diff --git a/crates/networking/src/lib.rs b/crates/networking/src/lib.rs index 96de798ca..3a911d302 100644 --- a/crates/networking/src/lib.rs +++ b/crates/networking/src/lib.rs @@ -1,19 +1,17 @@ #![cfg_attr(docsrs, feature(doc_auto_cfg))] -pub mod gossip; -pub mod handlers; -pub mod networking; -#[cfg(feature = "round-based-compat")] -pub mod round_based_compat; -#[cfg(feature = "round-based-compat")] -pub use round_based; - pub mod behaviours; +pub mod blueprint_protocol; +pub mod discovery; pub mod error; -pub mod setup; +pub mod service; pub mod types; +#[cfg(test)] +mod tests; + pub use key_types::*; +pub use service::{NetworkConfig, NetworkEvent, NetworkMessage, NetworkService}; #[cfg(all( feature = "sp-core-ecdsa", @@ -22,8 +20,8 @@ pub use key_types::*; ))] pub mod key_types { pub use gadget_crypto::sp_core::{ - SpEcdsa as Curve, SpEcdsaPair as GossipMsgKeyPair, SpEcdsaPublic as GossipMsgPublicKey, - SpEcdsaSignature as GossipSignedMsgSignature, + SpEcdsa as Curve, SpEcdsaPair as InstanceMsgKeyPair, SpEcdsaPublic as InstanceMsgPublicKey, + SpEcdsaSignature as InstanceSignedMsgSignature, }; } @@ -34,8 +32,8 @@ pub mod key_types { ))] pub mod key_types { pub use gadget_crypto::sp_core::{ - SpSr25519 as Curve, SpSr25519Pair as GossipMsgKeyPair, - SpSr25519Public as GossipMsgPublicKey, SpSr25519Signature as GossipSignedMsgSignature, + SpSr25519 as Curve, SpSr25519Pair as InstanceMsgKeyPair, + SpSr25519Public as InstanceMsgPublicKey, SpSr25519Signature as InstanceSignedMsgSignature, }; } @@ -46,8 +44,8 @@ pub mod key_types { ))] pub mod key_types { pub use gadget_crypto::sp_core::{ - SpEd25519 as Curve, SpEd25519Pair as GossipMsgKeyPair, - SpEd25519Public as GossipMsgPublicKey, SpEd25519Signature as GossipSignedMsgSignature, + SpEd25519 as Curve, SpEd25519Pair as InstanceMsgKeyPair, + SpEd25519Public as InstanceMsgPublicKey, SpEd25519Signature as InstanceSignedMsgSignature, }; } @@ -59,8 +57,8 @@ pub mod key_types { pub mod key_types { // Default to k256 ECDSA implementation pub use gadget_crypto::k256::{ - K256Ecdsa as Curve, K256Signature as GossipSignedMsgSignature, - K256SigningKey as GossipMsgKeyPair, K256VerifyingKey as GossipMsgPublicKey, + K256Ecdsa as Curve, K256Signature as InstanceSignedMsgSignature, + K256SigningKey as InstanceMsgKeyPair, K256VerifyingKey as InstanceMsgPublicKey, }; } diff --git a/crates/networking/src/networking.rs b/crates/networking/src/networking.rs deleted file mode 100644 index 877c93fd5..000000000 --- a/crates/networking/src/networking.rs +++ /dev/null @@ -1,552 +0,0 @@ -#[cfg(test)] -mod tests; - -use crate::error::Error; -use crate::key_types::GossipMsgPublicKey; -use crate::types::{IdentifierInfo, ParticipantInfo, ProtocolMessage, UserID}; -use async_trait::async_trait; -use dashmap::DashMap; -use futures::{Stream, StreamExt}; -use gadget_crypto::hashing::blake3_256; -use gadget_std as std; -use gadget_std::boxed::Box; -use gadget_std::cmp::Reverse; -use gadget_std::collections::{BinaryHeap, HashMap}; -use gadget_std::ops::{Deref, DerefMut}; -use gadget_std::pin::Pin; -use gadget_std::string::ToString; -use gadget_std::sync::Arc; -use gadget_std::task::{Context, Poll}; -use gadget_std::vec::Vec; -use serde::{Deserialize, Serialize}; -use tokio::sync::mpsc::UnboundedSender; -use tokio::sync::Mutex; - -#[async_trait] -#[auto_impl::auto_impl(&, Box, Arc)] -pub trait Network: Send + Sync + 'static { - async fn next_message(&self) -> Option; - async fn send_message(&self, message: ProtocolMessage) -> Result<(), Error>; - - fn public_id(&self) -> GossipMsgPublicKey; - - fn build_protocol_message( - &self, - identifier_info: IdentifierInfo, - from: UserID, - to: Option, - payload: &Payload, - to_network_id: Option, - ) -> ProtocolMessage { - assert!( - (u8::from(to.is_none()) + u8::from(to_network_id.is_none()) != 1), - "Either `to` must be Some AND `to_network_id` is Some, or, both None" - ); - - let sender_participant_info = ParticipantInfo { - user_id: from, - public_key: Some(self.public_id()), - }; - let receiver_participant_info = to.map(|to| ParticipantInfo { - user_id: to, - public_key: to_network_id, - }); - ProtocolMessage { - identifier_info, - sender: sender_participant_info, - recipient: receiver_participant_info, - payload: bincode::serialize(payload).expect("Failed to serialize message"), - } - } -} - -#[derive(Debug, Serialize, Deserialize)] -struct SequencedMessage { - sequence_number: u64, - payload: Vec, -} - -#[derive(Debug)] -struct PendingMessage { - sequence_number: u64, - message: ProtocolMessage, -} - -impl PartialEq for PendingMessage { - fn eq(&self, other: &Self) -> bool { - self.sequence_number == other.sequence_number - } -} - -impl Eq for PendingMessage {} - -impl PartialOrd for PendingMessage { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for PendingMessage { - fn cmp(&self, other: &Self) -> gadget_std::cmp::Ordering { - self.sequence_number.cmp(&other.sequence_number) - } -} - -#[derive(Debug, Serialize, Deserialize)] -pub struct MultiplexedMessage { - stream_id: StreamKey, - payload: SequencedMessage, -} - -pub struct NetworkMultiplexer { - to_receiving_streams: ActiveStreams, - unclaimed_receiving_streams: Arc>, - tx_to_networking_layer: MultiplexedSender, - sequence_numbers: Arc>, - pub my_id: GossipMsgPublicKey, -} - -type ActiveStreams = Arc>>; - -#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize, Default)] -pub struct StreamKey { - pub task_hash: [u8; 32], - pub round_id: i32, -} - -impl From for StreamKey { - fn from(identifier_info: IdentifierInfo) -> Self { - let str_repr = identifier_info.to_string(); - let task_hash = blake3_256(str_repr.as_bytes()); - Self { - task_hash, - round_id: -1, - } - } -} - -pub struct MultiplexedReceiver { - inner: tokio::sync::mpsc::UnboundedReceiver, - stream_id: StreamKey, - // For post-drop removal purposes - active_streams: ActiveStreams, -} - -#[derive(Clone)] -pub struct MultiplexedSender { - inner: tokio::sync::mpsc::UnboundedSender<(StreamKey, ProtocolMessage)>, - pub(crate) stream_id: StreamKey, -} - -impl MultiplexedSender { - /// Sends a protocol message through the multiplexed channel. - /// - /// # Arguments - /// * `message` - The protocol message to send - /// - /// # Returns - /// * `Ok(())` - If the message was successfully sent - /// * `Err(Error)` - If there was an error sending the message - /// - /// # Errors - /// Returns an error if the receiving end of the channel has been closed, - /// indicating that the network connection is no longer available. - pub fn send(&self, message: ProtocolMessage) -> Result<(), Error> { - self.inner - .send((self.stream_id, message)) - .map_err(|err| Error::Other(err.to_string())) - } -} - -impl Stream for MultiplexedReceiver { - type Item = ProtocolMessage; - - fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.get_mut().inner).poll_recv(cx) - } -} - -impl Deref for MultiplexedReceiver { - type Target = tokio::sync::mpsc::UnboundedReceiver; - - fn deref(&self) -> &Self::Target { - &self.inner - } -} - -impl DerefMut for MultiplexedReceiver { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.inner - } -} - -impl Drop for MultiplexedReceiver { - fn drop(&mut self) { - let _ = self.active_streams.remove(&self.stream_id); - } -} - -// Since a single stream can be used for multiple users, and, multiple users assign seq's independently, -// we need to make a key that is unique for each (send->dest) pair and stream. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] -struct CompoundStreamKey { - stream_key: StreamKey, - send_user: UserID, - recv_user: Option, -} - -impl NetworkMultiplexer { - /// Creates a new `NetworkMultiplexer` instance. - /// - /// # Arguments - /// * `network` - The underlying network implementation that implements the Network trait - /// - /// # Type Parameters - /// * `N` - The network type that implements the Network trait - /// - /// # Returns - /// * `Self` - A new `NetworkMultiplexer` instance - /// - /// # Panics - /// This function will panic if the internal receiver has already been taken, which should not happen. - #[allow(clippy::too_many_lines)] - pub fn new(network: N) -> Self { - let (tx_to_networking_layer, mut rx_from_substreams) = - tokio::sync::mpsc::unbounded_channel(); - let my_id = network.public_id(); - let this = NetworkMultiplexer { - to_receiving_streams: Arc::new(DashMap::new()), - unclaimed_receiving_streams: Arc::new(DashMap::new()), - tx_to_networking_layer: MultiplexedSender { - inner: tx_to_networking_layer, - stream_id: StreamKey::default(), - }, - sequence_numbers: Arc::new(DashMap::new()), - my_id, - }; - - let active_streams = this.to_receiving_streams.clone(); - let unclaimed_streams = this.unclaimed_receiving_streams.clone(); - let tx_to_networking_layer = this.tx_to_networking_layer.clone(); - let sequence_numbers = this.sequence_numbers.clone(); - - drop(tokio::spawn(async move { - let network_clone = &network; - - let task1 = async move { - while let Some((stream_id, msg)) = rx_from_substreams.recv().await { - let compound_key = CompoundStreamKey { - stream_key: stream_id, - send_user: msg.sender.user_id, - recv_user: msg.recipient.as_ref().map(|p| p.user_id), - }; - - let mut seq = sequence_numbers.entry(compound_key).or_insert(0); - let current_seq = *seq; - *seq += 1; - - gadget_logging::trace!( - "SEND SEQ {current_seq} FROM {} | StreamKey: {:?}", - msg.sender.user_id, - hex::encode(bincode::serialize(&compound_key).unwrap()) - ); - - let multiplexed_message = MultiplexedMessage { - stream_id, - payload: SequencedMessage { - sequence_number: current_seq, - payload: msg.payload, - }, - }; - - let message = ProtocolMessage { - identifier_info: msg.identifier_info, - sender: msg.sender, - recipient: msg.recipient, - payload: bincode::serialize(&multiplexed_message) - .expect("Failed to serialize message"), - }; - - if let Err(err) = network_clone.send_message(message).await { - gadget_logging::error!("Failed to send message to network: {err:?}"); - break; - } - } - }; - - let task2 = async move { - let mut pending_messages: HashMap< - CompoundStreamKey, - BinaryHeap>, - > = HashMap::default(); - let mut expected_seqs: HashMap = HashMap::default(); - - while let Some(mut msg) = network_clone.next_message().await { - if let Some(recv) = msg.recipient.as_ref() { - if let Some(recv_pk) = &recv.public_key { - if recv_pk != &my_id { - gadget_logging::warn!( - "Received a message not intended for the local user" - ); - } - } - } - - let Ok(multiplexed_message) = - bincode::deserialize::(&msg.payload) - else { - gadget_logging::error!("Failed to deserialize message (networking)"); - continue; - }; - - let stream_id = multiplexed_message.stream_id; - let compound_key = CompoundStreamKey { - stream_key: stream_id, - send_user: msg.sender.user_id, - recv_user: msg.recipient.as_ref().map(|p| p.user_id), - }; - let seq = multiplexed_message.payload.sequence_number; - msg.payload = multiplexed_message.payload.payload; - - // Get or create the pending heap for this stream - let pending = pending_messages.entry(compound_key).or_default(); - let expected_seq = expected_seqs.entry(compound_key).or_default(); - - let send_user = msg.sender.user_id; - let recv_user = msg.recipient.as_ref().map(|p| p.user_id); - let compound_key_hex = hex::encode(bincode::serialize(&compound_key).unwrap()); - gadget_logging::trace!( - "RECV SEQ {seq} FROM {} as user {:?} | Expecting: {} | StreamKey: {:?}", - send_user, - recv_user, - *expected_seq, - compound_key_hex, - ); - - // Add the message to pending - pending.push(Reverse(PendingMessage { - sequence_number: seq, - message: msg, - })); - - // Try to deliver messages in order - if let Some(active_receiver) = active_streams.get(&stream_id) { - while let Some(Reverse(PendingMessage { - sequence_number, - message: _, - })) = pending.peek() - { - if *sequence_number != *expected_seq { - gadget_logging::error!( - "Sequence number mismatch, expected {} but got {}", - *expected_seq, - sequence_number - ); - break; - } - - gadget_logging::trace!("DELIVERING SEQ {seq} FROM {} as user {:?} | Expecting: {} | StreamKey: {:?}", send_user, recv_user, *expected_seq, compound_key_hex); - - *expected_seq += 1; - - let message = pending.pop().unwrap().0.message; - - if let Err(err) = active_receiver.send(message) { - gadget_logging::error!(%err, "Failed to send message to receiver"); - let _ = active_streams.remove(&stream_id); - break; - } - } - } else { - let (tx, rx) = Self::create_multiplexed_stream_inner( - tx_to_networking_layer.clone(), - &active_streams, - stream_id, - ); - - // Deliver any pending messages in order - while let Some(Reverse(PendingMessage { - sequence_number, - message: _, - })) = pending.peek() - { - if *sequence_number != *expected_seq { - gadget_logging::error!( - "Sequence number mismatch, expected {} but got {}", - *expected_seq, - sequence_number - ); - break; - } - - gadget_logging::warn!("EARLY DELIVERY SEQ {seq} FROM {} as user {:?} | Expecting: {} | StreamKey: {:?}", send_user, recv_user, *expected_seq, compound_key_hex); - - *expected_seq += 1; - - let message = pending.pop().unwrap().0.message; - - if let Err(err) = tx.send(message) { - gadget_logging::error!(%err, "Failed to send message to receiver"); - break; - } - } - - let _ = unclaimed_streams.insert(stream_id, rx); - } - } - }; - - tokio::select! { - () = task1 => { - gadget_logging::error!("Task 1 exited"); - }, - () = task2 => { - gadget_logging::error!("Task 2 exited"); - } - } - })); - - this - } - - /// Creates a new multiplexed stream. - /// - /// # Arguments - /// * `id` - The ID of the stream to create - /// - /// # Returns - /// * `Self` - A new multiplexed stream - pub fn multiplex(&self, id: impl Into) -> SubNetwork { - let id = id.into(); - let my_id = self.my_id; - let mut tx_to_networking_layer = self.tx_to_networking_layer.clone(); - if let Some(unclaimed) = self.unclaimed_receiving_streams.remove(&id) { - tx_to_networking_layer.stream_id = id; - return SubNetwork { - tx: tx_to_networking_layer, - rx: Some(unclaimed.1.into()), - my_id, - }; - } - - let (tx, rx) = Self::create_multiplexed_stream_inner( - tx_to_networking_layer, - &self.to_receiving_streams, - id, - ); - - SubNetwork { - tx, - rx: Some(rx.into()), - my_id, - } - } - - /// Creates a subnetwork, and also forwards all messages to the given channel. The network cannot be used to - /// receive messages since the messages will be forwarded to the provided channel. - /// - /// # Panics - /// - /// This function will panic if the internal receiver has already been taken, which should not happen - /// under normal circumstances. - pub fn multiplex_with_forwarding( - &self, - id: impl Into, - forward_tx: tokio::sync::mpsc::UnboundedSender, - ) -> SubNetwork { - let mut network = self.multiplex(id); - let rx = network.rx.take().expect("Rx from network should be Some"); - let forwarding_task = async move { - let mut rx = rx.into_inner(); - while let Some(msg) = rx.recv().await { - gadget_logging::info!( - "Round {}: Received message from {} to {:?} (id: {})", - msg.identifier_info.round_id, - msg.sender.user_id, - msg.recipient.as_ref().map(|p| p.user_id), - msg.identifier_info.message_id, - ); - if let Err(err) = forward_tx.send(msg) { - gadget_logging::error!(%err, "Failed to forward message to network"); - // TODO: Add AtomicBool to make sending stop - break; - } - } - }; - - drop(tokio::spawn(forwarding_task)); - - network - } - - fn create_multiplexed_stream_inner( - mut tx_to_networking_layer: MultiplexedSender, - active_streams: &ActiveStreams, - stream_id: StreamKey, - ) -> (MultiplexedSender, MultiplexedReceiver) { - let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); - if active_streams.insert(stream_id, tx).is_some() { - gadget_logging::warn!( - "Stream ID {stream_id:?} already exists! Existing stream will be replaced" - ); - } - tx_to_networking_layer.stream_id = stream_id; - - ( - tx_to_networking_layer, - MultiplexedReceiver { - inner: rx, - stream_id, - active_streams: active_streams.clone(), - }, - ) - } -} - -impl From for NetworkMultiplexer { - fn from(network: N) -> Self { - Self::new(network) - } -} - -pub struct SubNetwork { - tx: MultiplexedSender, - rx: Option>, - my_id: GossipMsgPublicKey, -} - -impl SubNetwork { - /// Sends a protocol message through the subnetwork. - /// - /// # Arguments - /// * `message` - The protocol message to send - /// - /// # Returns - /// * `Ok(())` - If the message was successfully sent - /// * `Err(Error)` - If there was an error sending the message - /// - /// # Errors - /// * Returns an error if the underlying network connection is closed or unavailable - pub fn send(&self, message: ProtocolMessage) -> Result<(), Error> { - self.tx.send(message) - } - - pub async fn recv(&self) -> Option { - self.rx.as_ref()?.lock().await.next().await - } -} - -#[async_trait] -impl Network for SubNetwork { - async fn next_message(&self) -> Option { - self.recv().await - } - - async fn send_message(&self, message: ProtocolMessage) -> Result<(), Error> { - self.send(message) - } - - fn public_id(&self) -> GossipMsgPublicKey { - self.my_id - } -} diff --git a/crates/networking/src/round_based_compat.rs b/crates/networking/src/round_based_compat.rs deleted file mode 100644 index ea9dd9e8a..000000000 --- a/crates/networking/src/round_based_compat.rs +++ /dev/null @@ -1,253 +0,0 @@ -use crate::key_types::GossipMsgPublicKey; -use crate::networking::{NetworkMultiplexer, StreamKey, SubNetwork}; -use crate::types::{IdentifierInfo, ParticipantInfo, ProtocolMessage}; -use core::pin::Pin; -use core::sync::atomic::AtomicU64; -use core::task::{ready, Context, Poll}; -use futures::prelude::*; -use gadget_std::collections::{BTreeMap, HashMap}; -use gadget_std::string::ToString; -use gadget_std::sync::Arc; -use round_based::{Delivery, Incoming, MessageType, Outgoing}; -use round_based::{MessageDestination, MsgId, PartyIndex}; -use stream::{SplitSink, SplitStream}; - -pub struct NetworkDeliveryWrapper { - /// The wrapped network implementation. - network: NetworkWrapper, -} - -impl NetworkDeliveryWrapper -where - M: Clone + Send + Unpin + 'static, - M: serde::Serialize + serde::de::DeserializeOwned, -{ - /// Create a new `NetworkDeliveryWrapper` over a network implementation with the given party index. - #[must_use] - pub fn new( - mux: Arc, - i: PartyIndex, - task_hash: [u8; 32], - parties: BTreeMap, - ) -> Self { - let (tx_forward, rx) = tokio::sync::mpsc::unbounded_channel(); - // By default, we create 10 substreams for each party. - let mut sub_streams = HashMap::new(); - for x in 0..N { - let key = StreamKey { - task_hash, - round_id: x as i32, - }; - // Creates a multiplexed subnetwork, and also forwards all messages to the given channel - let _ = sub_streams.insert(key, mux.multiplex_with_forwarding(key, tx_forward.clone())); - } - - let network = NetworkWrapper { - me: i, - mux, - message_hashes: HashMap::new(), - sub_streams, - participants: parties, - task_hash, - tx_forward, - rx, - next_msg_id: Arc::new(NextMessageId::default()), - _phantom: std::marker::PhantomData, - }; - - NetworkDeliveryWrapper { network } - } -} - -/// A `NetworkWrapper` wraps a network implementation -/// and implements [`Stream`] and [`Sink`] for it. -pub struct NetworkWrapper { - /// The current party index. - me: PartyIndex, - /// Our network Multiplexer. - mux: Arc, - /// A Map of substreams for each round. - sub_streams: HashMap, - /// A map of message hashes to their corresponding message id. - /// This is used to deduplicate messages. - message_hashes: HashMap, - /// Participants in the network with their corresponding public keys. - /// Note: This is a `BTreeMap` to ensure that the participants are sorted by their party index. - participants: BTreeMap, - /// The next message id to use. - next_msg_id: Arc, - /// A channel for forwarding messages to the network. - tx_forward: tokio::sync::mpsc::UnboundedSender, - /// A channel for receiving messages from the network. - rx: tokio::sync::mpsc::UnboundedReceiver, - /// The task hash of the current task. - task_hash: [u8; 32], - /// A phantom data type to ensure that the network wrapper is generic over the message type. - _phantom: std::marker::PhantomData, -} - -impl Delivery for NetworkDeliveryWrapper -where - M: Clone + Send + Unpin + 'static, - M: serde::Serialize + serde::de::DeserializeOwned, - M: round_based::ProtocolMessage, -{ - type Send = SplitSink, Outgoing>; - type Receive = SplitStream>; - type SendError = crate::error::Error; - type ReceiveError = crate::error::Error; - - fn split(self) -> (Self::Receive, Self::Send) { - let (sink, stream) = self.network.split(); - (stream, sink) - } -} - -impl Stream for NetworkWrapper -where - M: serde::de::DeserializeOwned + Unpin, - M: round_based::ProtocolMessage, -{ - type Item = Result, crate::error::Error>; - - fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let this = self.get_mut(); - let res = ready!(this.rx.poll_recv(cx)); - if let Some(res) = res { - let msg_type = if res.recipient.is_some() { - MessageType::P2P - } else { - MessageType::Broadcast - }; - - let id = res.identifier_info.message_id; - - let msg: M = match serde_json::from_slice(&res.payload) { - Ok(msg) => msg, - Err(err) => { - gadget_logging::error!(%err, "Failed to deserialize message (round_based_compat)"); - return Poll::Ready(Some(Err(crate::error::Error::Other(err.to_string())))); - } - }; - - let message_hash = blake3::hash(&res.payload); - gadget_logging::debug!( - "Received message with hash {} from {} in round {}", - hex::encode(message_hash.as_bytes()), - res.sender.user_id, - res.identifier_info.round_id - ); - - if this.message_hashes.contains_key(&message_hash) { - gadget_logging::warn!( - "Received duplicate message with hash {} (id: {})", - hex::encode(message_hash.as_bytes()), - id - ); - return Poll::Ready(None); - } - - this.message_hashes.insert(message_hash, id); - - Poll::Ready(Some(Ok(Incoming { - msg, - sender: res.sender.user_id, - id, - msg_type, - }))) - } else { - Poll::Ready(None) - } - } -} - -impl Sink> for NetworkWrapper -where - M: Unpin + serde::Serialize, - M: round_based::ProtocolMessage, -{ - type Error = crate::error::Error; - - fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } - - fn start_send(self: Pin<&mut Self>, out: Outgoing) -> Result<(), Self::Error> { - let this = self.get_mut(); - let id = this.next_msg_id.next(); - - let round_id = out.msg.round(); - - gadget_logging::info!( - "Round {}: Sending message from {} to {:?} (id: {})", - round_id, - this.me, - out.recipient, - id, - ); - - // Get the substream to send the message to. - let key = StreamKey { - task_hash: this.task_hash, - round_id: i32::from(round_id), - }; - let substream = this.sub_streams.entry(key).or_insert_with(|| { - this.mux - .multiplex_with_forwarding(key, this.tx_forward.clone()) - }); - - let identifier_info = IdentifierInfo { - message_id: id, - round_id, - }; - let (to, to_network_id) = match out.recipient { - MessageDestination::AllParties => (None, None), - MessageDestination::OneParty(p) => (Some(p), this.participants.get(&p).copied()), - }; - - if matches!(out.recipient, MessageDestination::OneParty(_)) && to_network_id.is_none() { - gadget_logging::warn!("Recipient not found when required for {:?}", out.recipient); - return Err(crate::error::Error::Other( - "Recipient not found".to_string(), - )); - } - - // Manually construct a `ProtocolMessage` since rounds-based - // does not work well with bincode - let protocol_message = ProtocolMessage { - identifier_info, - sender: ParticipantInfo { - user_id: this.me, - public_key: this.participants.get(&this.me).copied(), - }, - recipient: to.map(|user_id| ParticipantInfo { - user_id, - public_key: to_network_id, - }), - payload: serde_json::to_vec(&out.msg).expect("Should be able to serialize message"), - }; - - match substream.send(protocol_message) { - Ok(()) => Ok(()), - Err(e) => Err(e), - } - } - - fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } - - fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } -} - -#[derive(Default)] -struct NextMessageId(AtomicU64); - -impl NextMessageId { - fn next(&self) -> MsgId { - self.0 - .fetch_add(1, gadget_std::sync::atomic::Ordering::Relaxed) - } -} diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs new file mode 100644 index 000000000..32cda2bb8 --- /dev/null +++ b/crates/networking/src/service.rs @@ -0,0 +1,376 @@ +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use crate::{ + behaviours::{GadgetBehaviour, GadgetBehaviourEvent}, + blueprint_protocol::{BlueprintProtocolEvent, InstanceMessageRequest, InstanceMessageResponse}, + discovery::{behaviour::DiscoveryEvent, PeerManager}, + error::Error, + key_types::{InstanceMsgKeyPair, InstanceMsgPublicKey, InstanceSignedMsgSignature}, +}; +use futures::{Stream, StreamExt}; +use libp2p::{ + core::{transport::Boxed, upgrade}, + gossipsub::{IdentTopic, Topic}, + identity::Keypair, + noise, + swarm::{ConnectionId, SwarmEvent}, + tcp, yamux, Multiaddr, PeerId, Swarm, SwarmBuilder, Transport, +}; +use tokio::sync::mpsc; +use tokio_stream::wrappers::{IntervalStream, UnboundedReceiverStream}; +use tracing::{debug, error, info, warn}; + +/// Events emitted by the network service +#[derive(Debug)] +pub enum NetworkEvent { + /// New request received from a peer + InstanceRequestInbound { + peer: PeerId, + request: InstanceMessageRequest, + }, + /// New response received from a peer + InstanceResponseInbound { + peer: PeerId, + response: InstanceMessageResponse, + }, + /// New request sent to a peer + InstanceRequestOutbound { + peer: PeerId, + request: InstanceMessageRequest, + }, + /// Response sent to a peer + InstanceResponseOutbound { + peer: PeerId, + response: InstanceMessageResponse, + }, + /// New gossip message received + GossipReceived { + source: PeerId, + topic: String, + message: Vec, + }, + /// New gossip message sent + GossipSent { topic: String, message: Vec }, + /// Peer connected + PeerConnected(PeerId), + /// Peer disconnected + PeerDisconnected(PeerId), + /// Handshake completed successfully + HandshakeCompleted { peer: PeerId }, + /// Handshake failed + HandshakeFailed { peer: PeerId, reason: String }, +} + +#[derive(Debug)] +pub enum NetworkMessage { + InstanceRequest { + peer: PeerId, + request: InstanceMessageRequest, + }, + InstanceResponse { + peer: PeerId, + response: InstanceMessageResponse, + }, + GossipMessage { + source: PeerId, + topic: String, + message: Vec, + }, + HandshakeRequest { + peer: PeerId, + public_key: InstanceMsgPublicKey, + signature: InstanceSignedMsgSignature, + }, + HandshakeResponse { + peer: PeerId, + public_key: InstanceMsgPublicKey, + signature: InstanceSignedMsgSignature, + }, +} + +/// Configuration for the network service +#[derive(Debug)] +pub struct NetworkConfig { + /// Network name/namespace + pub network_name: String, + /// Local keypair for authentication + pub local_key: Keypair, + /// Secret key for message signing + pub secret_key: InstanceMsgKeyPair, + /// Address to listen on + pub listen_addr: Multiaddr, + /// Target number of peers to maintain + pub target_peer_count: u64, + /// Bootstrap peers to connect to + pub bootstrap_peers: Vec<(PeerId, Multiaddr)>, + /// Whether to enable mDNS discovery + pub enable_mdns: bool, + /// Whether to enable Kademlia DHT + pub enable_kademlia: bool, +} + +pub struct NetworkService { + /// The libp2p swarm + swarm: Swarm, + /// Peer manager for tracking peer states + peer_manager: Arc, + /// Channel for sending messages to the network service + network_sender: mpsc::UnboundedSender, + /// Channel for receiving messages from the network service + network_receiver: UnboundedReceiverStream, + /// Channel for sending events to the network service + event_sender: mpsc::UnboundedSender, + /// Channel for receiving events from the network service + event_receiver: UnboundedReceiverStream, + /// Network name/namespace + network_name: String, + /// Bootstrap peers + bootstrap_peers: HashMap, +} + +impl NetworkService { + /// Create a new network service + pub async fn new( + config: NetworkConfig, + ) -> Result<(Self, impl Stream), Error> { + let NetworkConfig { + network_name, + local_key, + secret_key, + listen_addr, + target_peer_count, + bootstrap_peers, + enable_mdns, + enable_kademlia, + } = config; + + let peer_manager = Arc::new(PeerManager::default()); + + // Create the swarm + let behaviour = GadgetBehaviour::new( + &network_name, + &local_key, + target_peer_count, + peer_manager.clone(), + ); + + let mut swarm = SwarmBuilder::with_existing_identity(local_key) + .with_tokio() + .with_tcp( + libp2p::tcp::Config::default().nodelay(true), + libp2p::noise::Config::new, + libp2p::yamux::Config::default, + )? + .with_quic_config(|mut config| { + config.handshake_timeout = Duration::from_secs(30); + config + }) + .with_dns()? + .with_behaviour(|_| behaviour) + .unwrap() + .build(); + + // Start listening + swarm.listen_on(listen_addr)?; + + let (network_sender, network_receiver) = mpsc::unbounded_channel(); + let (event_sender, event_receiver) = mpsc::unbounded_channel(); + + let bootstrap_peers = bootstrap_peers.into_iter().collect(); + + let service = Self { + swarm, + peer_manager, + network_sender, + network_receiver: UnboundedReceiverStream::new(network_receiver), + event_sender, + event_receiver: UnboundedReceiverStream::new(event_receiver), + network_name, + bootstrap_peers, + }; + + // Spawn network task + tokio::spawn(service.run(event_sender)); + + Ok((service, UnboundedReceiverStream::new(event_receiver))) + } + + /// Get a sender to send messages to the network service + pub fn message_sender(&self) -> mpsc::UnboundedSender { + self.network_sender.clone() + } + + /// Run the network service + async fn run(mut self, event_sender: mpsc::UnboundedSender) { + info!("Starting network service"); + + // Bootstrap with Kademlia + if let Err(e) = self.swarm.behaviour_mut().bootstrap() { + warn!("Failed to bootstrap with Kademlia: {}", e); + } + + // Connect to bootstrap peers + for (peer_id, addr) in &self.bootstrap_peers { + debug!("Dialing bootstrap peer {} at {}", peer_id, addr); + if let Err(e) = self.swarm.dial(addr.clone()) { + warn!("Failed to dial bootstrap peer: {}", e); + } + } + + let mut swarm_stream = self.swarm.fuse(); + let mut network_stream = self.network_receiver.fuse(); + loop { + tokio::select! { + message = network_stream.next() => { + match message { + Some(msg) => self.handle_network_message(msg, &event_sender).await, + None => break, + } + } + event = swarm_stream.next() => { + match event { + Some(event) => self.handle_swarm_event(event, &event_sender).await, + None => break, + } + } + } + } + + info!("Network service stopped"); + } + + /// Handle a network message + async fn handle_network_message( + &mut self, + msg: NetworkMessage, + event_sender: &mpsc::UnboundedSender, + ) { + match msg { + NetworkMessage::InstanceRequest { peer, request } => event_sender + .send(NetworkEvent::InstanceRequestOutbound { peer, request }) + .unwrap(), + NetworkMessage::InstanceResponse { peer, response } => event_sender + .send(NetworkEvent::InstanceResponseOutbound { peer, response }) + .unwrap(), + NetworkMessage::GossipMessage { + source, + topic, + message, + } => event_sender.send(NetworkEvent::GossipSent { topic, message }), + NetworkMessage::HandshakeRequest { + peer, + public_key, + signature, + } => event_sender.send(NetworkEvent::HandshakeCompleted { peer }), + NetworkMessage::HandshakeResponse { + peer, + public_key, + signature, + } => event_sender.send(NetworkEvent::HandshakeCompleted { peer }), + NetworkMessage::HandshakeFailed { peer, reason } => { + event_sender.send(NetworkEvent::HandshakeFailed { peer, reason }) + } + } + } + + /// Handle a swarm event + async fn handle_swarm_event( + &mut self, + event: SwarmEvent, + event_sender: &mpsc::UnboundedSender, + ) { + match event { + SwarmEvent::Behaviour(behaviour_event) => { + self.handle_behaviour_event(behaviour_event, event_sender) + .await + } + _ => {} + } + } + + /// Handle a behaviour event + async fn handle_behaviour_event( + &mut self, + event: GadgetBehaviourEvent, + event_sender: &mpsc::UnboundedSender, + ) { + match event { + GadgetBehaviourEvent::ConnectionLimits(_) => {} + GadgetBehaviourEvent::Discovery(discovery_event) => { + self.handle_discovery_event(discovery_event, event_sender) + .await + } + GadgetBehaviourEvent::BlueprintProtocol(blueprint_event) => { + self.handle_blueprint_event(blueprint_event, event_sender) + .await + } + GadgetBehaviourEvent::Ping(ping_event) => { + self.handle_ping_event(ping_event, event_sender).await + } + } + } + + /// Handle a discovery event + async fn handle_discovery_event( + &mut self, + event: DiscoveryEvent, + event_sender: &mpsc::UnboundedSender, + ) { + match event { + DiscoveryEvent::PeerConnected(peer_id) => { + event_sender.send(NetworkEvent::PeerConnected(peer_id)) + } + DiscoveryEvent::PeerDisconnected(peer_id) => { + event_sender.send(NetworkEvent::PeerDisconnected(peer_id)) + } + DiscoveryEvent::Discovery(derived_discovery_event) => { + self.handle_discovery_event(derived_discovery_event, event_sender) + .await + } + } + } + + /// Handle a blueprint event + async fn handle_blueprint_event( + &mut self, + event: BlueprintProtocolEvent, + event_sender: &mpsc::UnboundedSender, + ) -> Result<(), Error> { + match event { + BlueprintProtocolEvent::Request { + peer, + request, + channel, + } => event_sender.send(NetworkEvent::InstanceRequestInbound { peer, request }), + BlueprintProtocolEvent::Response { + peer, + response, + request_id, + } => event_sender.send(NetworkEvent::InstanceResponseInbound { peer, response }), + BlueprintProtocolEvent::GossipMessage { + source, + topic, + message, + } => event_sender.send(NetworkEvent::GossipReceived { + source, + topic, + message, + }), + } + } + + /// Handle a ping event + async fn handle_ping_event( + &mut self, + event: PingEvent, + event_sender: &mpsc::UnboundedSender, + ) { + match event { + PingEvent::Ping(ping_event) => { + event_sender.send(NetworkEvent::PingSent { peer, request }) + } + PingEvent::Pong(pong_event) => { + event_sender.send(NetworkEvent::PongReceived { peer, response }) + } + } +} diff --git a/crates/networking/src/setup.rs b/crates/networking/src/setup.rs deleted file mode 100644 index f28a947a7..000000000 --- a/crates/networking/src/setup.rs +++ /dev/null @@ -1,345 +0,0 @@ -#![allow(unused_results, missing_docs)] - -use ::std::net::{Ipv4Addr, Ipv6Addr}; - -use crate::behaviours::MyBehaviour; -use crate::error::Error; -use crate::gossip::{GossipHandle, NetworkServiceWithoutSwarm}; -pub use crate::key_types::GossipMsgKeyPair; -use crate::types::{IntraNodePayload, MAX_MESSAGE_SIZE}; -use futures::StreamExt; -use gadget_std as std; -use gadget_std::collections::BTreeMap; -use gadget_std::io; -use gadget_std::string::String; -use gadget_std::sync::atomic::AtomicUsize; -use gadget_std::sync::Arc; -use gadget_std::time::Duration; -use gadget_std::vec; -use gadget_std::vec::Vec; -use libp2p::multiaddr::Protocol; -use libp2p::Multiaddr; -use libp2p::{ - gossipsub, gossipsub::IdentTopic, kad::store::MemoryStore, mdns, request_response, - swarm::dial_opts::DialOpts, StreamProtocol, -}; -use lru_mem::LruCache; -use tokio::select; -use tokio::sync::{Mutex, RwLock}; -use tokio::task::{spawn, JoinHandle}; - -/// The version of the gadget sdk -pub const AGENT_VERSION: &str = "tangle/gadget-sdk/1.0.0"; -/// The version of the client -pub const CLIENT_VERSION: &str = "1.0.0"; - -/// The base network configuration for a blueprint's `libp2p` network. -/// -/// This configuration is used to setup the `libp2p` network for a blueprint. -/// Construct using [`NetworkConfig::new`] for advanced users or -/// [`NetworkConfig::new_service_network`] ordinarily. -pub struct NetworkConfig { - pub identity: libp2p::identity::Keypair, - pub secret_key: GossipMsgKeyPair, - pub bootnodes: Vec, - pub bind_port: u16, - pub topics: Vec, -} - -impl gadget_std::fmt::Debug for NetworkConfig { - fn fmt(&self, f: &mut gadget_std::fmt::Formatter<'_>) -> gadget_std::fmt::Result { - f.debug_struct("NetworkConfig") - .field("identity", &self.identity) - .field("bootnodes", &self.bootnodes) - .field("bind_port", &self.bind_port) - .field("topics", &self.topics) - .finish_non_exhaustive() - } -} - -impl NetworkConfig { - /// For advanced use only. Use `NetworkConfig::new_service_network` for ordinary use. - /// This function allows for the creation of a network with multiple topics. - #[must_use] - pub fn new( - identity: libp2p::identity::Keypair, - secret_key: GossipMsgKeyPair, - bootnodes: Vec, - bind_port: u16, - topics: Vec, - ) -> Self { - Self { - identity, - secret_key, - bootnodes, - bind_port, - topics, - } - } - - /// When constructing a network for a single service, the service name is used as the network name. - /// Each service within a blueprint must have a unique network name. - pub fn new_service_network>( - identity: libp2p::identity::Keypair, - secret_key: GossipMsgKeyPair, - bootnodes: Vec, - bind_port: u16, - service_name: T, - ) -> Self { - Self::new( - identity, - secret_key, - bootnodes, - bind_port, - vec![service_name.into()], - ) - } -} - -/// Start a P2P network with the given configuration. -/// -/// Each service will only have one network. It is necessary that each service calling this function -/// uses a distinct network name, otherwise, the network will not be able to distinguish between -/// the different services. -/// -/// # Arguments -/// -/// * `config` - The network configuration. -/// -/// # Errors -/// -/// Returns an error if the network setup fails. -pub fn start_p2p_network(config: NetworkConfig) -> Result { - if config.topics.len() != 1 { - return Err(Error::TooManyTopics(config.topics.len())); - } - - let (networks, _) = multiplexed_libp2p_network(config)?; - let network = networks.into_iter().next().ok_or(Error::NoNetworkFound)?.1; - Ok(network) -} - -pub type NetworkResult = Result<(BTreeMap, JoinHandle<()>), Error>; - -#[allow(clippy::collapsible_else_if, clippy::too_many_lines)] -/// Starts the multiplexed libp2p network with the given configuration. -/// -/// # Arguments -/// -/// * `config` - The network configuration. -/// -/// # Errors -/// -/// Returns an error if the network setup fails. -/// -/// # Panics -/// -/// Panics if the network name is invalid. -pub fn multiplexed_libp2p_network(config: NetworkConfig) -> NetworkResult { - // Setup both QUIC (UDP) and TCP transports the increase the chances of NAT traversal - - use gadget_std::collections::BTreeMap; - gadget_logging::trace!("Building P2P Network with config: {config:?}"); - let NetworkConfig { - identity, - bootnodes, - bind_port, - topics, - secret_key, - } = config; - - // Ensure all topics are unique - let topics_unique = topics - .iter() - .cloned() - .collect::>() - .into_iter() - .collect::>(); - - if topics_unique.len() != topics.len() { - return Err(Error::DuplicateTopics); - } - - let networks = topics; - - let my_pk = secret_key.public(); - let my_id = identity.public().to_peer_id(); - - let mut swarm = libp2p::SwarmBuilder::with_existing_identity(identity) - .with_tokio() - .with_tcp( - libp2p::tcp::Config::default().nodelay(true), // Allow port reuse for TCP-hole punching - libp2p::noise::Config::new, - libp2p::yamux::Config::default, - )? - .with_quic_config(|mut config| { - config.handshake_timeout = Duration::from_secs(30); - config - }) - .with_dns()? - .with_behaviour(|key| { - // Set a custom gossipsub configuration - let gossipsub_config = gossipsub::ConfigBuilder::default() - .protocol_id_prefix("/tangle/gadget-binary-sdk/meshsub") - .max_transmit_size(MAX_MESSAGE_SIZE) - .validate_messages() - .validation_mode(gossipsub::ValidationMode::Strict) // This sets the kind of message validation. The default is Strict (enforce message signing) - .build() - .map_err(|msg| io::Error::new(io::ErrorKind::Other, msg))?; // Temporary hack because `build` does not return a proper `std::error::Error`. - - // Setup gossipsub network behaviour for broadcasting - let gossipsub = gossipsub::Behaviour::new( - gossipsub::MessageAuthenticity::Signed(key.clone()), - gossipsub_config, - )?; - - // Setup mDNS for peer discovery - let mdns = - mdns::tokio::Behaviour::new(mdns::Config::default(), key.public().to_peer_id())?; - - // Setup request-response for direct messaging - let p2p_config = request_response::Config::default(); - // StreamProtocols MUST begin with a forward slash - let protocols = networks - .iter() - .map(|n| { - ( - StreamProtocol::try_from_owned(n.clone()).expect("Invalid network name"), - request_response::ProtocolSupport::Full, - ) - }) - .collect::>(); - - let p2p = request_response::Behaviour::new(protocols, p2p_config); - - // Setup the identify protocol for peers to exchange information about each other, a requirement for kadmelia DHT - let identify = libp2p::identify::Behaviour::new( - libp2p::identify::Config::new(CLIENT_VERSION.into(), key.public()) - .with_agent_version(AGENT_VERSION.into()), - ); - - // Setup kadmelia for DHT for peer discovery over a larger network - let memory_db = MemoryStore::new(key.public().to_peer_id()); - let kadmelia = libp2p::kad::Behaviour::new(key.public().to_peer_id(), memory_db); - - // Setup dcutr for upgrading existing connections to use relay against the bootnodes when necessary - // This also provided hole-punching capabilities to attempt to seek a direct connection, and fallback to relaying - // otherwise. - // dcutr = direct connection upgrade through relay - let dcutr = libp2p::dcutr::Behaviour::new(key.public().to_peer_id()); - - // Setup relay for using the dcutr-upgraded connections to relay messages for other peers when required - let relay_config = libp2p::relay::Config::default(); - let relay = libp2p::relay::Behaviour::new(key.public().to_peer_id(), relay_config); - - // Setup ping for liveness checks between connections - let ping = libp2p::ping::Behaviour::new(libp2p::ping::Config::default()); - - // Setup autonat for NAT traversal - let autonat = libp2p::autonat::Behaviour::new( - key.public().to_peer_id(), - libp2p::autonat::Config::default(), - ); - - Ok(MyBehaviour { - gossipsub, - mdns, - p2p, - identify, - kadmelia, - dcutr, - relay, - ping, - autonat, - }) - })? - .with_swarm_config(|c| c.with_idle_connection_timeout(Duration::from_secs(60))) - .build(); - - gadget_logging::trace!("~~~ Starting P2P Network Setup Phase 1 ~~~"); - - // Subscribe to all networks - let mut inbound_mapping = Vec::new(); - let (tx_to_outbound, mut rx_to_outbound) = - tokio::sync::mpsc::unbounded_channel::(); - let public_key_to_libp2p_id = Arc::new(RwLock::new(BTreeMap::new())); - let mut handles_ret = BTreeMap::new(); - let connected_peers = Arc::new(AtomicUsize::new(0)); - for network in networks { - let topic = IdentTopic::new(network.clone()); - swarm.behaviour_mut().gossipsub.subscribe(&topic)?; - let (inbound_tx, inbound_rx) = tokio::sync::mpsc::unbounded_channel(); - inbound_mapping.push((topic.clone(), inbound_tx, connected_peers.clone())); - - handles_ret.insert( - network, - GossipHandle { - connected_peers: connected_peers.clone(), - topic, - tx_to_outbound: tx_to_outbound.clone(), - rx_from_inbound: Arc::new(Mutex::new(inbound_rx)), - public_key_to_libp2p_id: public_key_to_libp2p_id.clone(), - // Each key is 32 bytes, therefore 512 messages hashes can be stored in the set - recent_messages: LruCache::new(16 * 1024).into(), - my_id: my_pk, - }, - ); - } - - gadget_logging::trace!("~~~ Starting P2P Network Setup Phase 2 ~~~"); - - // Listen on all interfaces UDP - swarm.listen_on( - Multiaddr::empty() - .with(Protocol::from(Ipv4Addr::UNSPECIFIED)) - .with(Protocol::from(Ipv6Addr::UNSPECIFIED)) - .with(Protocol::Udp(bind_port)) - .with(Protocol::QuicV1), - )?; - - // Listen on all interfaces TCP - swarm.listen_on( - Multiaddr::empty() - .with(Protocol::from(Ipv4Addr::UNSPECIFIED)) - .with(Protocol::from(Ipv6Addr::UNSPECIFIED)) - .with(Protocol::Tcp(bind_port)), - )?; - - gadget_logging::trace!("~~~ Starting P2P Network Setup Phase 3 ~~~"); - // Dial all bootnodes - for bootnode in &bootnodes { - swarm.dial( - DialOpts::unknown_peer_id() - .address(bootnode.clone()) - .build(), - )?; - } - - let worker = async move { - let span = tracing::debug_span!("network_worker"); - let _enter = span.enter(); - let service = NetworkServiceWithoutSwarm { - inbound_mapping: &inbound_mapping, - connected_peers, - public_key_to_libp2p_id, - secret_key: &secret_key, - span: tracing::debug_span!(parent: &span, "network_service"), - my_id, - }; - - loop { - select! { - // Setup outbound channel - Some(msg) = rx_to_outbound.recv() => { - service.with_swarm(&mut swarm).handle_intra_node_payload(msg); - } - event = swarm.select_next_some() => { - service.with_swarm(&mut swarm).handle_swarm_event(event).await; - } - } - } - }; - - let spawn_handle = spawn(worker); - Ok((handles_ret, spawn_handle)) -} diff --git a/crates/networking/src/tests/mod.rs b/crates/networking/src/tests/mod.rs new file mode 100644 index 000000000..ab21b2d3c --- /dev/null +++ b/crates/networking/src/tests/mod.rs @@ -0,0 +1,773 @@ +use self::std::time::Duration; +use super::*; +use crate::{ + gossip::GossipHandle, + networking::{Network, NetworkMultiplexer, StreamKey}, + types::IdentifierInfo, +}; +use futures::{stream, StreamExt}; +use gadget_crypto::hashing::blake3_256; +use gadget_crypto::KeyType; +use gadget_logging::setup_log; +use gadget_std::collections::BTreeMap; +use gadget_std::sync::LazyLock; +use serde::{Deserialize, Serialize}; +use tokio::time::sleep; + +const TOPIC: &str = "/gadget/test/1.0.0"; + +fn deserialize<'a, T>(data: &'a [u8]) -> Result +where + T: Deserialize<'a>, +{ + bincode::deserialize(data).map_err(|err| Error::Other(err.to_string())) +} + +#[derive(Serialize, Deserialize, Debug)] +struct StressTestPayload { + value: u64, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +enum Msg { + Round1(Round1Msg), + Round2(Round2Msg), + Round3(Round3Msg), +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +struct Round1Msg { + pub power: u16, + pub hitpoints: u16, + pub armor: u16, + pub name: String, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +struct Round2Msg { + pub x: u16, + pub y: u16, + pub z: u16, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +struct Round3Msg { + rotation: u16, + velocity: (u16, u16, u16), +} + +async fn wait_for_nodes_connected(nodes: &[GossipHandle]) { + let node_count = nodes.len(); + + // wait for the nodes to connect to each other + let max_retries = 10 * node_count; + let mut retry = 0; + loop { + gadget_logging::debug!(%node_count, %max_retries, %retry, "Checking if all nodes are connected to each other"); + let connected = nodes + .iter() + .map(super::super::gossip::GossipHandle::connected_peers) + .collect::>(); + + let all_connected = connected + .iter() + .enumerate() + .inspect(|(node, peers)| { + gadget_logging::debug!("Node {node} has {peers} connected peers"); + }) + .all(|(_, &peers)| peers >= node_count - 1); + if all_connected { + gadget_logging::debug!("All nodes are connected to each other"); + return; + } + sleep(Duration::from_millis(300)).await; + retry += 1; + assert!( + retry <= max_retries, + "Failed to connect all nodes to each other" + ); + } +} + +#[tokio::test(flavor = "multi_thread")] +#[allow(clippy::cast_possible_truncation)] +async fn test_p2p() { + setup_log(); + let nodes = stream::iter(0..*NODE_COUNT) + .map(|_| node()) + .collect::>() + .await; + + wait_for_nodes_connected(&nodes).await; + + let mut mapping = BTreeMap::new(); + for (i, node) in nodes.iter().enumerate() { + mapping.insert(i as u16, node.my_id); + } + + let mut tasks = Vec::new(); + for (i, node) in nodes.into_iter().enumerate() { + let task = tokio::spawn(run_protocol(node, i as u16, mapping.clone())); + tasks.push(task); + } + // Wait for all tasks to finish + let results = futures::future::try_join_all(tasks) + .await + .expect("Failed to run protocol"); + // Assert that all are okay. + assert!( + results.iter().all(std::result::Result::is_ok), + "Some nodes failed to run protocol" + ); +} + +#[allow(clippy::too_many_lines, clippy::cast_possible_truncation)] +async fn run_protocol( + node: N, + i: u16, + mapping: BTreeMap, +) -> Result<(), crate::error::Error> { + let task_hash = [0u8; 32]; + // Safety note: We should be passed a NetworkMultiplexer, and all uses of the N: Network + // used throughout the program must also use the multiplexer to prevent mixed messages. + let multiplexer = NetworkMultiplexer::new(node); + + let round1_network = multiplexer.multiplex(StreamKey { + task_hash, // To differentiate between different instances of a running program (i.e., a task) + round_id: 0, // To differentiate between different subsets of a running task + }); + + let round2_network = multiplexer.multiplex(StreamKey { + task_hash, // To differentiate between different instances of a running program (i.e., a task) + round_id: 1, // To differentiate between different subsets of a running task + }); + + let round3_network = multiplexer.multiplex(StreamKey { + task_hash, // To differentiate between different instances of a running program (i.e., a task) + round_id: 2, // To differentiate between different subsets of a running task + }); + + //let (round1_tx, round1_rx) = node. + // Round 1 (broadcast) + let msg = { + let round = Round1Msg { + power: i * 100, + hitpoints: (i + 1) * 50, + armor: i + 2, + name: format!("Player {}", i), + }; + round1_network.build_protocol_message( + IdentifierInfo { + message_id: 0, + round_id: 0, + }, + i, + None, + &Msg::Round1(round), + None, + ) + }; + + gadget_logging::debug!("Broadcast Message"); + round1_network + .send(msg) + .map_err(|_| crate::error::Error::Other("Failed to send message".into()))?; + + // Wait for all other nodes to send their messages + let mut msgs = BTreeMap::new(); + while let Some(msg) = round1_network.recv().await { + let m = deserialize::(&msg.payload).unwrap(); + gadget_logging::debug!(from = %msg.sender.user_id, ?m, "Received message"); + // Expecting Round1 message + assert!( + matches!(m, Msg::Round1(_)), + "Expected Round1 message but got {:?} from node {}", + m, + msg.sender.user_id, + ); + let old = msgs.insert(msg.sender.user_id, m); + assert!( + old.is_none(), + "Duplicate message from node {}", + msg.sender.user_id, + ); + // Break if all messages are received + if msgs.len() == *NODE_COUNT - 1 { + break; + } + } + gadget_logging::debug!("Done r1 w/ {i}"); + + // Round 2 (P2P) + let msgs = (0..*NODE_COUNT) + .map(|r| r as u16) + .filter(|&j| j != i) + .map(|j| { + let peer_pk = mapping.get(&j).copied().unwrap(); + round2_network.build_protocol_message( + IdentifierInfo { + message_id: 0, + round_id: 0, + }, + i, + Some(j), + &Msg::Round2(Round2Msg { + x: i * 10, + y: (i + 1) * 20, + z: i + 2, + }), + Some(peer_pk), + ) + }) + .collect::>(); + for msg in msgs { + let to = msg + .recipient + .map(|r| r.user_id) + .expect("Recipient should be present for P2P message. This is a bug in the test code"); + gadget_logging::debug!(%to, "Send P2P Message"); + round2_network.send(msg)?; + } + + // Wait for all other nodes to send their messages + let mut msgs = BTreeMap::new(); + while let Some(msg) = round2_network.recv().await { + let m = deserialize::(&msg.payload).unwrap(); + gadget_logging::info!( + "[Node {}] Received message from {} | Intended Recipient: {}", + i, + msg.sender.user_id, + msg.recipient + .as_ref() + .map_or_else(|| "Broadcast".into(), |r| r.user_id.to_string()) + ); + // Expecting Round2 message + assert!( + matches!(m, Msg::Round2(_)), + "Expected Round2 message but got {:?} from node {}", + m, + msg.sender.user_id, + ); + let old = msgs.insert(msg.sender.user_id, m); + assert!( + old.is_none(), + "Duplicate message from node {}", + msg.sender.user_id, + ); + // Break if all messages are received + if msgs.len() == *NODE_COUNT - 1 { + break; + } + } + gadget_logging::debug!("Done r2 w/ {i}"); + + // Round 3 (broadcast) + + let msg = { + let round = Round3Msg { + rotation: i * 30, + velocity: (i + 1, i + 2, i + 3), + }; + round3_network.build_protocol_message( + IdentifierInfo { + message_id: 0, + round_id: 0, + }, + i, + None, + &Msg::Round3(round), + None, + ) + }; + + gadget_logging::debug!("Broadcast Message"); + round3_network.send(msg)?; + + // Wait for all other nodes to send their messages + let mut msgs = BTreeMap::new(); + while let Some(msg) = round3_network.recv().await { + let m = deserialize::(&msg.payload).unwrap(); + gadget_logging::debug!(from = %msg.sender.user_id, ?m, "Received message"); + // Expecting Round3 message + assert!( + matches!(m, Msg::Round3(_)), + "Expected Round3 message but got {:?} from node {}", + m, + msg.sender.user_id, + ); + let old = msgs.insert(msg.sender.user_id, m); + assert!( + old.is_none(), + "Duplicate message from node {}", + msg.sender.user_id, + ); + // Break if all messages are received + if msgs.len() == *NODE_COUNT - 1 { + break; + } + } + gadget_logging::debug!("Done r3 w/ {i}"); + + gadget_logging::info!(node = i, "Protocol completed"); + + Ok(()) +} + +fn node_with_id() -> (GossipHandle, crate::key_types::GossipMsgKeyPair) { + let identity = libp2p::identity::Keypair::generate_ed25519(); + let crypto_key = crate::key_types::Curve::generate_with_seed(None).unwrap(); + let bind_port = 0; + let handle = crate::setup::start_p2p_network(crate::setup::NetworkConfig::new_service_network( + identity, + crypto_key.clone(), + Vec::default(), + bind_port, + TOPIC, + )) + .unwrap(); + + (handle, crypto_key) +} + +fn node() -> GossipHandle { + node_with_id().0 +} + +static NODE_COUNT: LazyLock = + LazyLock::new(|| std::env::var("IN_CI").map_or_else(|_| 10, |_| 2)); +#[allow(dead_code)] +static MESSAGE_COUNT: LazyLock = + LazyLock::new(|| std::env::var("IN_CI").map_or_else(|_| 10, |_| 100)); + +#[tokio::test(flavor = "multi_thread")] +#[allow(clippy::cast_possible_truncation)] +async fn test_stress_test_multiplexer() { + setup_log(); + gadget_logging::info!("Starting test_stress_test_multiplexer"); + + let (network0, network1) = get_networks().await; + + let multiplexer0 = NetworkMultiplexer::new(network0); + let multiplexer1 = NetworkMultiplexer::new(network1); + + let stream_key = StreamKey { + task_hash: blake3_256(&[1]), + round_id: 0, + }; + + let _subnetwork0 = multiplexer0.multiplex(stream_key); + let _subnetwork1 = multiplexer1.multiplex(stream_key); + + // Create a channel for forwarding + let (forward_tx, mut forward_rx) = tokio::sync::mpsc::unbounded_channel(); + + // Create a subnetwork with forwarding + let subnetwork0 = multiplexer0.multiplex(stream_key); + let subnetwork1 = multiplexer1.multiplex_with_forwarding(stream_key, forward_tx); + + let payload = StressTestPayload { value: 42 }; + let msg = subnetwork0.build_protocol_message( + IdentifierInfo::default(), + 0, + Some(1), + &payload, + Some(subnetwork1.public_id()), + ); + + gadget_logging::info!("Sending message from subnetwork0"); + subnetwork0.send(msg.clone()).unwrap(); + + // Message should be forwarded to the forward_rx channel + let forwarded_msg = forward_rx.recv().await.unwrap(); + let received: StressTestPayload = deserialize(&forwarded_msg.payload).unwrap(); + assert_eq!(received.value, payload.value); +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_multiplexer() { + setup_log(); + gadget_logging::info!("Starting test_nested_multiplexer"); + let (network0, network1) = get_networks().await; + + nested_multiplex(0, 10, network0, network1).await; +} + +async fn get_networks() -> (GossipHandle, GossipHandle) { + let network0 = node(); + let network1 = node(); + + let mut gossip_networks = vec![network0, network1]; + + wait_for_nodes_connected(&gossip_networks).await; + + (gossip_networks.remove(0), gossip_networks.remove(0)) +} + +async fn nested_multiplex( + cur_depth: usize, + max_depth: usize, + network0: N, + network1: N, +) { + gadget_logging::info!("At nested depth = {cur_depth}/{max_depth}"); + + if cur_depth == max_depth { + return; + } + + let multiplexer0 = NetworkMultiplexer::new(network0); + let multiplexer1 = NetworkMultiplexer::new(network1); + + let stream_key = StreamKey { + #[allow(clippy::cast_possible_truncation)] + task_hash: blake3_256(&[(cur_depth % 255) as u8]), + round_id: 0, + }; + + let subnetwork0 = multiplexer0.multiplex(stream_key); + let subnetwork1 = multiplexer1.multiplex(stream_key); + let subnetwork1_id = subnetwork1.public_id(); + + // Send a message in the subnetwork0 to subnetwork1 and vice versa, assert values of message + let payload = StressTestPayload { value: 42 }; + let msg = subnetwork0.build_protocol_message( + IdentifierInfo::default(), + 0, + Some(1), + &payload, + Some(subnetwork1_id), + ); + + gadget_logging::info!("Sending message from subnetwork0"); + subnetwork0.send(msg.clone()).unwrap(); + + // Receive message + let received_msg = subnetwork1.recv().await.unwrap(); + let received: StressTestPayload = deserialize(&received_msg.payload).unwrap(); + assert_eq!(received.value, payload.value); + + let msg = subnetwork1.build_protocol_message( + IdentifierInfo::default(), + 1, + Some(0), + &payload, + Some(subnetwork0.public_id()), + ); + + gadget_logging::info!("Sending message from subnetwork1"); + subnetwork1.send(msg.clone()).unwrap(); + + // Receive message + let received_msg = subnetwork0.recv().await.unwrap(); + let received: StressTestPayload = deserialize(&received_msg.payload).unwrap(); + assert_eq!(received.value, payload.value); + tracing::info!("Done nested depth = {cur_depth}/{max_depth}"); + + Box::pin(nested_multiplex( + cur_depth + 1, + max_depth, + subnetwork0, + subnetwork1, + )) + .await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_closed_channel_handling() { + setup_log(); + let (network0, network1) = get_networks().await; + + let multiplexer0 = NetworkMultiplexer::new(network0); + let multiplexer1 = NetworkMultiplexer::new(network1); + + let stream_key = StreamKey { + task_hash: blake3_256(&[1]), + round_id: 0, + }; + + let subnetwork0 = multiplexer0.multiplex(stream_key); + // Drop subnetwork1's receiver to simulate closed channel + let subnetwork1 = multiplexer1.multiplex(stream_key); + drop(subnetwork1); + + let payload = StressTestPayload { value: 42 }; + let msg = + subnetwork0.build_protocol_message(IdentifierInfo::default(), 0, None, &payload, None); + + // Sending to a closed channel should return an error + assert!(subnetwork0.send(msg).is_ok()); // Changed to ok() since the message will be sent but not received +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_empty_payload() { + setup_log(); + let (network0, network1) = get_networks().await; + + let multiplexer0 = NetworkMultiplexer::new(network0); + let multiplexer1 = NetworkMultiplexer::new(network1); + + let stream_key = StreamKey { + task_hash: blake3_256(&[1]), + round_id: 0, + }; + + let subnetwork0 = multiplexer0.multiplex(stream_key); + let subnetwork1 = multiplexer1.multiplex(stream_key); + + // Test empty payload + let empty_payload = StressTestPayload { value: 0 }; + let msg = subnetwork0.build_protocol_message( + IdentifierInfo::default(), + 0, + Some(1), + &empty_payload, + Some(subnetwork1.public_id()), + ); + + gadget_logging::info!("Sending message from subnetwork0"); + subnetwork0.send(msg).unwrap(); + + // Receive message + let received_msg = subnetwork1.recv().await.unwrap(); + let received: StressTestPayload = deserialize(&received_msg.payload).unwrap(); + assert_eq!(received.value, empty_payload.value); +} + +#[tokio::test(flavor = "multi_thread")] +#[allow(clippy::cast_possible_truncation)] +async fn test_concurrent_messaging() { + setup_log(); + let (network0, network1) = get_networks().await; + + let multiplexer0 = NetworkMultiplexer::new(network0); + let multiplexer1 = NetworkMultiplexer::new(network1); + + let mut send_handles = Vec::new(); + let mut receive_handles = Vec::new(); + + // Create multiple messages to send concurrently + let message_count = 10; + + // Spawn tasks to send messages + for i in 0..message_count { + let stream_key = StreamKey { + task_hash: blake3_256(&[i]), + round_id: 0, + }; + + let subnetwork0 = multiplexer0.multiplex(stream_key); + let subnetwork1 = multiplexer1.multiplex(stream_key); + let subnetwork1_id = subnetwork1.public_id(); + + let i_u64: u64 = i.into(); + let payload = StressTestPayload { value: i_u64 }; + let send_subnetwork0 = subnetwork0; + let handle = tokio::spawn(async move { + let msg = send_subnetwork0.build_protocol_message( + IdentifierInfo::default(), + 0, + Some(1), + &payload, + Some(subnetwork1_id), + ); + send_subnetwork0.send(msg).unwrap(); + }); + + send_handles.push(handle); + + // Spawn tasks to receive messages + let handle = tokio::spawn(async move { + let msg = subnetwork1.recv().await.unwrap(); + let received: StressTestPayload = deserialize(&msg.payload).unwrap(); + received.value as u8 // Return the payload value for verification + }); + + receive_handles.push(handle); + } + + // Wait for all sends to complete + for handle in send_handles { + handle.await.unwrap(); + } + + // Wait for all receives and verify we got all messages + let mut received_values = Vec::new(); + for handle in receive_handles { + received_values.push(handle.await.unwrap()); + } + + received_values.sort_unstable(); + assert_eq!(received_values.len(), message_count as usize); + for i in 0..message_count { + assert_eq!(received_values[i as usize], i); + } +} + +#[tokio::test(flavor = "multi_thread")] +#[allow(clippy::cast_possible_truncation)] +async fn test_message_ordering() { + setup_log(); + let (network0, network1) = get_networks().await; + + let multiplexer0 = NetworkMultiplexer::new(network0); + let multiplexer1 = NetworkMultiplexer::new(network1); + + let stream_key = StreamKey { + task_hash: blake3_256(&[1]), + round_id: 0, + }; + + let subnetwork0 = multiplexer0.multiplex(stream_key); + let subnetwork1 = multiplexer1.multiplex(stream_key); + + // Send messages with sequential sequence numbers + let message_count = 10; + for i in 0..message_count { + let payload = StressTestPayload { value: i }; + let msg = subnetwork0.build_protocol_message( + IdentifierInfo { + message_id: i, + ..Default::default() + }, + 0, + Some(1), + &payload, + Some(subnetwork1.public_id()), + ); + subnetwork0.send(msg).unwrap(); + } + + // Verify messages are received in order + let mut last_seq = 0; + for _ in 0..message_count { + let msg = subnetwork1.recv().await.unwrap(); + assert!( + msg.identifier_info.message_id >= last_seq, + "Messages should be received in order or equal to last sequence number" + ); + last_seq = msg.identifier_info.message_id; + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_network_id_handling() { + setup_log(); + let (network0, network1) = get_networks().await; + let _network0_id = network0.public_id(); + let network1_id = network1.public_id(); + + let multiplexer0 = NetworkMultiplexer::new(network0); + let multiplexer1 = NetworkMultiplexer::new(network1); + + let stream_key = StreamKey { + task_hash: blake3_256(&[1]), + round_id: 0, + }; + + let subnetwork0 = multiplexer0.multiplex(stream_key); + let subnetwork1 = multiplexer1.multiplex(stream_key); + + // Test sending with correct network ID + let payload = StressTestPayload { value: 42 }; + let msg = subnetwork0.build_protocol_message( + IdentifierInfo::default(), + 0, + Some(1), + &payload, + Some(network1_id), + ); + gadget_logging::info!("Sending message from subnetwork0"); + subnetwork0.send(msg.clone()).unwrap(); + + // Receive message + let received_msg = subnetwork1.recv().await.unwrap(); + let received: StressTestPayload = deserialize(&received_msg.payload).unwrap(); + assert_eq!(received.value, payload.value); + + // Test sending with wrong network ID + let wrong_key = crate::key_types::Curve::generate_with_seed(None).unwrap(); + let msg = subnetwork0.build_protocol_message( + IdentifierInfo::default(), + 0, + Some(1), + &payload, + Some(wrong_key.public()), + ); + gadget_logging::info!("Sending message from subnetwork0"); + subnetwork0.send(msg).unwrap(); + + // Message with wrong network ID should not be received + let timeout = tokio::time::sleep(tokio::time::Duration::from_millis(100)); + tokio::select! { + () = timeout => (), + _ = subnetwork1.recv() => panic!("Should not receive message with wrong network ID"), + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_stream_isolation() { + setup_log(); + let (network0, network1) = get_networks().await; + + let multiplexer0 = NetworkMultiplexer::new(network0); + let multiplexer1 = NetworkMultiplexer::new(network1); + + // Create two different stream keys + let stream_key1 = StreamKey { + task_hash: blake3_256(&[1]), + round_id: 0, + }; + let stream_key2 = StreamKey { + task_hash: blake3_256(&[2]), + round_id: 0, + }; + + let subnetwork0_stream1 = multiplexer0.multiplex(stream_key1); + let subnetwork0_stream2 = multiplexer0.multiplex(stream_key2); + let subnetwork1_stream1 = multiplexer1.multiplex(stream_key1); + let subnetwork1_stream2 = multiplexer1.multiplex(stream_key2); + + // Send messages on both streams + let payload1 = StressTestPayload { value: 1 }; + let payload2 = StressTestPayload { value: 2 }; + + let msg1 = subnetwork0_stream1.build_protocol_message( + IdentifierInfo::default(), + 0, + Some(1), // Send to node 1 + &payload1, + Some(subnetwork1_stream1.public_id()), + ); + let msg2 = subnetwork0_stream2.build_protocol_message( + IdentifierInfo::default(), + 0, + Some(1), // Send to node 1 + &payload2, + Some(subnetwork1_stream2.public_id()), + ); + + gadget_logging::info!("Sending message from subnetwork0_stream1"); + subnetwork0_stream1.send(msg1.clone()).unwrap(); + gadget_logging::info!("Sending message from subnetwork0_stream2"); + subnetwork0_stream2.send(msg2.clone()).unwrap(); + + // Verify messages are received on correct streams + gadget_logging::info!("Waiting for message on subnetwork1_stream1"); + let received_msg1 = subnetwork1_stream1.recv().await.unwrap(); + gadget_logging::info!("Waiting for message on subnetwork1_stream2"); + let received_msg2 = subnetwork1_stream2.recv().await.unwrap(); + + let received1: StressTestPayload = deserialize(&received_msg1.payload).unwrap(); + let received2: StressTestPayload = deserialize(&received_msg2.payload).unwrap(); + + assert_eq!(received1.value, payload1.value); + assert_eq!(received2.value, payload2.value); + + // Verify no cross-stream message leakage + let timeout = tokio::time::sleep(tokio::time::Duration::from_millis(100)); + tokio::select! { + () = timeout => (), + _ = subnetwork1_stream1.recv() => panic!("Should not receive more messages on stream 1"), + _ = subnetwork1_stream2.recv() => panic!("Should not receive more messages on stream 2"), + } +} diff --git a/crates/networking/src/types.rs b/crates/networking/src/types.rs index d60ccc70b..63b7e5e07 100644 --- a/crates/networking/src/types.rs +++ b/crates/networking/src/types.rs @@ -1,78 +1,112 @@ -use crate::{behaviours::GossipOrRequestResponse, key_types::GossipMsgPublicKey}; +use crate::key_types::InstanceMsgPublicKey; use libp2p::{gossipsub::IdentTopic, PeerId}; use serde::{Deserialize, Serialize}; use std::fmt::Display; -/// Maximum allowed size for a Signed Message. +/// Maximum allowed size for a message payload pub const MAX_MESSAGE_SIZE: usize = 16 * 1024 * 1024; -/// Unique identifier for a party -pub type UserID = u16; +/// Unique identifier for a participant in the network +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct ParticipantId(pub u16); -/// Type of message to be sent -pub enum MessageType { - Broadcast, - P2P(PeerId), +impl Into for ParticipantId { + fn into(self) -> u16 { + self.0 + } +} + +/// Type of message delivery mechanism +#[derive(Debug, Clone)] +pub enum MessageDelivery { + /// Broadcast to all peers via gossipsub + Broadcast { + /// The topic to broadcast on + topic: IdentTopic, + }, + /// Direct P2P message to a specific peer + Direct { + /// The target peer ID + peer_id: PeerId, + }, } -#[derive(Debug, Serialize, Deserialize, Clone, Copy, Default)] -pub struct IdentifierInfo { +/// Message routing information +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MessageRouting { + /// Unique identifier for this message pub message_id: u64, + /// The round/sequence number this message belongs to pub round_id: u16, + /// The sender's information + pub sender: ParticipantInfo, + /// Optional recipient information for direct messages + pub recipient: Option, } -impl Display for IdentifierInfo { - fn fmt(&self, f: &mut gadget_std::fmt::Formatter<'_>) -> gadget_std::fmt::Result { - let message_id = format!("message_id: {}", self.message_id); - let round_id = format!("round_id: {}", self.round_id); - write!(f, "{} {}", message_id, round_id) - } +/// Information about a participant in the network +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ParticipantInfo { + /// The participant's unique ID + pub id: ParticipantId, + /// The participant's public key (if known) + pub public_key: Option, } -#[derive(Debug, Serialize, Deserialize, Clone, Copy)] -pub struct ParticipantInfo { - pub user_id: u16, - pub public_key: Option, +/// A protocol message that can be sent over the network +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ProtocolMessage { + /// Routing information for the message + pub routing: MessageRouting, + /// The actual message payload + pub payload: Vec, } -impl Display for ParticipantInfo { - fn fmt(&self, f: &mut gadget_std::fmt::Formatter<'_>) -> gadget_std::fmt::Result { - let public_key = self - .public_key - .map(|key| format!("public_key: {:?}", key)) - .unwrap_or_default(); - write!(f, "user_id: {}, {}", self.user_id, public_key) - } +/// Internal representation of a message to be sent +#[derive(Debug)] +pub struct OutboundMessage { + /// The message to be sent + pub message: ProtocolMessage, + /// How the message should be delivered + pub delivery: MessageDelivery, } -#[derive(Debug, Serialize, Deserialize, Clone)] -pub struct ProtocolMessage { - pub identifier_info: IdentifierInfo, - pub sender: ParticipantInfo, - pub recipient: Option, - pub payload: Vec, +impl Display for ParticipantId { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "P{}", self.0) + } } -impl Display for ProtocolMessage { - fn fmt(&self, f: &mut gadget_std::fmt::Formatter<'_>) -> gadget_std::fmt::Result { +impl Display for MessageRouting { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, - "identifier_info: {}, sender: {}, recipient: {:?}, payload: {:?}", - self.identifier_info, self.sender, self.recipient, self.payload + "msg={} round={} from={} to={:?}", + self.message_id, + self.round_id, + self.sender.id, + self.recipient.as_ref().map(|r| r.id) ) } } -pub struct IntraNodePayload { - pub topic: IdentTopic, - pub payload: GossipOrRequestResponse, - pub message_type: MessageType, +impl Display for ParticipantInfo { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{} key={}", + self.id, + if self.public_key.is_some() { + "yes" + } else { + "no" + } + ) + } } -impl gadget_std::fmt::Debug for IntraNodePayload { - fn fmt(&self, f: &mut gadget_std::fmt::Formatter<'_>) -> gadget_std::fmt::Result { - f.debug_struct("IntraNodePayload") - .field("topic", &self.topic) - .finish_non_exhaustive() +impl Display for ProtocolMessage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{} payload_size={}", self.routing, self.payload.len()) } } From 161d55e7c22d581dc080ac87efa12cd8e08c3ed1 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Tue, 18 Feb 2025 08:58:04 -0700 Subject: [PATCH 02/52] chore: updates --- crates/networking/src/behaviours.rs | 2 +- crates/networking/src/service.rs | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/crates/networking/src/behaviours.rs b/crates/networking/src/behaviours.rs index 6f3b4aa6c..0ee01c6f6 100644 --- a/crates/networking/src/behaviours.rs +++ b/crates/networking/src/behaviours.rs @@ -39,7 +39,7 @@ pub struct GadgetBehaviour { connection_limits: connection_limits::Behaviour, /// Discovery mechanisms (Kademlia, mDNS, etc) discovery: DiscoveryBehaviour, - /// Direct P2P messaging + /// Direct P2P messaging and gossip blueprint_protocol: BlueprintProtocolBehaviour, /// Connection liveness checks ping: ping::Behaviour, diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index 32cda2bb8..9749d241e 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -353,7 +353,7 @@ impl NetworkService { message, } => event_sender.send(NetworkEvent::GossipReceived { source, - topic, + topic: topic.to_string(), message, }), } @@ -371,6 +371,7 @@ impl NetworkService { } PingEvent::Pong(pong_event) => { event_sender.send(NetworkEvent::PongReceived { peer, response }) - } + } + } } } From d15b581619b3406d638c8963ed29d1b69d04a246 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Tue, 18 Feb 2025 11:24:35 -0700 Subject: [PATCH 03/52] chore: blueprint protocol handlers --- crates/networking/src/behaviours.rs | 3 + .../src/blueprint_protocol/behaviour.rs | 201 ++++++++------ .../src/blueprint_protocol/handler.rs | 249 ++++++++++++++++++ .../networking/src/blueprint_protocol/mod.rs | 1 + .../src/handlers/blueprint_protocol.rs | 1 + crates/networking/src/handlers/connections.rs | 141 ---------- crates/networking/src/handlers/dcutr.rs | 8 - crates/networking/src/handlers/discovery.rs | 1 + crates/networking/src/handlers/gossip.rs | 118 --------- crates/networking/src/handlers/identify.rs | 41 --- crates/networking/src/handlers/kadmelia.rs | 9 - crates/networking/src/handlers/mod.rs | 10 +- crates/networking/src/handlers/p2p.rs | 204 -------------- crates/networking/src/handlers/ping.rs | 7 - crates/networking/src/handlers/relay.rs | 8 - crates/networking/src/lib.rs | 1 + crates/networking/src/service.rs | 112 +++++--- 17 files changed, 453 insertions(+), 662 deletions(-) create mode 100644 crates/networking/src/blueprint_protocol/handler.rs create mode 100644 crates/networking/src/handlers/blueprint_protocol.rs delete mode 100644 crates/networking/src/handlers/connections.rs delete mode 100644 crates/networking/src/handlers/dcutr.rs create mode 100644 crates/networking/src/handlers/discovery.rs delete mode 100644 crates/networking/src/handlers/gossip.rs delete mode 100644 crates/networking/src/handlers/identify.rs delete mode 100644 crates/networking/src/handlers/kadmelia.rs delete mode 100644 crates/networking/src/handlers/p2p.rs delete mode 100644 crates/networking/src/handlers/relay.rs diff --git a/crates/networking/src/behaviours.rs b/crates/networking/src/behaviours.rs index 0ee01c6f6..ba7f0cbda 100644 --- a/crates/networking/src/behaviours.rs +++ b/crates/networking/src/behaviours.rs @@ -1,3 +1,4 @@ +use crate::key_types::{InstanceMsgKeyPair, InstanceMsgPublicKey}; use crate::{ blueprint_protocol::{BlueprintProtocolBehaviour, BlueprintProtocolEvent}, discovery::{ @@ -49,6 +50,8 @@ impl GadgetBehaviour { pub fn new( network_name: &str, local_key: &Keypair, + instance_secret_key: &InstanceMsgKeyPair, + instance_public_key: &InstanceMsgPublicKey, target_peer_count: u64, peer_manager: Arc, ) -> Self { diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index 133e939ab..26ef89d17 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -1,6 +1,6 @@ -use crate::{Curve, InstanceMsgPublicKey, InstanceSignedMsgSignature}; -use dashmap::DashMap; -use gadget_crypto::KeyType; +use crate::{Curve, InstanceMsgKeyPair, InstanceMsgPublicKey, InstanceSignedMsgSignature}; +use dashmap::{DashMap, DashSet}; +use gadget_crypto::{hashing::blake3_256, KeyType}; use gadget_logging::{debug, trace, warn}; use libp2p::{ core::transport::PortUse, @@ -23,6 +23,15 @@ use crate::discovery::PeerManager; use super::{InstanceMessageRequest, InstanceMessageResponse}; +#[derive(NetworkBehaviour)] +pub struct DerivedBlueprintProtocolBehaviour { + /// Request/response protocol for direct messaging + request_response: + request_response::cbor::Behaviour, + /// Gossipsub for broadcast messaging + gossipsub: gossipsub::Behaviour, +} + /// Events emitted by the BlueprintProtocolBehaviour #[derive(Debug)] pub enum BlueprintProtocolEvent { @@ -49,21 +58,34 @@ pub enum BlueprintProtocolEvent { /// Behaviour that handles the blueprint protocol request/response and gossip pub struct BlueprintProtocolBehaviour { /// Request/response protocol for direct messaging - request_response: - request_response::cbor::Behaviour, - /// Gossipsub for broadcast messaging - gossipsub: gossipsub::Behaviour, + blueprint_protocol: DerivedBlueprintProtocolBehaviour, /// Peer manager for tracking peer states - peer_manager: Arc, - /// Peers with outstanding handshake requests - pending_handshake_peers: DashMap, + pub(crate) peer_manager: Arc, + /// Libp2p peer ID + pub(crate) local_peer_id: PeerId, + /// Instance public key for handshakes and blueprint_protocol + pub(crate) instance_public_key: InstanceMsgPublicKey, + /// Instance secret key for handshakes and blueprint_protocol + pub(crate) instance_secret_key: InstanceMsgKeyPair, + /// Peers with pending inbound handshakes + pub(crate) inbound_handshakes: DashMap, + /// Peers with pending outbound handshakes + pub(crate) outbound_handshakes: DashMap, + /// Peers that have completed handshake verification + pub(crate) verified_peers: DashSet, /// Active response channels - response_channels: DashMap>, + pub(crate) response_channels: + DashMap>, } impl BlueprintProtocolBehaviour { /// Create a new blueprint protocol behaviour - pub fn new(local_key: &Keypair, peer_manager: Arc) -> Self { + pub fn new( + local_key: &Keypair, + instance_secret_key: &InstanceMsgKeyPair, + instance_public_key: &InstanceMsgPublicKey, + peer_manager: Arc, + ) -> Self { let protocols = vec![( StreamProtocol::new("/gadget/blueprint_protocol/1.0.0"), request_response::ProtocolSupport::Full, @@ -82,19 +104,38 @@ impl BlueprintProtocolBehaviour { ) .expect("Valid gossipsub behaviour"); - let config = libp2p::request_response::Config::default() + let config = request_response::Config::default() .with_request_timeout(Duration::from_secs(30)) .with_max_concurrent_streams(50); - Self { + let blueprint_protocol = DerivedBlueprintProtocolBehaviour { request_response: request_response::cbor::Behaviour::new(protocols, config), gossipsub, + }; + + let local_peer_id = local_key.public().to_peer_id(); + + Self { + blueprint_protocol, peer_manager, - pending_handshake_peers: DashMap::new(), + local_peer_id, + instance_public_key: instance_public_key.clone(), + instance_secret_key: instance_secret_key.clone(), + inbound_handshakes: DashMap::new(), + outbound_handshakes: DashMap::new(), + verified_peers: DashSet::new(), response_channels: DashMap::new(), } } + /// Sign a handshake message for a peer + pub(crate) fn sign_handshake(&self, peer: &PeerId) -> InstanceSignedMsgSignature { + let msg = peer.to_bytes(); + let msg_hash = blake3_256(&msg); + let signature = self.instance_secret_key.sign_prehashed(&msg_hash); + InstanceSignedMsgSignature(signature) + } + /// Send a request to a peer pub fn send_request( &mut self, @@ -102,7 +143,9 @@ impl BlueprintProtocolBehaviour { request: InstanceMessageRequest, ) -> OutboundRequestId { debug!(%peer, ?request, "sending request"); - self.request_response.send_request(peer, request) + self.blueprint_protocol + .request_response + .send_request(peer, request) } /// Send a response through a response channel @@ -112,13 +155,15 @@ impl BlueprintProtocolBehaviour { response: InstanceMessageResponse, ) -> Result<(), InstanceMessageResponse> { debug!(?response, "sending response"); - self.request_response.send_response(channel, response) + self.blueprint_protocol + .request_response + .send_response(channel, response) } /// Subscribe to a gossip topic pub fn subscribe(&mut self, topic: &str) -> Result { let topic = Sha256Topic::new(topic); - self.gossipsub.subscribe(&topic) + self.blueprint_protocol.gossipsub.subscribe(&topic) } /// Publish a message to a gossip topic @@ -128,7 +173,7 @@ impl BlueprintProtocolBehaviour { data: impl Into>, ) -> Result { let topic = Sha256Topic::new(topic); - self.gossipsub.publish(topic, data) + self.blueprint_protocol.gossipsub.publish(topic, data) } /// Verify and handle a handshake with a peer @@ -182,15 +227,10 @@ impl BlueprintProtocolBehaviour { } impl NetworkBehaviour for BlueprintProtocolBehaviour { - type ConnectionHandler = as NetworkBehaviour>::ConnectionHandler; + type ConnectionHandler = + ::ConnectionHandler; - type ToSwarm = as NetworkBehaviour>::ToSwarm; + type ToSwarm = BlueprintProtocolEvent; fn handle_established_inbound_connection( &mut self, @@ -199,12 +239,8 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { local_addr: &libp2p::Multiaddr, remote_addr: &libp2p::Multiaddr, ) -> Result, ConnectionDenied> { - self.request_response.handle_established_inbound_connection( - connection_id, - peer, - local_addr, - remote_addr, - ) + self.blueprint_protocol + .handle_established_inbound_connection(connection_id, peer, local_addr, remote_addr) } fn handle_established_outbound_connection( @@ -215,7 +251,7 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { role_override: libp2p::core::Endpoint, port_use: PortUse, ) -> Result, ConnectionDenied> { - self.request_response + self.blueprint_protocol .handle_established_outbound_connection( connection_id, peer, @@ -231,7 +267,7 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { local_addr: &libp2p::Multiaddr, remote_addr: &libp2p::Multiaddr, ) -> Result<(), ConnectionDenied> { - self.request_response.handle_pending_inbound_connection( + self.blueprint_protocol.handle_pending_inbound_connection( connection_id, local_addr, remote_addr, @@ -245,7 +281,7 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { addresses: &[libp2p::Multiaddr], effective_role: libp2p::core::Endpoint, ) -> Result, ConnectionDenied> { - self.request_response.handle_pending_outbound_connection( + self.blueprint_protocol.handle_pending_outbound_connection( connection_id, maybe_peer, addresses, @@ -259,69 +295,72 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { connection_id: ConnectionId, event: THandlerOutEvent, ) { - self.request_response + self.blueprint_protocol .on_connection_handler_event(peer_id, connection_id, event) } fn on_swarm_event(&mut self, event: FromSwarm<'_>) { if let FromSwarm::ConnectionEstablished(e) = &event { if e.other_established == 0 { - self.pending_handshake_peers - .insert(e.peer_id, Instant::now()); + self.inbound_handshakes.insert(e.peer_id, Instant::now()); } } - self.request_response.on_swarm_event(event) + self.blueprint_protocol.on_swarm_event(event) } fn poll( &mut self, cx: &mut std::task::Context<'_>, ) -> Poll>> { - if let Poll::Ready(ev) = self.request_response.poll(cx) { - // Remove a peer from `pending_handshake_peers` when its handshake request is received. - if let ToSwarm::GenerateEvent(request_response::Event::Message { - peer, - message: - request_response::Message::Request { - request: - InstanceMessageRequest::Handshake { - public_key, - signature, - }, - .. - }, - .. - }) = &ev - { - match self.handle_handshake(peer, public_key, signature) { - Ok(_) => { - self.pending_handshake_peers.remove(&peer); - } - Err(e) => { - self.handle_handshake_failure(peer, "Handshake request failed"); + while let Poll::Ready(ev) = self.blueprint_protocol.poll(cx) { + match ev { + ToSwarm::GenerateEvent(ev) => match &ev { + DerivedBlueprintProtocolBehaviourEvent::RequestResponse( + blueprint_protocol_event, + ) => self.handle_request_response_event(blueprint_protocol_event), + DerivedBlueprintProtocolBehaviourEvent::Gossipsub(gossip_event) => { + self.handle_gossipsub_event(gossip_event) } + }, + ToSwarm::Dial { opts } => { + return Poll::Ready(ToSwarm::Dial { opts }); + } + ToSwarm::NotifyHandler { + peer_id, + handler, + event, + } => { + return Poll::Ready(ToSwarm::NotifyHandler { + peer_id, + handler, + event, + }) } + ToSwarm::CloseConnection { + peer_id, + connection, + } => { + return Poll::Ready(ToSwarm::CloseConnection { + peer_id, + connection, + }) + } + ToSwarm::ListenOn { opts } => return Poll::Ready(ToSwarm::ListenOn { opts }), + ToSwarm::RemoveListener { id } => { + return Poll::Ready(ToSwarm::RemoveListener { id }) + } + ToSwarm::NewExternalAddrCandidate(addr) => { + return Poll::Ready(ToSwarm::NewExternalAddrCandidate(addr)) + } + ToSwarm::ExternalAddrConfirmed(addr) => { + return Poll::Ready(ToSwarm::ExternalAddrConfirmed(addr)) + } + ToSwarm::ExternalAddrExpired(addr) => { + return Poll::Ready(ToSwarm::ExternalAddrExpired(addr)) + } + _ => {} } - - return Poll::Ready(ev); - } - - // Track peers whose handshake requests have timed out - const INBOUND_HANDSHAKE_WAIT_TIMEOUT: Duration = Duration::from_secs(30); - let now = Instant::now(); - if let Some((&peer_id, _)) = self - .pending_handshake_peers - .clone() - .into_read_only() - .iter() - .find(|(_, &connected_instant)| { - now.duration_since(connected_instant) > INBOUND_HANDSHAKE_WAIT_TIMEOUT - }) - { - self.pending_handshake_peers.remove(&peer_id); - self.handle_handshake_failure(&peer_id, "Handshake request timeout"); - debug!(peer=%peer_id, "Handshake request timed out, marking failure"); } Poll::Pending diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs new file mode 100644 index 000000000..f9103395b --- /dev/null +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -0,0 +1,249 @@ +use std::time::{Duration, Instant}; + +use gadget_crypto::tangle_pair_signer::sp_core::offchain::Duration as GadgetDuration; +use libp2p::{gossipsub, request_response, PeerId}; +use tracing::{debug, warn}; + +use crate::key_types::InstanceMsgPublicKey; + +use super::{BlueprintProtocolBehaviour, InstanceMessageRequest, InstanceMessageResponse}; + +const INBOUND_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30); +const OUTBOUND_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30); + +impl BlueprintProtocolBehaviour { + pub fn handle_request_response_event( + &mut self, + event: request_response::Event, + ) { + match event { + request_response::Event::Message { + peer, + message: + request_response::Message::Request { + request: + InstanceMessageRequest::Handshake { + public_key, + signature, + }, + channel, + .. + }, + .. + } => { + debug!(%peer, "Received handshake request"); + + // Check if we already have a pending outbound handshake + if let Some(outbound_time) = self.outbound_handshakes.get(&peer) { + // If we have an outbound handshake and their peer_id is less than ours, + // we should wait for their response instead of responding to their request + if peer < &self.local_peer_id { + debug!(%peer, "Deferring inbound handshake - waiting for outbound response"); + return; + } + // If we have an outbound handshake and their peer_id is greater than ours, + // we should handle their request and cancel our outbound attempt + self.outbound_handshakes.remove(&peer); + } + + // Verify the handshake + match self.verify_handshake(&peer, &public_key, &signature) { + Ok(()) => { + // Store the handshake request + self.inbound_handshakes.insert(peer, Instant::now()); + + // Send handshake response + let response = InstanceMessageResponse::Handshake { + public_key: self.instance_public_key.clone(), + signature: self.sign_handshake(&peer), + }; + + if let Err(e) = self.send_response(channel, response) { + warn!(%peer, "Failed to send handshake response: {:?}", e); + self.handle_handshake_failure(&peer, "Failed to send response"); + return; + } + + // Mark handshake as completed + self.complete_handshake(&peer, &public_key); + } + Err(e) => { + warn!(%peer, "Invalid handshake request: {:?}", e); + let response = InstanceMessageResponse::Error { + code: 400, + message: format!("Invalid handshake: {:?}", e), + }; + if let Err(e) = self.send_response(channel, response) { + warn!(%peer, "Failed to send error response: {:?}", e); + } + self.handle_handshake_failure(&peer, "Invalid handshake"); + } + } + } + request_response::Event::Message { + peer, + message: + request_response::Message::Response { + response: + InstanceMessageResponse::Handshake { + public_key, + signature, + }, + .. + }, + .. + } => { + debug!(%peer, "Received handshake response"); + + // Verify we have a pending outbound handshake + if !self.outbound_handshakes.contains_key(&peer) { + warn!(%peer, "Received unexpected handshake response"); + return; + } + + // Verify the handshake + match self.verify_handshake(&peer, &public_key, &signature) { + Ok(()) => { + // Mark handshake as completed + self.complete_handshake(&peer, &public_key); + } + Err(e) => { + warn!(%peer, "Invalid handshake response: {:?}", e); + self.handle_handshake_failure(&peer, "Invalid handshake response"); + } + } + + // Remove the outbound handshake + self.outbound_handshakes.remove(&peer); + } + request_response::Event::Message { + peer, + message: + request_response::Message::Response { + response: InstanceMessageResponse::Error { code, message }, + .. + }, + .. + } => { + warn!(%peer, code, %message, "Received error response"); + self.handle_handshake_failure(&peer, &message); + } + request_response::Event::Message { + peer, + message: + request_response::Message::Request { + request: + InstanceMessageRequest::Protocol { + protocol, + payload, + metadata, + }, + channel, + .. + }, + .. + } => { + // Only accept protocol messages from peers we've completed handshakes with + if !self.is_peer_verified(&peer) { + warn!(%peer, "Received protocol message from unverified peer"); + let response = InstanceMessageResponse::Error { + code: 403, + message: "Handshake required".to_string(), + }; + if let Err(e) = self.send_response(channel, response) { + warn!(%peer, "Failed to send error response: {:?}", e); + } + return; + } + + debug!(%peer, %protocol, "Received protocol request"); + // Handle protocol message... + } + _ => {} + } + + // Check for expired handshakes + self.check_expired_handshakes(); + } + + /// Check for and remove expired handshakes + fn check_expired_handshakes(&mut self) { + let now = Instant::now(); + + // Check inbound handshakes + let expired_inbound: Vec<_> = self + .inbound_handshakes + .clone() + .into_read_only() + .iter() + .filter(|(_, &time)| now.duration_since(time) > INBOUND_HANDSHAKE_TIMEOUT) + .map(|(peer_id, _)| *peer_id) + .collect(); + + for peer_id in expired_inbound { + self.inbound_handshakes.remove(&peer_id); + self.handle_handshake_failure(&peer_id, "Inbound handshake timeout"); + } + + // Check outbound handshakes + let expired_outbound: Vec<_> = self + .outbound_handshakes + .clone() + .into_read_only() + .iter() + .filter(|(_, &time)| now.duration_since(time) > OUTBOUND_HANDSHAKE_TIMEOUT) + .map(|(peer_id, _)| *peer_id) + .collect(); + + for peer_id in expired_outbound { + self.outbound_handshakes.remove(&peer_id); + self.handle_handshake_failure(&peer_id, "Outbound handshake timeout"); + } + } + + /// Complete a successful handshake + fn complete_handshake(&mut self, peer: &PeerId, public_key: &InstanceMsgPublicKey) { + debug!(%peer, "Completed handshake"); + + // Remove from pending handshakes + self.inbound_handshakes.remove(peer); + self.outbound_handshakes.remove(peer); + + // Add to verified peers + self.verified_peers.insert(*peer); + + // Update peer manager + self.peer_manager + .add_peer_id_to_public_key(peer, public_key); + } + + /// Check if a peer has completed handshake verification + fn is_peer_verified(&self, peer: &PeerId) -> bool { + self.verified_peers.contains(peer) + } + + pub fn handle_gossipsub_event(&mut self, event: &gossipsub::Event) { + match event { + gossipsub::Event::Message { + propagation_source, + message_id, + message, + } => { + // Only accept gossip from verified peers + if !self.is_peer_verified(propagation_source) { + warn!(%propagation_source, "Received gossip from unverified peer"); + return; + } + + debug!(%propagation_source, "Received gossip message"); + } + gossipsub::Event::Subscribed { peer_id, topic } => { + debug!(%peer_id, %topic, "Peer subscribed to topic"); + } + gossipsub::Event::Unsubscribed { peer_id, topic } => { + debug!(%peer_id, %topic, "Peer unsubscribed from topic"); + } + _ => {} + } + } +} diff --git a/crates/networking/src/blueprint_protocol/mod.rs b/crates/networking/src/blueprint_protocol/mod.rs index 4de601a41..9f4efe078 100644 --- a/crates/networking/src/blueprint_protocol/mod.rs +++ b/crates/networking/src/blueprint_protocol/mod.rs @@ -1,4 +1,5 @@ mod behaviour; +mod handler; pub use behaviour::{BlueprintProtocolBehaviour, BlueprintProtocolEvent}; diff --git a/crates/networking/src/handlers/blueprint_protocol.rs b/crates/networking/src/handlers/blueprint_protocol.rs new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/crates/networking/src/handlers/blueprint_protocol.rs @@ -0,0 +1 @@ + diff --git a/crates/networking/src/handlers/connections.rs b/crates/networking/src/handlers/connections.rs deleted file mode 100644 index a511ba066..000000000 --- a/crates/networking/src/handlers/connections.rs +++ /dev/null @@ -1,141 +0,0 @@ -#![allow(unused_results, clippy::used_underscore_binding)] - -use crate::behaviours::Peer2PeerRequest; -use crate::gossip::NetworkService; -use crate::key_types::Curve; -use gadget_crypto::KeyType; -use gadget_std as std; -use itertools::Itertools; -use libp2p::PeerId; - -impl NetworkService<'_> { - #[tracing::instrument(skip(self))] - pub(crate) async fn handle_connection_established( - &mut self, - peer_id: PeerId, - _num_established: u32, - ) { - gadget_logging::debug!("Connection established"); - if !self - .public_key_to_libp2p_id - .read() - .await - .iter() - .any(|(_, id)| id == &peer_id) - { - let my_peer_id = *self.swarm.local_peer_id(); - let msg = my_peer_id.to_bytes(); - match ::sign_with_secret(&mut self.secret_key.clone(), &msg) { - Ok(signature) => { - let handshake = Peer2PeerRequest::Handshake { - public_key: self.secret_key.public(), - signature, - }; - self.swarm - .behaviour_mut() - .p2p - .send_request(&peer_id, handshake); - self.swarm - .behaviour_mut() - .gossipsub - .add_explicit_peer(&peer_id); - gadget_logging::info!("Sent handshake from {my_peer_id} to {peer_id}"); - } - Err(e) => { - gadget_logging::error!("Failed to sign handshake: {e}"); - } - } - } - } - - #[tracing::instrument(skip(self))] - pub(crate) async fn handle_connection_closed( - &mut self, - peer_id: PeerId, - num_established: u32, - _cause: Option, - ) { - gadget_logging::trace!("Connection closed"); - if num_established == 0 { - self.swarm - .behaviour_mut() - .gossipsub - .remove_explicit_peer(&peer_id); - let mut pub_key_to_libp2p_id = self.public_key_to_libp2p_id.write().await; - let len_initial = 0; - pub_key_to_libp2p_id.retain(|_, id| *id != peer_id); - if pub_key_to_libp2p_id.len() == len_initial + 1 { - self.connected_peers - .fetch_sub(1, std::sync::atomic::Ordering::Relaxed); - } - } - } - - #[tracing::instrument(skip(self))] - pub(crate) async fn handle_incoming_connection( - &mut self, - _connection_id: libp2p::swarm::ConnectionId, - _local_addr: libp2p::Multiaddr, - _send_back_addr: libp2p::Multiaddr, - ) { - gadget_logging::trace!("Incoming connection"); - } - - #[tracing::instrument(skip(self))] - async fn handle_outgoing_connection( - &mut self, - peer_id: PeerId, - _connection_id: libp2p::swarm::ConnectionId, - ) { - gadget_logging::trace!("Outgoing connection to peer: {peer_id}"); - } - - #[tracing::instrument(skip(self, error))] - pub(crate) async fn handle_incoming_connection_error( - &mut self, - _connection_id: libp2p::swarm::ConnectionId, - _local_addr: libp2p::Multiaddr, - _send_back_addr: libp2p::Multiaddr, - error: libp2p::swarm::ListenError, - ) { - gadget_logging::error!("Incoming connection error: {error}"); - } - - #[tracing::instrument(skip(self, error))] - pub(crate) async fn handle_outgoing_connection_error( - &mut self, - _connection_id: libp2p::swarm::ConnectionId, - _peer_id: Option, - error: libp2p::swarm::DialError, - ) { - if let libp2p::swarm::DialError::Transport(addrs) = error { - let read = self.public_key_to_libp2p_id.read().await; - for (addr, err) in addrs { - if let Some(peer_id) = get_peer_id_from_multiaddr(&addr) { - if !read.values().contains(&peer_id) { - gadget_logging::warn!( - "Outgoing connection error to peer: {peer_id} at {addr}: {err}", - peer_id = peer_id, - addr = addr, - err = err - ); - } - } - } - } else { - gadget_logging::error!("Outgoing connection error to peer: {error}"); - } - } -} - -fn get_peer_id_from_multiaddr(addr: &libp2p::Multiaddr) -> Option { - addr.iter() - .find_map(|proto| { - if let libp2p::multiaddr::Protocol::P2p(peer_id) = proto { - Some(Some(peer_id)) - } else { - None - } - }) - .flatten() -} diff --git a/crates/networking/src/handlers/dcutr.rs b/crates/networking/src/handlers/dcutr.rs deleted file mode 100644 index f0fd3497f..000000000 --- a/crates/networking/src/handlers/dcutr.rs +++ /dev/null @@ -1,8 +0,0 @@ -use crate::gossip::NetworkService; - -impl NetworkService<'_> { - #[tracing::instrument(skip(self, event))] - pub async fn handle_dcutr_event(&mut self, event: libp2p::dcutr::Event) { - gadget_logging::trace!("DCUTR event: {event:?}"); - } -} diff --git a/crates/networking/src/handlers/discovery.rs b/crates/networking/src/handlers/discovery.rs new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/crates/networking/src/handlers/discovery.rs @@ -0,0 +1 @@ + diff --git a/crates/networking/src/handlers/gossip.rs b/crates/networking/src/handlers/gossip.rs deleted file mode 100644 index 67a68c913..000000000 --- a/crates/networking/src/handlers/gossip.rs +++ /dev/null @@ -1,118 +0,0 @@ -#![allow(unused_results)] - -use crate::behaviours::GossipMessage; -use crate::gossip::NetworkService; -use gadget_std::string::ToString; -use gadget_std::sync::atomic::AtomicUsize; -use gadget_std::sync::Arc; -use libp2p::gossipsub::{Event, TopicHash}; -use libp2p::{gossipsub, PeerId}; - -impl NetworkService<'_> { - #[tracing::instrument(skip(self, event))] - pub(crate) async fn handle_gossip(&mut self, event: gossipsub::Event) { - let with_connected_peers = |topic: &TopicHash, f: fn(&Arc)| { - let maybe_mapping = self - .inbound_mapping - .iter() - .find(|r| r.0.to_string() == topic.to_string()); - match maybe_mapping { - Some((_, _, connected_peers)) => { - f(connected_peers); - true - } - None => false, - } - }; - match event { - Event::Message { - propagation_source, - message_id, - message, - } => { - self.handle_gossip_message(propagation_source, message_id, message) - .await; - } - Event::Subscribed { peer_id, topic } => { - let added = with_connected_peers(&topic, |_connected_peers| { - // Code commented out because each peer needs to do a request-response - // direct P2P handshake, which is where the connected_peers counter is - // incremented. Adding here will just add twice, which is undesirable. - // connected_peers.fetch_add(1, gadget_std::sync::atomic::Ordering::Relaxed); - }); - if added { - gadget_logging::trace!("{peer_id} subscribed to {topic}"); - } else { - gadget_logging::error!("{peer_id} subscribed to unknown topic: {topic}"); - } - } - Event::Unsubscribed { peer_id, topic } => { - let removed = with_connected_peers(&topic, |_connected_peers| { - // Code commented out because each peer needs to do a request-response - // direct P2P handshake, which is where the connected_peers counter is - // decremented. Subbing here will just sub twice, which is undesirable. - // connected_peers.fetch_sub(1, gadget_std::sync::atomic::Ordering::Relaxed); - }); - if removed { - gadget_logging::trace!("{peer_id} unsubscribed from {topic}"); - } else { - gadget_logging::error!("{peer_id} unsubscribed from unknown topic: {topic}"); - } - } - Event::GossipsubNotSupported { peer_id } => { - gadget_logging::trace!("{peer_id} does not support gossipsub!"); - } - Event::SlowPeer { - peer_id, - failed_messages: _, - } => { - gadget_logging::error!("{peer_id} wasn't able to download messages in time!"); - } - } - } - - #[tracing::instrument( - skip(self, message), - fields( - %_message_id, - %_propagation_source, - source = ?message.source - ) - )] - async fn handle_gossip_message( - &mut self, - _propagation_source: PeerId, - _message_id: gossipsub::MessageId, - message: gossipsub::Message, - ) { - let Some(origin) = message.source else { - gadget_logging::error!("Got message from unknown peer"); - return; - }; - - // Reject messages from self - if origin == self.my_id { - return; - } - - gadget_logging::trace!("Got message from peer: {origin}"); - match bincode::deserialize::(&message.data) { - Ok(GossipMessage { topic, raw_payload }) => { - if let Some((_, tx, _)) = self - .inbound_mapping - .iter() - .find(|r| r.0.to_string() == topic) - { - if let Err(e) = tx.send(raw_payload) { - gadget_logging::warn!("Failed to send message to worker: {e}"); - } - } else { - gadget_logging::error!("No registered worker for topic: {topic}!"); - } - } - Err(e) => { - gadget_logging::error!("Failed to deserialize message (handlers/gossip): {e}"); - } - } - } -} diff --git a/crates/networking/src/handlers/identify.rs b/crates/networking/src/handlers/identify.rs deleted file mode 100644 index 017be564b..000000000 --- a/crates/networking/src/handlers/identify.rs +++ /dev/null @@ -1,41 +0,0 @@ -use crate::gossip::NetworkService; -use gadget_std::format; - -impl NetworkService<'_> { - #[tracing::instrument(skip(self, event))] - pub(crate) async fn handle_identify_event(&mut self, event: libp2p::identify::Event) { - use libp2p::identify::Event::{Error, Pushed, Received, Sent}; - match event { - Received { peer_id, info, .. } => { - // TODO: Verify the peer info, for example the protocol version, agent version, etc. - let info_lines = [ - format!("Protocol Version: {}", info.protocol_version), - format!("Agent Version: {}", info.agent_version), - format!("Supported Protocols: {:?}", info.protocols), - ]; - let info_lines = info_lines.join(", "); - gadget_logging::trace!( - "Received identify event from peer: {peer_id} with info: {info_lines}" - ); - self.swarm.add_external_address(info.observed_addr); - } - Sent { peer_id, .. } => { - gadget_logging::trace!("Sent identify event to peer: {peer_id}"); - } - Pushed { peer_id, info, .. } => { - let info_lines = [ - format!("Protocol Version: {}", info.protocol_version), - format!("Agent Version: {}", info.agent_version), - format!("Supported Protocols: {:?}", info.protocols), - ]; - let info_lines = info_lines.join(", "); - gadget_logging::trace!( - "Pushed identify event to peer: {peer_id} with info: {info_lines}" - ); - } - Error { peer_id, error, .. } => { - gadget_logging::error!("Identify error from peer: {peer_id} with error: {error}"); - } - } - } -} diff --git a/crates/networking/src/handlers/kadmelia.rs b/crates/networking/src/handlers/kadmelia.rs deleted file mode 100644 index fec6a8a39..000000000 --- a/crates/networking/src/handlers/kadmelia.rs +++ /dev/null @@ -1,9 +0,0 @@ -use crate::gossip::NetworkService; - -impl NetworkService<'_> { - #[tracing::instrument(skip(self, event))] - async fn handle_kadmelia_event(&mut self, event: libp2p::kad::Event) { - // TODO: Handle kadmelia events - gadget_logging::trace!("Kadmelia event: {event:?}"); - } -} diff --git a/crates/networking/src/handlers/mod.rs b/crates/networking/src/handlers/mod.rs index c9e63fff3..4e5668a5d 100644 --- a/crates/networking/src/handlers/mod.rs +++ b/crates/networking/src/handlers/mod.rs @@ -1,9 +1,3 @@ -#[cfg(not(target_family = "wasm"))] -pub mod connections; -pub mod dcutr; -pub mod gossip; -pub mod identify; -pub mod kadmelia; -pub mod p2p; +pub mod blueprint_protocol; +pub mod discovery; pub mod ping; -pub mod relay; diff --git a/crates/networking/src/handlers/p2p.rs b/crates/networking/src/handlers/p2p.rs deleted file mode 100644 index ee27af024..000000000 --- a/crates/networking/src/handlers/p2p.rs +++ /dev/null @@ -1,204 +0,0 @@ -#![allow(unused_results)] - -use crate::behaviours::{Peer2PeerRequest, Peer2PeerResponse}; -use crate::gossip::NetworkService; -use crate::key_types::Curve; -use gadget_crypto::KeyType; -use gadget_std::string::ToString; -use gadget_std::sync::atomic::Ordering; -use libp2p::gossipsub::IdentTopic; -use libp2p::{request_response, PeerId}; - -impl NetworkService<'_> { - #[tracing::instrument(skip(self, event))] - pub(crate) async fn handle_p2p( - &mut self, - event: request_response::Event, - ) { - use request_response::Event::{InboundFailure, Message, OutboundFailure, ResponseSent}; - match event { - Message { - peer, - message, - connection_id: _, - } => { - gadget_logging::trace!("Received P2P message from: {peer}"); - self.handle_p2p_message(peer, message).await; - } - OutboundFailure { - peer, - request_id, - error, - connection_id: _, - } => { - gadget_logging::error!("Failed to send message to peer: {peer} with request_id: {request_id} and error: {error}"); - } - InboundFailure { - peer, - request_id, - error, - connection_id: _, - } => { - gadget_logging::error!("Failed to receive message from peer: {peer} with request_id: {request_id} and error: {error}"); - } - ResponseSent { - peer, - request_id, - connection_id: _, - } => { - gadget_logging::debug!( - "Sent response to peer: {peer} with request_id: {request_id}" - ); - } - } - } - - #[tracing::instrument(skip(self, message))] - async fn handle_p2p_message( - &mut self, - peer: PeerId, - message: request_response::Message, - ) { - use request_response::Message::{Request, Response}; - match message { - Request { - request, - channel, - request_id, - } => { - gadget_logging::trace!( - "Received request with request_id: {request_id} from peer: {peer}" - ); - self.handle_p2p_request(peer, request_id, request, channel) - .await; - } - Response { - response, - request_id, - } => { - gadget_logging::trace!( - "Received response from peer: {peer} with request_id: {request_id}" - ); - self.handle_p2p_response(peer, request_id, response).await; - } - } - } - - #[tracing::instrument(skip(self, req, channel))] - async fn handle_p2p_request( - &mut self, - peer: PeerId, - request_id: request_response::InboundRequestId, - req: Peer2PeerRequest, - channel: request_response::ResponseChannel, - ) { - let result = match req { - Peer2PeerRequest::Handshake { - public_key, - signature, - } => { - gadget_logging::trace!("Received handshake from peer: {peer}"); - // Verify the signature - let msg = peer.to_bytes(); - let valid = ::verify(&public_key, &msg, &signature); - if !valid { - gadget_logging::warn!("Invalid initial handshake signature from peer: {peer}"); - let _ = self.swarm.disconnect_peer_id(peer); - return; - } - if self - .public_key_to_libp2p_id - .write() - .await - .insert(public_key, peer) - .is_none() - { - let _ = self.connected_peers.fetch_add(1, Ordering::Relaxed); - } - // Send response with our public key - let my_peer_id = self.swarm.local_peer_id(); - let msg = my_peer_id.to_bytes(); - match ::sign_with_secret(&mut self.secret_key.clone(), &msg) { - Ok(signature) => self.swarm.behaviour_mut().p2p.send_response( - channel, - Peer2PeerResponse::Handshaked { - public_key: self.secret_key.public(), - signature, - }, - ), - Err(e) => { - gadget_logging::error!("Failed to sign message: {e}"); - return; - } - } - } - Peer2PeerRequest::Message { topic, raw_payload } => { - // Reject messages from self - if peer == self.my_id { - return; - } - - let topic = IdentTopic::new(topic); - if let Some((_, tx, _)) = self - .inbound_mapping - .iter() - .find(|r| r.0.to_string() == topic.to_string()) - { - if let Err(e) = tx.send(raw_payload) { - gadget_logging::warn!("Failed to send message to worker: {e}"); - } - } else { - gadget_logging::error!("No registered worker for topic: {topic}!"); - } - self.swarm - .behaviour_mut() - .p2p - .send_response(channel, Peer2PeerResponse::MessageHandled) - } - }; - if result.is_err() { - gadget_logging::error!("Failed to send response for {request_id}"); - } - } - - #[tracing::instrument(skip(self, message))] - async fn handle_p2p_response( - &mut self, - peer: PeerId, - request_id: request_response::OutboundRequestId, - message: Peer2PeerResponse, - ) { - match message { - Peer2PeerResponse::Handshaked { - public_key, - signature, - } => { - gadget_logging::trace!("Received handshake-ack message from peer: {peer}"); - let msg = peer.to_bytes(); - let valid = ::verify(&public_key, &msg, &signature); - if !valid { - gadget_logging::warn!( - "Invalid handshake-acknowledgement signature from peer: {peer}" - ); - // TODO: report this peer. - self.public_key_to_libp2p_id - .write() - .await - .remove(&public_key); - let _ = self.swarm.disconnect_peer_id(peer); - return; - } - if self - .public_key_to_libp2p_id - .write() - .await - .insert(public_key, peer) - .is_none() - { - let _ = self.connected_peers.fetch_add(1, Ordering::Relaxed); - } - } - Peer2PeerResponse::MessageHandled => {} - } - } -} diff --git a/crates/networking/src/handlers/ping.rs b/crates/networking/src/handlers/ping.rs index af93d89fd..8b1378917 100644 --- a/crates/networking/src/handlers/ping.rs +++ b/crates/networking/src/handlers/ping.rs @@ -1,8 +1 @@ -use crate::gossip::NetworkService; -impl NetworkService<'_> { - #[tracing::instrument(skip(self, _event))] - pub async fn handle_ping_event(&mut self, _event: libp2p::ping::Event) { - //gadget_logging::trace!("Ping event: {event:?}"); - } -} diff --git a/crates/networking/src/handlers/relay.rs b/crates/networking/src/handlers/relay.rs deleted file mode 100644 index 320132086..000000000 --- a/crates/networking/src/handlers/relay.rs +++ /dev/null @@ -1,8 +0,0 @@ -use crate::gossip::NetworkService; - -impl NetworkService<'_> { - #[tracing::instrument(skip(self, event))] - pub async fn handle_relay_event(&mut self, event: libp2p::relay::Event) { - gadget_logging::trace!("Relay event: {event:?}"); - } -} diff --git a/crates/networking/src/lib.rs b/crates/networking/src/lib.rs index 3a911d302..e2bfe502e 100644 --- a/crates/networking/src/lib.rs +++ b/crates/networking/src/lib.rs @@ -4,6 +4,7 @@ pub mod behaviours; pub mod blueprint_protocol; pub mod discovery; pub mod error; +pub mod handlers; pub mod service; pub mod types; diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index 9749d241e..131f76f90 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -3,16 +3,20 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use crate::{ behaviours::{GadgetBehaviour, GadgetBehaviourEvent}, blueprint_protocol::{BlueprintProtocolEvent, InstanceMessageRequest, InstanceMessageResponse}, - discovery::{behaviour::DiscoveryEvent, PeerManager}, + discovery::{ + behaviour::{DerivedDiscoveryBehaviour, DerivedDiscoveryBehaviourEvent, DiscoveryEvent}, + PeerManager, + }, error::Error, key_types::{InstanceMsgKeyPair, InstanceMsgPublicKey, InstanceSignedMsgSignature}, }; use futures::{Stream, StreamExt}; +use gadget_logging::trace; use libp2p::{ core::{transport::Boxed, upgrade}, gossipsub::{IdentTopic, Topic}, identity::Keypair, - noise, + noise, ping, swarm::{ConnectionId, SwarmEvent}, tcp, yamux, Multiaddr, PeerId, Swarm, SwarmBuilder, Transport, }; @@ -93,10 +97,12 @@ pub enum NetworkMessage { pub struct NetworkConfig { /// Network name/namespace pub network_name: String, + /// Instance secret key for blueprint protocol + pub instance_secret_key: InstanceMsgKeyPair, + /// Instance public key for blueprint protocol + pub instance_public_key: InstanceMsgPublicKey, /// Local keypair for authentication pub local_key: Keypair, - /// Secret key for message signing - pub secret_key: InstanceMsgKeyPair, /// Address to listen on pub listen_addr: Multiaddr, /// Target number of peers to maintain @@ -135,8 +141,9 @@ impl NetworkService { ) -> Result<(Self, impl Stream), Error> { let NetworkConfig { network_name, + instance_secret_key, + instance_public_key, local_key, - secret_key, listen_addr, target_peer_count, bootstrap_peers, @@ -150,6 +157,8 @@ impl NetworkService { let behaviour = GadgetBehaviour::new( &network_name, &local_key, + &instance_secret_key, + &instance_public_key, target_peer_count, peer_manager.clone(), ); @@ -244,14 +253,14 @@ impl NetworkService { &mut self, msg: NetworkMessage, event_sender: &mpsc::UnboundedSender, - ) { + ) -> Result<(), Error> { match msg { - NetworkMessage::InstanceRequest { peer, request } => event_sender - .send(NetworkEvent::InstanceRequestOutbound { peer, request }) - .unwrap(), - NetworkMessage::InstanceResponse { peer, response } => event_sender - .send(NetworkEvent::InstanceResponseOutbound { peer, response }) - .unwrap(), + NetworkMessage::InstanceRequest { peer, request } => { + event_sender.send(NetworkEvent::InstanceRequestOutbound { peer, request })? + } + NetworkMessage::InstanceResponse { peer, response } => { + event_sender.send(NetworkEvent::InstanceResponseOutbound { peer, response })? + } NetworkMessage::GossipMessage { source, topic, @@ -268,9 +277,11 @@ impl NetworkService { signature, } => event_sender.send(NetworkEvent::HandshakeCompleted { peer }), NetworkMessage::HandshakeFailed { peer, reason } => { - event_sender.send(NetworkEvent::HandshakeFailed { peer, reason }) + event_sender.send(NetworkEvent::HandshakeFailed { peer, reason })? } } + + Ok(()) } /// Handle a swarm event @@ -298,40 +309,53 @@ impl NetworkService { GadgetBehaviourEvent::ConnectionLimits(_) => {} GadgetBehaviourEvent::Discovery(discovery_event) => { self.handle_discovery_event(discovery_event, event_sender) - .await + .await? } GadgetBehaviourEvent::BlueprintProtocol(blueprint_event) => { - self.handle_blueprint_event(blueprint_event, event_sender) - .await + self.handle_blueprint_protocol_event(blueprint_event, event_sender) + .await? } GadgetBehaviourEvent::Ping(ping_event) => { - self.handle_ping_event(ping_event, event_sender).await + self.handle_ping_event(ping_event, event_sender).await? } } + + Ok(()) } /// Handle a discovery event async fn handle_discovery_event( &mut self, - event: DiscoveryEvent, + event: Box, event_sender: &mpsc::UnboundedSender, - ) { + ) -> Result<(), Error> { + let event = event.as_ref(); match event { - DiscoveryEvent::PeerConnected(peer_id) => { - event_sender.send(NetworkEvent::PeerConnected(peer_id)) + DerivedDiscoveryBehaviourEvent::Kademlia(event) => { + // Handle Kademlia discovery events } - DiscoveryEvent::PeerDisconnected(peer_id) => { - event_sender.send(NetworkEvent::PeerDisconnected(peer_id)) + DerivedDiscoveryBehaviourEvent::Mdns(event) => { + // Handle mDNS discovery events } - DiscoveryEvent::Discovery(derived_discovery_event) => { - self.handle_discovery_event(derived_discovery_event, event_sender) - .await + DerivedDiscoveryBehaviourEvent::Identify(event) => { + // Handle Identify protocol events + } + DerivedDiscoveryBehaviourEvent::Autonat(event) => { + // Handle AutoNAT events + } + DerivedDiscoveryBehaviourEvent::Upnp(event) => { + // Handle UPnP events + } + DerivedDiscoveryBehaviourEvent::Relay(event) => { + // Handle relay events } } + + Ok(()) } /// Handle a blueprint event - async fn handle_blueprint_event( + async fn handle_blueprint_protocol_event( &mut self, event: BlueprintProtocolEvent, event_sender: &mpsc::UnboundedSender, @@ -341,12 +365,12 @@ impl NetworkService { peer, request, channel, - } => event_sender.send(NetworkEvent::InstanceRequestInbound { peer, request }), + } => event_sender.send(NetworkEvent::InstanceRequestInbound { peer, request })?, BlueprintProtocolEvent::Response { peer, response, request_id, - } => event_sender.send(NetworkEvent::InstanceResponseInbound { peer, response }), + } => event_sender.send(NetworkEvent::InstanceResponseInbound { peer, response })?, BlueprintProtocolEvent::GossipMessage { source, topic, @@ -355,23 +379,37 @@ impl NetworkService { source, topic: topic.to_string(), message, - }), + })?, } + + Ok(()) } /// Handle a ping event async fn handle_ping_event( &mut self, - event: PingEvent, + event: ping::Event, event_sender: &mpsc::UnboundedSender, - ) { - match event { - PingEvent::Ping(ping_event) => { - event_sender.send(NetworkEvent::PingSent { peer, request }) + ) -> Result<(), Error> { + match event.result { + Ok(rtt) => { + trace!( + "PingSuccess::Ping rtt to {} is {} ms", + event.peer, + rtt.as_millis() + ); + } + Err(ping::Failure::Unsupported) => { + debug!(peer=%event.peer, "Ping protocol unsupported"); } - PingEvent::Pong(pong_event) => { - event_sender.send(NetworkEvent::PongReceived { peer, response }) + Err(ping::Failure::Timeout) => { + debug!("Ping timeout: {}", event.peer); + } + Err(ping::Failure::Other { error }) => { + debug!("Ping failure: {error}"); } } + + Ok(()) } } From 64dfa9101f191d065ba19e99e1a31c7eb1bf63d2 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Tue, 18 Feb 2025 11:27:59 -0700 Subject: [PATCH 04/52] chore: updates --- crates/networking/src/behaviours.rs | 7 ++++++- crates/networking/src/blueprint_protocol/behaviour.rs | 2 +- crates/networking/src/blueprint_protocol/handler.rs | 4 ++-- crates/networking/src/service.rs | 8 +++----- 4 files changed, 12 insertions(+), 9 deletions(-) diff --git a/crates/networking/src/behaviours.rs b/crates/networking/src/behaviours.rs index ba7f0cbda..b1152dc51 100644 --- a/crates/networking/src/behaviours.rs +++ b/crates/networking/src/behaviours.rs @@ -81,7 +81,12 @@ impl GadgetBehaviour { .build() .unwrap(); - let blueprint_protocol = BlueprintProtocolBehaviour::new(local_key, peer_manager); + let blueprint_protocol = BlueprintProtocolBehaviour::new( + local_key, + instance_secret_key, + instance_public_key, + peer_manager, + ); Self { connection_limits, diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index 26ef89d17..0cfe8d1e4 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -315,7 +315,7 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { ) -> Poll>> { while let Poll::Ready(ev) = self.blueprint_protocol.poll(cx) { match ev { - ToSwarm::GenerateEvent(ev) => match &ev { + ToSwarm::GenerateEvent(ev) => match ev { DerivedBlueprintProtocolBehaviourEvent::RequestResponse( blueprint_protocol_event, ) => self.handle_request_response_event(blueprint_protocol_event), diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index f9103395b..f41686681 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -222,7 +222,7 @@ impl BlueprintProtocolBehaviour { self.verified_peers.contains(peer) } - pub fn handle_gossipsub_event(&mut self, event: &gossipsub::Event) { + pub fn handle_gossipsub_event(&mut self, event: gossipsub::Event) { match event { gossipsub::Event::Message { propagation_source, @@ -230,7 +230,7 @@ impl BlueprintProtocolBehaviour { message, } => { // Only accept gossip from verified peers - if !self.is_peer_verified(propagation_source) { + if !self.is_peer_verified(&propagation_source) { warn!(%propagation_source, "Received gossip from unverified peer"); return; } diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index 131f76f90..e280ef786 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -136,9 +136,7 @@ pub struct NetworkService { impl NetworkService { /// Create a new network service - pub async fn new( - config: NetworkConfig, - ) -> Result<(Self, impl Stream), Error> { + pub async fn new(config: NetworkConfig) -> Result { let NetworkConfig { network_name, instance_secret_key, @@ -192,7 +190,7 @@ impl NetworkService { peer_manager, network_sender, network_receiver: UnboundedReceiverStream::new(network_receiver), - event_sender, + event_sender: event_sender.clone(), event_receiver: UnboundedReceiverStream::new(event_receiver), network_name, bootstrap_peers, @@ -201,7 +199,7 @@ impl NetworkService { // Spawn network task tokio::spawn(service.run(event_sender)); - Ok((service, UnboundedReceiverStream::new(event_receiver))) + Ok(service) } /// Get a sender to send messages to the network service From 9c8f61e1c1573ca4af9169afd56afdb342988768 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Tue, 18 Feb 2025 11:50:58 -0700 Subject: [PATCH 05/52] chore: fix lints --- .../src/blueprint_protocol/behaviour.rs | 17 +- .../src/blueprint_protocol/handler.rs | 20 +- crates/networking/src/discovery/behaviour.rs | 7 +- crates/networking/src/discovery/peers.rs | 34 +- crates/networking/src/service.rs | 52 +- crates/networking/src/tests/mod.rs | 772 ------------------ 6 files changed, 65 insertions(+), 837 deletions(-) diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index 0cfe8d1e4..13b6b8fd5 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -25,7 +25,7 @@ use super::{InstanceMessageRequest, InstanceMessageResponse}; #[derive(NetworkBehaviour)] pub struct DerivedBlueprintProtocolBehaviour { - /// Request/response protocol for direct messaging + /// Request/response protocol for p2p messaging request_response: request_response::cbor::Behaviour, /// Gossipsub for broadcast messaging @@ -35,18 +35,18 @@ pub struct DerivedBlueprintProtocolBehaviour { /// Events emitted by the BlueprintProtocolBehaviour #[derive(Debug)] pub enum BlueprintProtocolEvent { - /// Response received from a peer - Response { - peer: PeerId, - request_id: OutboundRequestId, - response: InstanceMessageResponse, - }, /// Request received from a peer Request { peer: PeerId, request: InstanceMessageRequest, channel: ResponseChannel, }, + /// Response received from a peer + Response { + peer: PeerId, + request_id: OutboundRequestId, + response: InstanceMessageResponse, + }, /// Gossip message received GossipMessage { source: PeerId, @@ -71,8 +71,6 @@ pub struct BlueprintProtocolBehaviour { pub(crate) inbound_handshakes: DashMap, /// Peers with pending outbound handshakes pub(crate) outbound_handshakes: DashMap, - /// Peers that have completed handshake verification - pub(crate) verified_peers: DashSet, /// Active response channels pub(crate) response_channels: DashMap>, @@ -123,7 +121,6 @@ impl BlueprintProtocolBehaviour { instance_secret_key: instance_secret_key.clone(), inbound_handshakes: DashMap::new(), outbound_handshakes: DashMap::new(), - verified_peers: DashSet::new(), response_channels: DashMap::new(), } } diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index f41686681..e0bdb8844 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -37,10 +37,11 @@ impl BlueprintProtocolBehaviour { if let Some(outbound_time) = self.outbound_handshakes.get(&peer) { // If we have an outbound handshake and their peer_id is less than ours, // we should wait for their response instead of responding to their request - if peer < &self.local_peer_id { - debug!(%peer, "Deferring inbound handshake - waiting for outbound response"); - return; - } + // TODO: Fix + // if peer < &self.local_peer_id { + // debug!(%peer, "Deferring inbound handshake - waiting for outbound response"); + // return; + // } // If we have an outbound handshake and their peer_id is greater than ours, // we should handle their request and cancel our outbound attempt self.outbound_handshakes.remove(&peer); @@ -144,7 +145,7 @@ impl BlueprintProtocolBehaviour { .. } => { // Only accept protocol messages from peers we've completed handshakes with - if !self.is_peer_verified(&peer) { + if !self.peer_manager.is_peer_verified(&peer) { warn!(%peer, "Received protocol message from unverified peer"); let response = InstanceMessageResponse::Error { code: 403, @@ -210,18 +211,13 @@ impl BlueprintProtocolBehaviour { self.outbound_handshakes.remove(peer); // Add to verified peers - self.verified_peers.insert(*peer); + self.peer_manager.verify_peer(peer); // Update peer manager self.peer_manager .add_peer_id_to_public_key(peer, public_key); } - /// Check if a peer has completed handshake verification - fn is_peer_verified(&self, peer: &PeerId) -> bool { - self.verified_peers.contains(peer) - } - pub fn handle_gossipsub_event(&mut self, event: gossipsub::Event) { match event { gossipsub::Event::Message { @@ -230,7 +226,7 @@ impl BlueprintProtocolBehaviour { message, } => { // Only accept gossip from verified peers - if !self.is_peer_verified(&propagation_source) { + if !self.peer_manager.is_peer_verified(&propagation_source) { warn!(%propagation_source, "Received gossip from unverified peer"); return; } diff --git a/crates/networking/src/discovery/behaviour.rs b/crates/networking/src/discovery/behaviour.rs index af521d080..cabb5cdc7 100644 --- a/crates/networking/src/discovery/behaviour.rs +++ b/crates/networking/src/discovery/behaviour.rs @@ -80,11 +80,12 @@ pub struct DiscoveryBehaviour { } impl DiscoveryBehaviour { + /// Bootstrap Kademlia network pub fn bootstrap(&mut self) -> Result { - if let Some(kademlia) = self.discovery.kademlia.as_mut() { - kademlia.bootstrap().map_err(|e| e.to_string()) + if let Some(active_kad) = self.discovery.kademlia.as_mut() { + active_kad.bootstrap().map_err(|e| e.to_string()) } else { - Err("Kademlia is not enabled".to_string()) + Err("Kademlia is not activated".to_string()) } } diff --git a/crates/networking/src/discovery/peers.rs b/crates/networking/src/discovery/peers.rs index 01e427127..9921410a6 100644 --- a/crates/networking/src/discovery/peers.rs +++ b/crates/networking/src/discovery/peers.rs @@ -62,14 +62,12 @@ pub enum PeerEvent { pub struct PeerManager { /// Active peers and their information peers: DashMap, - /// Completed handshakes - completed_handshakes: DashSet, + /// Verified peers from completed handshakes + verified_peers: DashSet, /// Handshake keys to peer ids public_keys_to_peer_ids: Arc>>, /// Banned peers with optional expiration time banned_peers: DashMap>, - /// Protected peers that won't be banned - protected_peers: DashSet, /// Event sender for peer updates event_tx: broadcast::Sender, } @@ -80,8 +78,7 @@ impl Default for PeerManager { Self { peers: Default::default(), banned_peers: Default::default(), - protected_peers: Default::default(), - completed_handshakes: Default::default(), + verified_peers: Default::default(), public_keys_to_peer_ids: Arc::new(RwLock::new(BTreeMap::new())), event_tx, } @@ -118,13 +115,18 @@ impl PeerManager { } } + /// Verify a peer + pub fn verify_peer(&self, peer_id: &PeerId) { + self.verified_peers.insert(*peer_id); + } + + /// Check if a peer is verified + pub fn is_peer_verified(&self, peer_id: &PeerId) -> bool { + self.verified_peers.contains(peer_id) + } + /// Ban a peer with optional expiration pub fn ban_peer(&self, peer_id: PeerId, reason: impl Into, duration: Option) { - // Don't ban protected peers - if self.protected_peers.contains(&peer_id) { - return; - } - let expires_at = duration.map(|d| Instant::now() + d); // Remove from active peers @@ -175,16 +177,6 @@ impl PeerManager { } } - /// Protect a peer from being banned - pub fn protect_peer(&self, peer_id: PeerId) { - self.protected_peers.insert(peer_id); - } - - /// Remove protection from a peer - pub fn unprotect_peer(&self, peer_id: &PeerId) { - self.protected_peers.remove(peer_id); - } - /// Get peer information pub fn get_peer_info(&self, peer_id: &PeerId) -> Option { self.peers.get(peer_id).map(|info| info.value().clone()) diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index e280ef786..fffad0e52 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -93,7 +93,7 @@ pub enum NetworkMessage { } /// Configuration for the network service -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct NetworkConfig { /// Network name/namespace pub network_name: String, @@ -196,9 +196,6 @@ impl NetworkService { bootstrap_peers, }; - // Spawn network task - tokio::spawn(service.run(event_sender)); - Ok(service) } @@ -230,13 +227,23 @@ impl NetworkService { tokio::select! { message = network_stream.next() => { match message { - Some(msg) => self.handle_network_message(msg, &event_sender).await, + Some(msg) => match self.handle_network_message(msg, &event_sender).await { + Ok(_) => {} + Err(e) => { + warn!("Failed to handle network message: {}", e); + } + }, None => break, } } event = swarm_stream.next() => { match event { - Some(event) => self.handle_swarm_event(event, &event_sender).await, + Some(event) => match self.handle_swarm_event(event, &event_sender).await { + Ok(_) => {} + Err(e) => { + warn!("Failed to handle swarm event: {}", e); + } + }, None => break, } } @@ -263,20 +270,17 @@ impl NetworkService { source, topic, message, - } => event_sender.send(NetworkEvent::GossipSent { topic, message }), + } => event_sender.send(NetworkEvent::GossipSent { topic, message })?, NetworkMessage::HandshakeRequest { peer, public_key, signature, - } => event_sender.send(NetworkEvent::HandshakeCompleted { peer }), + } => event_sender.send(NetworkEvent::HandshakeCompleted { peer })?, NetworkMessage::HandshakeResponse { peer, public_key, signature, - } => event_sender.send(NetworkEvent::HandshakeCompleted { peer }), - NetworkMessage::HandshakeFailed { peer, reason } => { - event_sender.send(NetworkEvent::HandshakeFailed { peer, reason })? - } + } => event_sender.send(NetworkEvent::HandshakeCompleted { peer })?, } Ok(()) @@ -287,14 +291,16 @@ impl NetworkService { &mut self, event: SwarmEvent, event_sender: &mpsc::UnboundedSender, - ) { + ) -> Result<(), Error> { match event { SwarmEvent::Behaviour(behaviour_event) => { self.handle_behaviour_event(behaviour_event, event_sender) - .await + .await? } _ => {} } + + Ok(()) } /// Handle a behaviour event @@ -302,13 +308,21 @@ impl NetworkService { &mut self, event: GadgetBehaviourEvent, event_sender: &mpsc::UnboundedSender, - ) { + ) -> Result<(), Error> { match event { GadgetBehaviourEvent::ConnectionLimits(_) => {} - GadgetBehaviourEvent::Discovery(discovery_event) => { - self.handle_discovery_event(discovery_event, event_sender) - .await? - } + GadgetBehaviourEvent::Discovery(discovery_event) => match discovery_event { + DiscoveryEvent::Discovery(derived_event) => { + self.handle_discovery_event(derived_event, event_sender) + .await? + } + DiscoveryEvent::PeerConnected(peer) => { + event_sender.send(NetworkEvent::PeerConnected(peer))? + } + DiscoveryEvent::PeerDisconnected(peer) => { + event_sender.send(NetworkEvent::PeerDisconnected(peer))? + } + }, GadgetBehaviourEvent::BlueprintProtocol(blueprint_event) => { self.handle_blueprint_protocol_event(blueprint_event, event_sender) .await? diff --git a/crates/networking/src/tests/mod.rs b/crates/networking/src/tests/mod.rs index ab21b2d3c..8b1378917 100644 --- a/crates/networking/src/tests/mod.rs +++ b/crates/networking/src/tests/mod.rs @@ -1,773 +1 @@ -use self::std::time::Duration; -use super::*; -use crate::{ - gossip::GossipHandle, - networking::{Network, NetworkMultiplexer, StreamKey}, - types::IdentifierInfo, -}; -use futures::{stream, StreamExt}; -use gadget_crypto::hashing::blake3_256; -use gadget_crypto::KeyType; -use gadget_logging::setup_log; -use gadget_std::collections::BTreeMap; -use gadget_std::sync::LazyLock; -use serde::{Deserialize, Serialize}; -use tokio::time::sleep; -const TOPIC: &str = "/gadget/test/1.0.0"; - -fn deserialize<'a, T>(data: &'a [u8]) -> Result -where - T: Deserialize<'a>, -{ - bincode::deserialize(data).map_err(|err| Error::Other(err.to_string())) -} - -#[derive(Serialize, Deserialize, Debug)] -struct StressTestPayload { - value: u64, -} - -#[derive(Debug, Serialize, Deserialize, Clone)] -enum Msg { - Round1(Round1Msg), - Round2(Round2Msg), - Round3(Round3Msg), -} - -#[derive(Debug, Serialize, Deserialize, Clone)] -struct Round1Msg { - pub power: u16, - pub hitpoints: u16, - pub armor: u16, - pub name: String, -} - -#[derive(Debug, Serialize, Deserialize, Clone)] -struct Round2Msg { - pub x: u16, - pub y: u16, - pub z: u16, -} - -#[derive(Debug, Serialize, Deserialize, Clone)] -struct Round3Msg { - rotation: u16, - velocity: (u16, u16, u16), -} - -async fn wait_for_nodes_connected(nodes: &[GossipHandle]) { - let node_count = nodes.len(); - - // wait for the nodes to connect to each other - let max_retries = 10 * node_count; - let mut retry = 0; - loop { - gadget_logging::debug!(%node_count, %max_retries, %retry, "Checking if all nodes are connected to each other"); - let connected = nodes - .iter() - .map(super::super::gossip::GossipHandle::connected_peers) - .collect::>(); - - let all_connected = connected - .iter() - .enumerate() - .inspect(|(node, peers)| { - gadget_logging::debug!("Node {node} has {peers} connected peers"); - }) - .all(|(_, &peers)| peers >= node_count - 1); - if all_connected { - gadget_logging::debug!("All nodes are connected to each other"); - return; - } - sleep(Duration::from_millis(300)).await; - retry += 1; - assert!( - retry <= max_retries, - "Failed to connect all nodes to each other" - ); - } -} - -#[tokio::test(flavor = "multi_thread")] -#[allow(clippy::cast_possible_truncation)] -async fn test_p2p() { - setup_log(); - let nodes = stream::iter(0..*NODE_COUNT) - .map(|_| node()) - .collect::>() - .await; - - wait_for_nodes_connected(&nodes).await; - - let mut mapping = BTreeMap::new(); - for (i, node) in nodes.iter().enumerate() { - mapping.insert(i as u16, node.my_id); - } - - let mut tasks = Vec::new(); - for (i, node) in nodes.into_iter().enumerate() { - let task = tokio::spawn(run_protocol(node, i as u16, mapping.clone())); - tasks.push(task); - } - // Wait for all tasks to finish - let results = futures::future::try_join_all(tasks) - .await - .expect("Failed to run protocol"); - // Assert that all are okay. - assert!( - results.iter().all(std::result::Result::is_ok), - "Some nodes failed to run protocol" - ); -} - -#[allow(clippy::too_many_lines, clippy::cast_possible_truncation)] -async fn run_protocol( - node: N, - i: u16, - mapping: BTreeMap, -) -> Result<(), crate::error::Error> { - let task_hash = [0u8; 32]; - // Safety note: We should be passed a NetworkMultiplexer, and all uses of the N: Network - // used throughout the program must also use the multiplexer to prevent mixed messages. - let multiplexer = NetworkMultiplexer::new(node); - - let round1_network = multiplexer.multiplex(StreamKey { - task_hash, // To differentiate between different instances of a running program (i.e., a task) - round_id: 0, // To differentiate between different subsets of a running task - }); - - let round2_network = multiplexer.multiplex(StreamKey { - task_hash, // To differentiate between different instances of a running program (i.e., a task) - round_id: 1, // To differentiate between different subsets of a running task - }); - - let round3_network = multiplexer.multiplex(StreamKey { - task_hash, // To differentiate between different instances of a running program (i.e., a task) - round_id: 2, // To differentiate between different subsets of a running task - }); - - //let (round1_tx, round1_rx) = node. - // Round 1 (broadcast) - let msg = { - let round = Round1Msg { - power: i * 100, - hitpoints: (i + 1) * 50, - armor: i + 2, - name: format!("Player {}", i), - }; - round1_network.build_protocol_message( - IdentifierInfo { - message_id: 0, - round_id: 0, - }, - i, - None, - &Msg::Round1(round), - None, - ) - }; - - gadget_logging::debug!("Broadcast Message"); - round1_network - .send(msg) - .map_err(|_| crate::error::Error::Other("Failed to send message".into()))?; - - // Wait for all other nodes to send their messages - let mut msgs = BTreeMap::new(); - while let Some(msg) = round1_network.recv().await { - let m = deserialize::(&msg.payload).unwrap(); - gadget_logging::debug!(from = %msg.sender.user_id, ?m, "Received message"); - // Expecting Round1 message - assert!( - matches!(m, Msg::Round1(_)), - "Expected Round1 message but got {:?} from node {}", - m, - msg.sender.user_id, - ); - let old = msgs.insert(msg.sender.user_id, m); - assert!( - old.is_none(), - "Duplicate message from node {}", - msg.sender.user_id, - ); - // Break if all messages are received - if msgs.len() == *NODE_COUNT - 1 { - break; - } - } - gadget_logging::debug!("Done r1 w/ {i}"); - - // Round 2 (P2P) - let msgs = (0..*NODE_COUNT) - .map(|r| r as u16) - .filter(|&j| j != i) - .map(|j| { - let peer_pk = mapping.get(&j).copied().unwrap(); - round2_network.build_protocol_message( - IdentifierInfo { - message_id: 0, - round_id: 0, - }, - i, - Some(j), - &Msg::Round2(Round2Msg { - x: i * 10, - y: (i + 1) * 20, - z: i + 2, - }), - Some(peer_pk), - ) - }) - .collect::>(); - for msg in msgs { - let to = msg - .recipient - .map(|r| r.user_id) - .expect("Recipient should be present for P2P message. This is a bug in the test code"); - gadget_logging::debug!(%to, "Send P2P Message"); - round2_network.send(msg)?; - } - - // Wait for all other nodes to send their messages - let mut msgs = BTreeMap::new(); - while let Some(msg) = round2_network.recv().await { - let m = deserialize::(&msg.payload).unwrap(); - gadget_logging::info!( - "[Node {}] Received message from {} | Intended Recipient: {}", - i, - msg.sender.user_id, - msg.recipient - .as_ref() - .map_or_else(|| "Broadcast".into(), |r| r.user_id.to_string()) - ); - // Expecting Round2 message - assert!( - matches!(m, Msg::Round2(_)), - "Expected Round2 message but got {:?} from node {}", - m, - msg.sender.user_id, - ); - let old = msgs.insert(msg.sender.user_id, m); - assert!( - old.is_none(), - "Duplicate message from node {}", - msg.sender.user_id, - ); - // Break if all messages are received - if msgs.len() == *NODE_COUNT - 1 { - break; - } - } - gadget_logging::debug!("Done r2 w/ {i}"); - - // Round 3 (broadcast) - - let msg = { - let round = Round3Msg { - rotation: i * 30, - velocity: (i + 1, i + 2, i + 3), - }; - round3_network.build_protocol_message( - IdentifierInfo { - message_id: 0, - round_id: 0, - }, - i, - None, - &Msg::Round3(round), - None, - ) - }; - - gadget_logging::debug!("Broadcast Message"); - round3_network.send(msg)?; - - // Wait for all other nodes to send their messages - let mut msgs = BTreeMap::new(); - while let Some(msg) = round3_network.recv().await { - let m = deserialize::(&msg.payload).unwrap(); - gadget_logging::debug!(from = %msg.sender.user_id, ?m, "Received message"); - // Expecting Round3 message - assert!( - matches!(m, Msg::Round3(_)), - "Expected Round3 message but got {:?} from node {}", - m, - msg.sender.user_id, - ); - let old = msgs.insert(msg.sender.user_id, m); - assert!( - old.is_none(), - "Duplicate message from node {}", - msg.sender.user_id, - ); - // Break if all messages are received - if msgs.len() == *NODE_COUNT - 1 { - break; - } - } - gadget_logging::debug!("Done r3 w/ {i}"); - - gadget_logging::info!(node = i, "Protocol completed"); - - Ok(()) -} - -fn node_with_id() -> (GossipHandle, crate::key_types::GossipMsgKeyPair) { - let identity = libp2p::identity::Keypair::generate_ed25519(); - let crypto_key = crate::key_types::Curve::generate_with_seed(None).unwrap(); - let bind_port = 0; - let handle = crate::setup::start_p2p_network(crate::setup::NetworkConfig::new_service_network( - identity, - crypto_key.clone(), - Vec::default(), - bind_port, - TOPIC, - )) - .unwrap(); - - (handle, crypto_key) -} - -fn node() -> GossipHandle { - node_with_id().0 -} - -static NODE_COUNT: LazyLock = - LazyLock::new(|| std::env::var("IN_CI").map_or_else(|_| 10, |_| 2)); -#[allow(dead_code)] -static MESSAGE_COUNT: LazyLock = - LazyLock::new(|| std::env::var("IN_CI").map_or_else(|_| 10, |_| 100)); - -#[tokio::test(flavor = "multi_thread")] -#[allow(clippy::cast_possible_truncation)] -async fn test_stress_test_multiplexer() { - setup_log(); - gadget_logging::info!("Starting test_stress_test_multiplexer"); - - let (network0, network1) = get_networks().await; - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let stream_key = StreamKey { - task_hash: blake3_256(&[1]), - round_id: 0, - }; - - let _subnetwork0 = multiplexer0.multiplex(stream_key); - let _subnetwork1 = multiplexer1.multiplex(stream_key); - - // Create a channel for forwarding - let (forward_tx, mut forward_rx) = tokio::sync::mpsc::unbounded_channel(); - - // Create a subnetwork with forwarding - let subnetwork0 = multiplexer0.multiplex(stream_key); - let subnetwork1 = multiplexer1.multiplex_with_forwarding(stream_key, forward_tx); - - let payload = StressTestPayload { value: 42 }; - let msg = subnetwork0.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), - &payload, - Some(subnetwork1.public_id()), - ); - - gadget_logging::info!("Sending message from subnetwork0"); - subnetwork0.send(msg.clone()).unwrap(); - - // Message should be forwarded to the forward_rx channel - let forwarded_msg = forward_rx.recv().await.unwrap(); - let received: StressTestPayload = deserialize(&forwarded_msg.payload).unwrap(); - assert_eq!(received.value, payload.value); -} - -#[tokio::test(flavor = "multi_thread")] -async fn test_nested_multiplexer() { - setup_log(); - gadget_logging::info!("Starting test_nested_multiplexer"); - let (network0, network1) = get_networks().await; - - nested_multiplex(0, 10, network0, network1).await; -} - -async fn get_networks() -> (GossipHandle, GossipHandle) { - let network0 = node(); - let network1 = node(); - - let mut gossip_networks = vec![network0, network1]; - - wait_for_nodes_connected(&gossip_networks).await; - - (gossip_networks.remove(0), gossip_networks.remove(0)) -} - -async fn nested_multiplex( - cur_depth: usize, - max_depth: usize, - network0: N, - network1: N, -) { - gadget_logging::info!("At nested depth = {cur_depth}/{max_depth}"); - - if cur_depth == max_depth { - return; - } - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let stream_key = StreamKey { - #[allow(clippy::cast_possible_truncation)] - task_hash: blake3_256(&[(cur_depth % 255) as u8]), - round_id: 0, - }; - - let subnetwork0 = multiplexer0.multiplex(stream_key); - let subnetwork1 = multiplexer1.multiplex(stream_key); - let subnetwork1_id = subnetwork1.public_id(); - - // Send a message in the subnetwork0 to subnetwork1 and vice versa, assert values of message - let payload = StressTestPayload { value: 42 }; - let msg = subnetwork0.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), - &payload, - Some(subnetwork1_id), - ); - - gadget_logging::info!("Sending message from subnetwork0"); - subnetwork0.send(msg.clone()).unwrap(); - - // Receive message - let received_msg = subnetwork1.recv().await.unwrap(); - let received: StressTestPayload = deserialize(&received_msg.payload).unwrap(); - assert_eq!(received.value, payload.value); - - let msg = subnetwork1.build_protocol_message( - IdentifierInfo::default(), - 1, - Some(0), - &payload, - Some(subnetwork0.public_id()), - ); - - gadget_logging::info!("Sending message from subnetwork1"); - subnetwork1.send(msg.clone()).unwrap(); - - // Receive message - let received_msg = subnetwork0.recv().await.unwrap(); - let received: StressTestPayload = deserialize(&received_msg.payload).unwrap(); - assert_eq!(received.value, payload.value); - tracing::info!("Done nested depth = {cur_depth}/{max_depth}"); - - Box::pin(nested_multiplex( - cur_depth + 1, - max_depth, - subnetwork0, - subnetwork1, - )) - .await; -} - -#[tokio::test(flavor = "multi_thread")] -async fn test_closed_channel_handling() { - setup_log(); - let (network0, network1) = get_networks().await; - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let stream_key = StreamKey { - task_hash: blake3_256(&[1]), - round_id: 0, - }; - - let subnetwork0 = multiplexer0.multiplex(stream_key); - // Drop subnetwork1's receiver to simulate closed channel - let subnetwork1 = multiplexer1.multiplex(stream_key); - drop(subnetwork1); - - let payload = StressTestPayload { value: 42 }; - let msg = - subnetwork0.build_protocol_message(IdentifierInfo::default(), 0, None, &payload, None); - - // Sending to a closed channel should return an error - assert!(subnetwork0.send(msg).is_ok()); // Changed to ok() since the message will be sent but not received -} - -#[tokio::test(flavor = "multi_thread")] -async fn test_empty_payload() { - setup_log(); - let (network0, network1) = get_networks().await; - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let stream_key = StreamKey { - task_hash: blake3_256(&[1]), - round_id: 0, - }; - - let subnetwork0 = multiplexer0.multiplex(stream_key); - let subnetwork1 = multiplexer1.multiplex(stream_key); - - // Test empty payload - let empty_payload = StressTestPayload { value: 0 }; - let msg = subnetwork0.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), - &empty_payload, - Some(subnetwork1.public_id()), - ); - - gadget_logging::info!("Sending message from subnetwork0"); - subnetwork0.send(msg).unwrap(); - - // Receive message - let received_msg = subnetwork1.recv().await.unwrap(); - let received: StressTestPayload = deserialize(&received_msg.payload).unwrap(); - assert_eq!(received.value, empty_payload.value); -} - -#[tokio::test(flavor = "multi_thread")] -#[allow(clippy::cast_possible_truncation)] -async fn test_concurrent_messaging() { - setup_log(); - let (network0, network1) = get_networks().await; - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let mut send_handles = Vec::new(); - let mut receive_handles = Vec::new(); - - // Create multiple messages to send concurrently - let message_count = 10; - - // Spawn tasks to send messages - for i in 0..message_count { - let stream_key = StreamKey { - task_hash: blake3_256(&[i]), - round_id: 0, - }; - - let subnetwork0 = multiplexer0.multiplex(stream_key); - let subnetwork1 = multiplexer1.multiplex(stream_key); - let subnetwork1_id = subnetwork1.public_id(); - - let i_u64: u64 = i.into(); - let payload = StressTestPayload { value: i_u64 }; - let send_subnetwork0 = subnetwork0; - let handle = tokio::spawn(async move { - let msg = send_subnetwork0.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), - &payload, - Some(subnetwork1_id), - ); - send_subnetwork0.send(msg).unwrap(); - }); - - send_handles.push(handle); - - // Spawn tasks to receive messages - let handle = tokio::spawn(async move { - let msg = subnetwork1.recv().await.unwrap(); - let received: StressTestPayload = deserialize(&msg.payload).unwrap(); - received.value as u8 // Return the payload value for verification - }); - - receive_handles.push(handle); - } - - // Wait for all sends to complete - for handle in send_handles { - handle.await.unwrap(); - } - - // Wait for all receives and verify we got all messages - let mut received_values = Vec::new(); - for handle in receive_handles { - received_values.push(handle.await.unwrap()); - } - - received_values.sort_unstable(); - assert_eq!(received_values.len(), message_count as usize); - for i in 0..message_count { - assert_eq!(received_values[i as usize], i); - } -} - -#[tokio::test(flavor = "multi_thread")] -#[allow(clippy::cast_possible_truncation)] -async fn test_message_ordering() { - setup_log(); - let (network0, network1) = get_networks().await; - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let stream_key = StreamKey { - task_hash: blake3_256(&[1]), - round_id: 0, - }; - - let subnetwork0 = multiplexer0.multiplex(stream_key); - let subnetwork1 = multiplexer1.multiplex(stream_key); - - // Send messages with sequential sequence numbers - let message_count = 10; - for i in 0..message_count { - let payload = StressTestPayload { value: i }; - let msg = subnetwork0.build_protocol_message( - IdentifierInfo { - message_id: i, - ..Default::default() - }, - 0, - Some(1), - &payload, - Some(subnetwork1.public_id()), - ); - subnetwork0.send(msg).unwrap(); - } - - // Verify messages are received in order - let mut last_seq = 0; - for _ in 0..message_count { - let msg = subnetwork1.recv().await.unwrap(); - assert!( - msg.identifier_info.message_id >= last_seq, - "Messages should be received in order or equal to last sequence number" - ); - last_seq = msg.identifier_info.message_id; - } -} - -#[tokio::test(flavor = "multi_thread")] -async fn test_network_id_handling() { - setup_log(); - let (network0, network1) = get_networks().await; - let _network0_id = network0.public_id(); - let network1_id = network1.public_id(); - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let stream_key = StreamKey { - task_hash: blake3_256(&[1]), - round_id: 0, - }; - - let subnetwork0 = multiplexer0.multiplex(stream_key); - let subnetwork1 = multiplexer1.multiplex(stream_key); - - // Test sending with correct network ID - let payload = StressTestPayload { value: 42 }; - let msg = subnetwork0.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), - &payload, - Some(network1_id), - ); - gadget_logging::info!("Sending message from subnetwork0"); - subnetwork0.send(msg.clone()).unwrap(); - - // Receive message - let received_msg = subnetwork1.recv().await.unwrap(); - let received: StressTestPayload = deserialize(&received_msg.payload).unwrap(); - assert_eq!(received.value, payload.value); - - // Test sending with wrong network ID - let wrong_key = crate::key_types::Curve::generate_with_seed(None).unwrap(); - let msg = subnetwork0.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), - &payload, - Some(wrong_key.public()), - ); - gadget_logging::info!("Sending message from subnetwork0"); - subnetwork0.send(msg).unwrap(); - - // Message with wrong network ID should not be received - let timeout = tokio::time::sleep(tokio::time::Duration::from_millis(100)); - tokio::select! { - () = timeout => (), - _ = subnetwork1.recv() => panic!("Should not receive message with wrong network ID"), - } -} - -#[tokio::test(flavor = "multi_thread")] -async fn test_stream_isolation() { - setup_log(); - let (network0, network1) = get_networks().await; - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - // Create two different stream keys - let stream_key1 = StreamKey { - task_hash: blake3_256(&[1]), - round_id: 0, - }; - let stream_key2 = StreamKey { - task_hash: blake3_256(&[2]), - round_id: 0, - }; - - let subnetwork0_stream1 = multiplexer0.multiplex(stream_key1); - let subnetwork0_stream2 = multiplexer0.multiplex(stream_key2); - let subnetwork1_stream1 = multiplexer1.multiplex(stream_key1); - let subnetwork1_stream2 = multiplexer1.multiplex(stream_key2); - - // Send messages on both streams - let payload1 = StressTestPayload { value: 1 }; - let payload2 = StressTestPayload { value: 2 }; - - let msg1 = subnetwork0_stream1.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), // Send to node 1 - &payload1, - Some(subnetwork1_stream1.public_id()), - ); - let msg2 = subnetwork0_stream2.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), // Send to node 1 - &payload2, - Some(subnetwork1_stream2.public_id()), - ); - - gadget_logging::info!("Sending message from subnetwork0_stream1"); - subnetwork0_stream1.send(msg1.clone()).unwrap(); - gadget_logging::info!("Sending message from subnetwork0_stream2"); - subnetwork0_stream2.send(msg2.clone()).unwrap(); - - // Verify messages are received on correct streams - gadget_logging::info!("Waiting for message on subnetwork1_stream1"); - let received_msg1 = subnetwork1_stream1.recv().await.unwrap(); - gadget_logging::info!("Waiting for message on subnetwork1_stream2"); - let received_msg2 = subnetwork1_stream2.recv().await.unwrap(); - - let received1: StressTestPayload = deserialize(&received_msg1.payload).unwrap(); - let received2: StressTestPayload = deserialize(&received_msg2.payload).unwrap(); - - assert_eq!(received1.value, payload1.value); - assert_eq!(received2.value, payload2.value); - - // Verify no cross-stream message leakage - let timeout = tokio::time::sleep(tokio::time::Duration::from_millis(100)); - tokio::select! { - () = timeout => (), - _ = subnetwork1_stream1.recv() => panic!("Should not receive more messages on stream 1"), - _ = subnetwork1_stream2.recv() => panic!("Should not receive more messages on stream 2"), - } -} From 7b39059aee36ef6d7f3857dca178e4ae8f2eec0e Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Tue, 18 Feb 2025 11:56:42 -0700 Subject: [PATCH 06/52] chore: lints --- crates/networking/src/service.rs | 377 +++++++++++++++++-------------- 1 file changed, 212 insertions(+), 165 deletions(-) diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index fffad0e52..83f5ff4de 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -227,7 +227,14 @@ impl NetworkService { tokio::select! { message = network_stream.next() => { match message { - Some(msg) => match self.handle_network_message(msg, &event_sender).await { + Some(msg) => match handle_network_message( + swarm_stream.get_mut(), + msg, + &self.peer_manager, + &self.event_sender, + ) + .await + { Ok(_) => {} Err(e) => { warn!("Failed to handle network message: {}", e); @@ -236,192 +243,232 @@ impl NetworkService { None => break, } } - event = swarm_stream.next() => { - match event { - Some(event) => match self.handle_swarm_event(event, &event_sender).await { - Ok(_) => {} - Err(e) => { - warn!("Failed to handle swarm event: {}", e); - } - }, - None => break, - } - } + swarm_event = swarm_stream.next() => match swarm_event { + // outbound events + Some(SwarmEvent::Behaviour(event)) => { + handle_behaviour_event( + swarm_stream.get_mut(), + &self.peer_manager, + event, + &self.event_sender, + &self.network_sender, + ) + .await; + }, + None => { break; }, + _ => { }, + }, } } info!("Network service stopped"); } +} - /// Handle a network message - async fn handle_network_message( - &mut self, - msg: NetworkMessage, - event_sender: &mpsc::UnboundedSender, - ) -> Result<(), Error> { - match msg { - NetworkMessage::InstanceRequest { peer, request } => { - event_sender.send(NetworkEvent::InstanceRequestOutbound { peer, request })? - } - NetworkMessage::InstanceResponse { peer, response } => { - event_sender.send(NetworkEvent::InstanceResponseOutbound { peer, response })? - } - NetworkMessage::GossipMessage { - source, - topic, - message, - } => event_sender.send(NetworkEvent::GossipSent { topic, message })?, - NetworkMessage::HandshakeRequest { - peer, - public_key, - signature, - } => event_sender.send(NetworkEvent::HandshakeCompleted { peer })?, - NetworkMessage::HandshakeResponse { - peer, - public_key, - signature, - } => event_sender.send(NetworkEvent::HandshakeCompleted { peer })?, +/// Handle a swarm event +async fn handle_swarm_event( + swarm: &mut Swarm, + peer_manager: &Arc, + event: SwarmEvent, + event_sender: &mpsc::UnboundedSender, + network_sender: &mpsc::UnboundedSender, +) -> Result<(), Error> { + match event { + SwarmEvent::Behaviour(behaviour_event) => { + handle_behaviour_event( + swarm, + peer_manager, + behaviour_event, + event_sender, + network_sender, + ) + .await? } - - Ok(()) + _ => {} } - /// Handle a swarm event - async fn handle_swarm_event( - &mut self, - event: SwarmEvent, - event_sender: &mpsc::UnboundedSender, - ) -> Result<(), Error> { - match event { - SwarmEvent::Behaviour(behaviour_event) => { - self.handle_behaviour_event(behaviour_event, event_sender) - .await? - } - _ => {} - } - - Ok(()) - } + Ok(()) +} - /// Handle a behaviour event - async fn handle_behaviour_event( - &mut self, - event: GadgetBehaviourEvent, - event_sender: &mpsc::UnboundedSender, - ) -> Result<(), Error> { - match event { - GadgetBehaviourEvent::ConnectionLimits(_) => {} - GadgetBehaviourEvent::Discovery(discovery_event) => match discovery_event { - DiscoveryEvent::Discovery(derived_event) => { - self.handle_discovery_event(derived_event, event_sender) - .await? - } - DiscoveryEvent::PeerConnected(peer) => { - event_sender.send(NetworkEvent::PeerConnected(peer))? - } - DiscoveryEvent::PeerDisconnected(peer) => { - event_sender.send(NetworkEvent::PeerDisconnected(peer))? - } - }, - GadgetBehaviourEvent::BlueprintProtocol(blueprint_event) => { - self.handle_blueprint_protocol_event(blueprint_event, event_sender) - .await? +/// Handle a behaviour event +async fn handle_behaviour_event( + swarm: &mut Swarm, + peer_manager: &Arc, + event: GadgetBehaviourEvent, + event_sender: &mpsc::UnboundedSender, + network_sender: &mpsc::UnboundedSender, +) -> Result<(), Error> { + match event { + GadgetBehaviourEvent::ConnectionLimits(_) => {} + GadgetBehaviourEvent::Discovery(discovery_event) => match discovery_event { + DiscoveryEvent::Discovery(derived_event) => { + handle_discovery_event( + swarm, + peer_manager, + derived_event, + event_sender, + network_sender, + ) + .await? + } + DiscoveryEvent::PeerConnected(peer) => { + event_sender.send(NetworkEvent::PeerConnected(peer))? } - GadgetBehaviourEvent::Ping(ping_event) => { - self.handle_ping_event(ping_event, event_sender).await? + DiscoveryEvent::PeerDisconnected(peer) => { + event_sender.send(NetworkEvent::PeerDisconnected(peer))? } + }, + GadgetBehaviourEvent::BlueprintProtocol(blueprint_event) => { + handle_blueprint_protocol_event( + swarm, + peer_manager, + blueprint_event, + event_sender, + network_sender, + ) + .await? + } + GadgetBehaviourEvent::Ping(ping_event) => { + handle_ping_event( + swarm, + peer_manager, + ping_event, + event_sender, + network_sender, + ) + .await? } - - Ok(()) } - /// Handle a discovery event - async fn handle_discovery_event( - &mut self, - event: Box, - event_sender: &mpsc::UnboundedSender, - ) -> Result<(), Error> { - let event = event.as_ref(); - match event { - DerivedDiscoveryBehaviourEvent::Kademlia(event) => { - // Handle Kademlia discovery events - } - DerivedDiscoveryBehaviourEvent::Mdns(event) => { - // Handle mDNS discovery events - } - DerivedDiscoveryBehaviourEvent::Identify(event) => { - // Handle Identify protocol events - } - DerivedDiscoveryBehaviourEvent::Autonat(event) => { - // Handle AutoNAT events - } - DerivedDiscoveryBehaviourEvent::Upnp(event) => { - // Handle UPnP events - } - DerivedDiscoveryBehaviourEvent::Relay(event) => { - // Handle relay events - } - } + Ok(()) +} - Ok(()) +/// Handle a discovery event +async fn handle_discovery_event( + swarm: &mut Swarm, + peer_manager: &Arc, + event: Box, + event_sender: &mpsc::UnboundedSender, + network_sender: &mpsc::UnboundedSender, +) -> Result<(), Error> { + let event = event.as_ref(); + match event { + DerivedDiscoveryBehaviourEvent::Kademlia(event) => { + // Handle Kademlia discovery events + } + DerivedDiscoveryBehaviourEvent::Mdns(event) => { + // Handle mDNS discovery events + } + DerivedDiscoveryBehaviourEvent::Identify(event) => { + // Handle Identify protocol events + } + DerivedDiscoveryBehaviourEvent::Autonat(event) => { + // Handle AutoNAT events + } + DerivedDiscoveryBehaviourEvent::Upnp(event) => { + // Handle UPnP events + } + DerivedDiscoveryBehaviourEvent::Relay(event) => { + // Handle relay events + } } - /// Handle a blueprint event - async fn handle_blueprint_protocol_event( - &mut self, - event: BlueprintProtocolEvent, - event_sender: &mpsc::UnboundedSender, - ) -> Result<(), Error> { - match event { - BlueprintProtocolEvent::Request { - peer, - request, - channel, - } => event_sender.send(NetworkEvent::InstanceRequestInbound { peer, request })?, - BlueprintProtocolEvent::Response { - peer, - response, - request_id, - } => event_sender.send(NetworkEvent::InstanceResponseInbound { peer, response })?, - BlueprintProtocolEvent::GossipMessage { - source, - topic, - message, - } => event_sender.send(NetworkEvent::GossipReceived { - source, - topic: topic.to_string(), - message, - })?, - } + Ok(()) +} - Ok(()) +/// Handle a blueprint event +async fn handle_blueprint_protocol_event( + swarm: &mut Swarm, + peer_manager: &Arc, + event: BlueprintProtocolEvent, + event_sender: &mpsc::UnboundedSender, + network_sender: &mpsc::UnboundedSender, +) -> Result<(), Error> { + match event { + BlueprintProtocolEvent::Request { + peer, + request, + channel, + } => event_sender.send(NetworkEvent::InstanceRequestInbound { peer, request })?, + BlueprintProtocolEvent::Response { + peer, + response, + request_id, + } => event_sender.send(NetworkEvent::InstanceResponseInbound { peer, response })?, + BlueprintProtocolEvent::GossipMessage { + source, + topic, + message, + } => event_sender.send(NetworkEvent::GossipReceived { + source, + topic: topic.to_string(), + message, + })?, } - /// Handle a ping event - async fn handle_ping_event( - &mut self, - event: ping::Event, - event_sender: &mpsc::UnboundedSender, - ) -> Result<(), Error> { - match event.result { - Ok(rtt) => { - trace!( - "PingSuccess::Ping rtt to {} is {} ms", - event.peer, - rtt.as_millis() - ); - } - Err(ping::Failure::Unsupported) => { - debug!(peer=%event.peer, "Ping protocol unsupported"); - } - Err(ping::Failure::Timeout) => { - debug!("Ping timeout: {}", event.peer); - } - Err(ping::Failure::Other { error }) => { - debug!("Ping failure: {error}"); - } + Ok(()) +} + +/// Handle a ping event +async fn handle_ping_event( + swarm: &mut Swarm, + peer_manager: &Arc, + event: ping::Event, + event_sender: &mpsc::UnboundedSender, + network_sender: &mpsc::UnboundedSender, +) -> Result<(), Error> { + match event.result { + Ok(rtt) => { + trace!( + "PingSuccess::Ping rtt to {} is {} ms", + event.peer, + rtt.as_millis() + ); + } + Err(ping::Failure::Unsupported) => { + debug!(peer=%event.peer, "Ping protocol unsupported"); + } + Err(ping::Failure::Timeout) => { + debug!("Ping timeout: {}", event.peer); + } + Err(ping::Failure::Other { error }) => { + debug!("Ping failure: {error}"); } + } + + Ok(()) +} - Ok(()) +/// Handle a network message +async fn handle_network_message( + swarm: &mut Swarm, + msg: NetworkMessage, + peer_manager: &Arc, + event_sender: &mpsc::UnboundedSender, +) -> Result<(), Error> { + match msg { + NetworkMessage::InstanceRequest { peer, request } => { + event_sender.send(NetworkEvent::InstanceRequestOutbound { peer, request })? + } + NetworkMessage::InstanceResponse { peer, response } => { + event_sender.send(NetworkEvent::InstanceResponseOutbound { peer, response })? + } + NetworkMessage::GossipMessage { + source, + topic, + message, + } => event_sender.send(NetworkEvent::GossipSent { topic, message })?, + NetworkMessage::HandshakeRequest { + peer, + public_key, + signature, + } => event_sender.send(NetworkEvent::HandshakeCompleted { peer })?, + NetworkMessage::HandshakeResponse { + peer, + public_key, + signature, + } => event_sender.send(NetworkEvent::HandshakeCompleted { peer })?, } + + Ok(()) } From de082453a38f67f87d070934a29c01ba38f17288 Mon Sep 17 00:00:00 2001 From: Serial <69764315+Serial-ATA@users.noreply.github.com> Date: Tue, 18 Feb 2025 14:24:46 -0500 Subject: [PATCH 07/52] fix: get building --- crates/networking/Cargo.toml | 2 +- .../src/blueprint_protocol/behaviour.rs | 5 ++-- .../src/blueprint_protocol/handler.rs | 1 - crates/networking/src/error.rs | 4 ++- crates/networking/src/lib.rs | 28 +++++++++++++++++++ 5 files changed, 34 insertions(+), 6 deletions(-) diff --git a/crates/networking/Cargo.toml b/crates/networking/Cargo.toml index 26e5e4993..5e245f343 100644 --- a/crates/networking/Cargo.toml +++ b/crates/networking/Cargo.toml @@ -19,7 +19,7 @@ blake3 = { workspace = true } dashmap = { workspace = true } libp2p = { workspace = true } tokio = { workspace = true, features = ["macros"] } -tokio-stream = { workspace = true } +tokio-stream = { workspace = true, features = ["time"] } futures = { workspace = true } tracing = { workspace = true } bincode = { workspace = true } diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index 13b6b8fd5..ef39062c0 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -1,4 +1,4 @@ -use crate::{Curve, InstanceMsgKeyPair, InstanceMsgPublicKey, InstanceSignedMsgSignature}; +use crate::{Curve, InstanceMsgKeyPair, InstanceMsgPublicKey, InstanceSignedMsgSignature, KeySignExt}; use dashmap::{DashMap, DashSet}; use gadget_crypto::{hashing::blake3_256, KeyType}; use gadget_logging::{debug, trace, warn}; @@ -129,8 +129,7 @@ impl BlueprintProtocolBehaviour { pub(crate) fn sign_handshake(&self, peer: &PeerId) -> InstanceSignedMsgSignature { let msg = peer.to_bytes(); let msg_hash = blake3_256(&msg); - let signature = self.instance_secret_key.sign_prehashed(&msg_hash); - InstanceSignedMsgSignature(signature) + self.instance_secret_key.sign_prehash(&msg_hash) } /// Send a request to a peer diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index e0bdb8844..40ab5b2a0 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -1,6 +1,5 @@ use std::time::{Duration, Instant}; -use gadget_crypto::tangle_pair_signer::sp_core::offchain::Duration as GadgetDuration; use libp2p::{gossipsub, request_response, PeerId}; use tracing::{debug, warn}; diff --git a/crates/networking/src/error.rs b/crates/networking/src/error.rs index 3dc77e3ed..e2e4b056f 100644 --- a/crates/networking/src/error.rs +++ b/crates/networking/src/error.rs @@ -1,4 +1,6 @@ -use crate::{service::NetworkMessage, NetworkEvent}; +use crate::NetworkEvent; + +pub type Result = core::result::Result; #[derive(Debug, thiserror::Error)] pub enum Error { diff --git a/crates/networking/src/lib.rs b/crates/networking/src/lib.rs index e2bfe502e..a085802ea 100644 --- a/crates/networking/src/lib.rs +++ b/crates/networking/src/lib.rs @@ -24,6 +24,12 @@ pub mod key_types { SpEcdsa as Curve, SpEcdsaPair as InstanceMsgKeyPair, SpEcdsaPublic as InstanceMsgPublicKey, SpEcdsaSignature as InstanceSignedMsgSignature, }; + + impl super::KeySignExt for InstanceMsgKeyPair { + fn sign_prehash(&self, prehash: &[u8; 32]) -> InstanceSignedMsgSignature { + InstanceSignedMsgSignature(self.0.sign_prehashed(prehash)) + } + } } #[cfg(all( @@ -36,6 +42,12 @@ pub mod key_types { SpSr25519 as Curve, SpSr25519Pair as InstanceMsgKeyPair, SpSr25519Public as InstanceMsgPublicKey, SpSr25519Signature as InstanceSignedMsgSignature, }; + + impl super::KeySignExt for InstanceMsgKeyPair { + fn sign_prehash(&self, prehash: &[u8; 32]) -> InstanceSignedMsgSignature { + InstanceSignedMsgSignature(self.0.sign_prehashed(prehash)) + } + } } #[cfg(all( @@ -48,6 +60,12 @@ pub mod key_types { SpEd25519 as Curve, SpEd25519Pair as InstanceMsgKeyPair, SpEd25519Public as InstanceMsgPublicKey, SpEd25519Signature as InstanceSignedMsgSignature, }; + + impl super::KeySignExt for InstanceMsgKeyPair { + fn sign_prehash(&self, prehash: &[u8; 32]) -> InstanceSignedMsgSignature { + InstanceSignedMsgSignature(self.0.sign_prehashed(prehash)) + } + } } #[cfg(all( @@ -61,6 +79,12 @@ pub mod key_types { K256Ecdsa as Curve, K256Signature as InstanceSignedMsgSignature, K256SigningKey as InstanceMsgKeyPair, K256VerifyingKey as InstanceMsgPublicKey, }; + + impl super::KeySignExt for InstanceMsgKeyPair { + fn sign_prehash(&self, prehash: &[u8; 32]) -> InstanceSignedMsgSignature { + self.sign_prehash(prehash) + } + } } // Compile-time assertion to ensure only one feature is enabled @@ -72,3 +96,7 @@ pub mod key_types { compile_error!( "Only one of 'sp-core-ecdsa', 'sp-core-sr25519', or 'sp-core-ed25519' features can be enabled at a time" ); + +pub(crate) trait KeySignExt { + fn sign_prehash(&self, prehash: &[u8; 32]) -> InstanceSignedMsgSignature; +} \ No newline at end of file From a08b221eb59de5466b2b68fd84d947b49e861041 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Tue, 18 Feb 2025 12:30:09 -0700 Subject: [PATCH 08/52] chore: networking build --- crates/networking/src/behaviours.rs | 6 +- .../src/blueprint_protocol/behaviour.rs | 12 ++- crates/networking/src/discovery/peers.rs | 6 ++ crates/networking/src/service.rs | 90 ++++++++++--------- 4 files changed, 70 insertions(+), 44 deletions(-) diff --git a/crates/networking/src/behaviours.rs b/crates/networking/src/behaviours.rs index b1152dc51..3c843d0e1 100644 --- a/crates/networking/src/behaviours.rs +++ b/crates/networking/src/behaviours.rs @@ -39,9 +39,9 @@ pub struct GadgetBehaviour { /// Connection limits to prevent DoS connection_limits: connection_limits::Behaviour, /// Discovery mechanisms (Kademlia, mDNS, etc) - discovery: DiscoveryBehaviour, + pub(super) discovery: DiscoveryBehaviour, /// Direct P2P messaging and gossip - blueprint_protocol: BlueprintProtocolBehaviour, + pub(super) blueprint_protocol: BlueprintProtocolBehaviour, /// Connection liveness checks ping: ping::Behaviour, } @@ -49,6 +49,7 @@ pub struct GadgetBehaviour { impl GadgetBehaviour { pub fn new( network_name: &str, + blueprint_protocol_name: &str, local_key: &Keypair, instance_secret_key: &InstanceMsgKeyPair, instance_public_key: &InstanceMsgPublicKey, @@ -86,6 +87,7 @@ impl GadgetBehaviour { instance_secret_key, instance_public_key, peer_manager, + blueprint_protocol_name, ); Self { diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index ef39062c0..86603a6be 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -1,4 +1,6 @@ -use crate::{Curve, InstanceMsgKeyPair, InstanceMsgPublicKey, InstanceSignedMsgSignature, KeySignExt}; +use crate::{ + Curve, InstanceMsgKeyPair, InstanceMsgPublicKey, InstanceSignedMsgSignature, KeySignExt, +}; use dashmap::{DashMap, DashSet}; use gadget_crypto::{hashing::blake3_256, KeyType}; use gadget_logging::{debug, trace, warn}; @@ -59,6 +61,8 @@ pub enum BlueprintProtocolEvent { pub struct BlueprintProtocolBehaviour { /// Request/response protocol for direct messaging blueprint_protocol: DerivedBlueprintProtocolBehaviour, + /// Name of the blueprint protocol + pub(crate) blueprint_protocol_name: String, /// Peer manager for tracking peer states pub(crate) peer_manager: Arc, /// Libp2p peer ID @@ -83,9 +87,12 @@ impl BlueprintProtocolBehaviour { instance_secret_key: &InstanceMsgKeyPair, instance_public_key: &InstanceMsgPublicKey, peer_manager: Arc, + blueprint_protocol_name: &str, ) -> Self { + let blueprint_protocol_name = blueprint_protocol_name.to_string(); let protocols = vec![( - StreamProtocol::new("/gadget/blueprint_protocol/1.0.0"), + StreamProtocol::try_from_owned(blueprint_protocol_name.to_string()) + .unwrap_or_else(|_| StreamProtocol::new("/blueprint_protocol/1.0.0")), request_response::ProtocolSupport::Full, )]; @@ -115,6 +122,7 @@ impl BlueprintProtocolBehaviour { Self { blueprint_protocol, + blueprint_protocol_name, peer_manager, local_peer_id, instance_public_key: instance_public_key.clone(), diff --git a/crates/networking/src/discovery/peers.rs b/crates/networking/src/discovery/peers.rs index 9921410a6..a3d7679cd 100644 --- a/crates/networking/src/discovery/peers.rs +++ b/crates/networking/src/discovery/peers.rs @@ -144,6 +144,12 @@ impl PeerManager { }); } + /// Bans a peer with the default duration(`1h`) + pub async fn ban_peer_with_default_duration(&self, peer: PeerId, reason: impl Into) { + const BAN_PEER_DURATION: Duration = Duration::from_secs(60 * 60); //1h + self.ban_peer(peer, reason, Some(BAN_PEER_DURATION)) + } + /// Unban a peer pub fn unban_peer(&self, peer_id: &PeerId) { if self.banned_peers.remove(peer_id).is_some() { diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index 83f5ff4de..27fafb4a5 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -1,11 +1,15 @@ -use std::{collections::HashMap, sync::Arc, time::Duration}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, + time::Duration, +}; use crate::{ behaviours::{GadgetBehaviour, GadgetBehaviourEvent}, blueprint_protocol::{BlueprintProtocolEvent, InstanceMessageRequest, InstanceMessageResponse}, discovery::{ behaviour::{DerivedDiscoveryBehaviour, DerivedDiscoveryBehaviourEvent, DiscoveryEvent}, - PeerManager, + PeerInfo, PeerManager, }, error::Error, key_types::{InstanceMsgKeyPair, InstanceMsgPublicKey, InstanceSignedMsgSignature}, @@ -15,6 +19,7 @@ use gadget_logging::trace; use libp2p::{ core::{transport::Boxed, upgrade}, gossipsub::{IdentTopic, Topic}, + identify, identity::Keypair, noise, ping, swarm::{ConnectionId, SwarmEvent}, @@ -97,6 +102,8 @@ pub enum NetworkMessage { pub struct NetworkConfig { /// Network name/namespace pub network_name: String, + /// Instance id for blueprint protocol + pub instance_id: String, /// Instance secret key for blueprint protocol pub instance_secret_key: InstanceMsgKeyPair, /// Instance public key for blueprint protocol @@ -139,6 +146,7 @@ impl NetworkService { pub async fn new(config: NetworkConfig) -> Result { let NetworkConfig { network_name, + instance_id, instance_secret_key, instance_public_key, local_key, @@ -151,9 +159,12 @@ impl NetworkService { let peer_manager = Arc::new(PeerManager::default()); + let blueprint_protocol_name = format!("/blueprint_protocol/{}/1.0.0", instance_id); + // Create the swarm let behaviour = GadgetBehaviour::new( &network_name, + &blueprint_protocol_name, &local_key, &instance_secret_key, &instance_public_key, @@ -300,24 +311,17 @@ async fn handle_behaviour_event( ) -> Result<(), Error> { match event { GadgetBehaviourEvent::ConnectionLimits(_) => {} - GadgetBehaviourEvent::Discovery(discovery_event) => match discovery_event { - DiscoveryEvent::Discovery(derived_event) => { - handle_discovery_event( - swarm, - peer_manager, - derived_event, - event_sender, - network_sender, - ) - .await? - } - DiscoveryEvent::PeerConnected(peer) => { - event_sender.send(NetworkEvent::PeerConnected(peer))? - } - DiscoveryEvent::PeerDisconnected(peer) => { - event_sender.send(NetworkEvent::PeerDisconnected(peer))? - } - }, + GadgetBehaviourEvent::Discovery(discovery_event) => { + handle_discovery_event( + &swarm.behaviour().discovery.peer_info, + peer_manager, + discovery_event, + event_sender, + network_sender, + &swarm.behaviour().blueprint_protocol.blueprint_protocol_name, + ) + .await? + } GadgetBehaviourEvent::BlueprintProtocol(blueprint_event) => { handle_blueprint_protocol_event( swarm, @@ -345,33 +349,39 @@ async fn handle_behaviour_event( /// Handle a discovery event async fn handle_discovery_event( - swarm: &mut Swarm, + peer_info_map: &HashMap, peer_manager: &Arc, - event: Box, + event: DiscoveryEvent, event_sender: &mpsc::UnboundedSender, network_sender: &mpsc::UnboundedSender, + blueprint_protocol_name: &str, ) -> Result<(), Error> { - let event = event.as_ref(); match event { - DerivedDiscoveryBehaviourEvent::Kademlia(event) => { - // Handle Kademlia discovery events - } - DerivedDiscoveryBehaviourEvent::Mdns(event) => { - // Handle mDNS discovery events - } - DerivedDiscoveryBehaviourEvent::Identify(event) => { - // Handle Identify protocol events - } - DerivedDiscoveryBehaviourEvent::Autonat(event) => { - // Handle AutoNAT events + DiscoveryEvent::PeerConnected(peer_id) => { + trace!("Peer connected, {peer_id}"); + event_sender.send(NetworkEvent::PeerConnected(peer_id))?; } - DerivedDiscoveryBehaviourEvent::Upnp(event) => { - // Handle UPnP events + DiscoveryEvent::PeerDisconnected(peer_id) => { + trace!("Peer disconnected, {peer_id}"); + event_sender.send(NetworkEvent::PeerDisconnected(peer_id))?; } - DerivedDiscoveryBehaviourEvent::Relay(event) => { - // Handle relay events - } - } + DiscoveryEvent::Discovery(discovery_event) => match &*discovery_event { + DerivedDiscoveryBehaviourEvent::Identify(identify::Event::Received { + peer_id, + info, + .. + }) => { + let protocols: HashSet = + HashSet::from_iter(info.protocols.iter().map(|p| p.to_string())); + if !protocols.contains(blueprint_protocol_name) { + peer_manager + .ban_peer_with_default_duration(*peer_id, "hello protocol unsupported"); + } + } + DerivedDiscoveryBehaviourEvent::Identify(_) => {} + _ => {} + }, + }; Ok(()) } From 82e13cb1cb2e44f0d335e7b93fa56fd8b8672265 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Tue, 18 Feb 2025 13:33:26 -0700 Subject: [PATCH 09/52] feat: handshake requests and responses --- crates/crypto/k256/src/lib.rs | 6 +++ crates/crypto/sp-core/src/lib.rs | 6 +++ crates/crypto/sp-core/src/sp_core_util.rs | 1 - .../src/blueprint_protocol/behaviour.rs | 52 +++++++++++++++++-- .../src/blueprint_protocol/handler.rs | 21 ++------ crates/networking/src/discovery/peers.rs | 17 +++--- crates/networking/src/service.rs | 10 +++- 7 files changed, 81 insertions(+), 32 deletions(-) delete mode 100644 crates/crypto/sp-core/src/sp_core_util.rs diff --git a/crates/crypto/k256/src/lib.rs b/crates/crypto/k256/src/lib.rs index ed85d095f..c848bcf01 100644 --- a/crates/crypto/k256/src/lib.rs +++ b/crates/crypto/k256/src/lib.rs @@ -51,6 +51,12 @@ impl Ord for K256VerifyingKey { } } +impl gadget_std::hash::Hash for K256VerifyingKey { + fn hash(&self, state: &mut H) { + self.0.to_sec1_bytes().hash(state); + } +} + macro_rules! impl_serde_bytes { ($wrapper:ident, $inner:path) => { #[derive(Clone, PartialEq, Eq, Debug)] diff --git a/crates/crypto/sp-core/src/lib.rs b/crates/crypto/sp-core/src/lib.rs index ce2cbd02c..b78610066 100644 --- a/crates/crypto/sp-core/src/lib.rs +++ b/crates/crypto/sp-core/src/lib.rs @@ -77,6 +77,12 @@ macro_rules! impl_sp_core_pair_public { #[derive(Clone, serde::Serialize, serde::Deserialize)] pub struct [](pub <$pair_type as sp_core::Pair>::Public); + impl gadget_std::hash::Hash for [] { + fn hash(&self, state: &mut H) { + self.0.to_raw_vec().hash(state); + } + } + impl KeyEncoding for [] { fn to_bytes(&self) -> Vec { self.0.to_raw_vec() diff --git a/crates/crypto/sp-core/src/sp_core_util.rs b/crates/crypto/sp-core/src/sp_core_util.rs deleted file mode 100644 index 8b1378917..000000000 --- a/crates/crypto/sp-core/src/sp_core_util.rs +++ /dev/null @@ -1 +0,0 @@ - diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index 86603a6be..fa5e14056 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -3,7 +3,7 @@ use crate::{ }; use dashmap::{DashMap, DashSet}; use gadget_crypto::{hashing::blake3_256, KeyType}; -use gadget_logging::{debug, trace, warn}; +use gadget_logging::{debug, info, trace, warn}; use libp2p::{ core::transport::PortUse, gossipsub::{self, IdentTopic, MessageAuthenticity, MessageId, Sha256Topic}, @@ -304,10 +304,54 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { } fn on_swarm_event(&mut self, event: FromSwarm<'_>) { - if let FromSwarm::ConnectionEstablished(e) = &event { - if e.other_established == 0 { - self.inbound_handshakes.insert(e.peer_id, Instant::now()); + match &event { + FromSwarm::ConnectionEstablished(e) if e.other_established == 0 => { + // Start handshake if this peer is not verified + if !self.peer_manager.is_peer_verified(&e.peer_id) { + debug!( + "Established connection with unverified peer {:?}, sending handshake", + e.peer_id + ); + self.send_request( + &e.peer_id, + InstanceMessageRequest::Handshake { + public_key: self.instance_public_key, + signature: self.sign_handshake(&e.peer_id), + }, + ); + self.outbound_handshakes.insert(e.peer_id, Instant::now()); + info!( + "Established connection to {:?}, sending handshake", + e.peer_id + ); + } + + self.blueprint_protocol + .gossipsub + .add_explicit_peer(&e.peer_id); } + FromSwarm::ConnectionClosed(e) if e.remaining_established == 0 => { + if self.inbound_handshakes.contains_key(&e.peer_id) { + self.inbound_handshakes.remove(&e.peer_id); + } + + if self.outbound_handshakes.contains_key(&e.peer_id) { + self.outbound_handshakes.remove(&e.peer_id); + } + + if self.peer_manager.is_peer_verified(&e.peer_id) { + self.peer_manager + .remove_peer(&e.peer_id, "connection closed"); + } + + self.blueprint_protocol + .gossipsub + .remove_explicit_peer(&e.peer_id); + + self.peer_manager.remove_peer_id_from_public_key(&e.peer_id); + } + + _ => {} } self.blueprint_protocol.on_swarm_event(event) diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index 40ab5b2a0..9c9a0eb11 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -32,18 +32,11 @@ impl BlueprintProtocolBehaviour { } => { debug!(%peer, "Received handshake request"); - // Check if we already have a pending outbound handshake - if let Some(outbound_time) = self.outbound_handshakes.get(&peer) { - // If we have an outbound handshake and their peer_id is less than ours, - // we should wait for their response instead of responding to their request - // TODO: Fix - // if peer < &self.local_peer_id { - // debug!(%peer, "Deferring inbound handshake - waiting for outbound response"); - // return; - // } - // If we have an outbound handshake and their peer_id is greater than ours, - // we should handle their request and cancel our outbound attempt - self.outbound_handshakes.remove(&peer); + // Check if we already sent a handshake request to this peer + if self.outbound_handshakes.contains_key(&peer) { + // If we have an outbound handshake pending, we should still respond to their request + // This ensures both sides complete their handshakes even if messages cross on the wire + debug!(%peer, "Responding to inbound handshake request while outbound is pending"); } // Verify the handshake @@ -60,12 +53,8 @@ impl BlueprintProtocolBehaviour { if let Err(e) = self.send_response(channel, response) { warn!(%peer, "Failed to send handshake response: {:?}", e); - self.handle_handshake_failure(&peer, "Failed to send response"); return; } - - // Mark handshake as completed - self.complete_handshake(&peer, &public_key); } Err(e) => { warn!(%peer, "Invalid handshake request: {:?}", e); diff --git a/crates/networking/src/discovery/peers.rs b/crates/networking/src/discovery/peers.rs index a3d7679cd..ed20b5539 100644 --- a/crates/networking/src/discovery/peers.rs +++ b/crates/networking/src/discovery/peers.rs @@ -65,7 +65,7 @@ pub struct PeerManager { /// Verified peers from completed handshakes verified_peers: DashSet, /// Handshake keys to peer ids - public_keys_to_peer_ids: Arc>>, + public_keys_to_peer_ids: DashMap, /// Banned peers with optional expiration time banned_peers: DashMap>, /// Event sender for peer updates @@ -79,7 +79,7 @@ impl Default for PeerManager { peers: Default::default(), banned_peers: Default::default(), verified_peers: Default::default(), - public_keys_to_peer_ids: Arc::new(RwLock::new(BTreeMap::new())), + public_keys_to_peer_ids: Default::default(), event_tx, } } @@ -224,16 +224,15 @@ impl PeerManager { } /// Add a peer id to the public key to peer id map after verifying handshake - pub async fn add_peer_id_to_public_key( - &self, - peer_id: &PeerId, - public_key: &InstanceMsgPublicKey, - ) { + pub fn add_peer_id_to_public_key(&self, peer_id: &PeerId, public_key: &InstanceMsgPublicKey) { self.public_keys_to_peer_ids - .write() - .await .insert(public_key.clone(), *peer_id); } + + /// Remove a peer id from the public key to peer id map + pub fn remove_peer_id_from_public_key(&self, peer_id: &PeerId) { + self.public_keys_to_peer_ids.retain(|_, id| id != peer_id); + } } /// Update the average response time for a peer diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index 27fafb4a5..b71cebab4 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -257,14 +257,20 @@ impl NetworkService { swarm_event = swarm_stream.next() => match swarm_event { // outbound events Some(SwarmEvent::Behaviour(event)) => { - handle_behaviour_event( + match handle_behaviour_event( swarm_stream.get_mut(), &self.peer_manager, event, &self.event_sender, &self.network_sender, ) - .await; + .await + { + Ok(_) => {} + Err(e) => { + warn!("Failed to handle swarm event: {}", e); + } + } }, None => { break; }, _ => { }, From 9bf1006ae73497a7e874297fa9657a56e4c6cd63 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Tue, 18 Feb 2025 13:51:35 -0700 Subject: [PATCH 10/52] chore: add mermaid diagrams --- crates/networking/README.md | 129 ++++++++++++++++++ .../src/blueprint_protocol/handler.rs | 8 +- crates/networking/src/service.rs | 2 +- 3 files changed, 135 insertions(+), 4 deletions(-) create mode 100644 crates/networking/README.md diff --git a/crates/networking/README.md b/crates/networking/README.md new file mode 100644 index 000000000..71c8d1fb1 --- /dev/null +++ b/crates/networking/README.md @@ -0,0 +1,129 @@ +# Networking Protocol Documentation + +This document outlines the key protocols used in the networking layer. + +## Handshake Protocol + +The handshake protocol ensures mutual authentication between peers before allowing protocol messages. + +```mermaid +sequenceDiagram + participant A as Peer A + participant B as Peer B + + Note over A,B: Connection Established + + A->>B: HandshakeRequest { + public_key: A_pub, + signature: sign(B_peer_id) + } + + Note over B: Verify A's signature + Note over B: Store A's public key + + B->>A: HandshakeResponse { + public_key: B_pub, + signature: sign(A_peer_id) + } + + Note over A: Verify B's signature + Note over A: Store B's public key + + Note over A,B: Both peers verified + Note over A,B: Protocol messages allowed +``` + +### Handshake States + +```mermaid +stateDiagram-v2 + [*] --> Connected: Connection Established + Connected --> OutboundPending: Send Handshake + Connected --> InboundPending: Receive Handshake + + OutboundPending --> Verified: Valid Response + OutboundPending --> Failed: Invalid/Timeout + + InboundPending --> Verified: Valid Request + InboundPending --> Failed: Invalid/Timeout + + Verified --> [*]: Connection Closed + Failed --> [*]: Connection Closed +``` + +## Blueprint Protocol Instance Request/Response + +After handshake completion, peers can exchange protocol messages. + +```mermaid +sequenceDiagram + participant A as Peer A (Verified) + participant B as Peer B (Verified) + + Note over A,B: Handshake Completed + + A->>B: InstanceMessageRequest::Protocol { + protocol: String, + payload: Vec, + metadata: Option> + } + + alt Success Response + B->>A: InstanceMessageResponse::Success { + data: Option> + } + else Protocol Response + B->>A: InstanceMessageResponse::Protocol { + data: Vec + } + else Error Response + B->>A: InstanceMessageResponse::Error { + code: u16, + message: String + } + end +``` + +### Message Flow States + +```mermaid +stateDiagram-v2 + [*] --> Handshaked: Peers Verified + + Handshaked --> RequestPending: Send Request + RequestPending --> Processing: Request Received + + Processing --> ResponseSent: Success/Protocol + Processing --> ErrorSent: Error + + ResponseSent --> Handshaked: Complete + ErrorSent --> Handshaked: Complete + + Handshaked --> [*]: Connection Closed +``` + +## Protocol Details + +### Handshake Protocol + +- Initiated on first connection +- Mutual authentication using public key cryptography +- Signatures verify peer identity +- Timeouts after 30 seconds +- Handles concurrent handshakes gracefully + +### Blueprint Protocol + +- Requires completed handshake +- Supports protocol-specific messages +- Includes metadata for routing/handling +- Error responses for protocol violations +- Supports both direct and broadcast messaging + +### Security Features + +- Peer verification before message acceptance +- Signature verification for handshakes +- Banned peer tracking +- Connection limits +- Protocol version validation diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index 9c9a0eb11..b29ef8da2 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -44,6 +44,8 @@ impl BlueprintProtocolBehaviour { Ok(()) => { // Store the handshake request self.inbound_handshakes.insert(peer, Instant::now()); + self.peer_manager + .add_peer_id_to_public_key(&peer, &public_key); // Send handshake response let response = InstanceMessageResponse::Handshake { @@ -198,12 +200,12 @@ impl BlueprintProtocolBehaviour { self.inbound_handshakes.remove(peer); self.outbound_handshakes.remove(peer); - // Add to verified peers - self.peer_manager.verify_peer(peer); - // Update peer manager self.peer_manager .add_peer_id_to_public_key(peer, public_key); + + // Add to verified peers + self.peer_manager.verify_peer(peer); } pub fn handle_gossipsub_event(&mut self, event: gossipsub::Event) { diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index b71cebab4..87ea5b787 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -216,7 +216,7 @@ impl NetworkService { } /// Run the network service - async fn run(mut self, event_sender: mpsc::UnboundedSender) { + async fn run(mut self) { info!("Starting network service"); // Bootstrap with Kademlia From bc88d2c3ab29fff2ef697f7ff930597e9a23e7af Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Tue, 18 Feb 2025 17:41:11 -0700 Subject: [PATCH 11/52] feat: add compiling networking and round-based-extension --- Cargo.lock | 27 + Cargo.toml | 4 +- .../macros/blueprint-proc-macro/src/shared.rs | 2 +- .../Cargo.toml | 24 + .../src/lib.rs | 280 +++++++ crates/networking/Cargo.toml | 1 + crates/networking/src/behaviours.rs | 12 +- .../src/blueprint_protocol/behaviour.rs | 25 +- .../src/blueprint_protocol/handler.rs | 85 +- .../networking/src/blueprint_protocol/mod.rs | 7 +- crates/networking/src/discovery/behaviour.rs | 12 +- crates/networking/src/discovery/config.rs | 9 +- crates/networking/src/discovery/mod.rs | 1 + crates/networking/src/discovery/peers.rs | 55 +- crates/networking/src/error.rs | 3 + crates/networking/src/lib.rs | 5 +- crates/networking/src/networking/tests.rs | 768 ------------------ crates/networking/src/service.rs | 225 +++-- crates/networking/src/service_handle.rs | 134 +++ crates/networking/src/types.rs | 8 +- 20 files changed, 730 insertions(+), 957 deletions(-) create mode 100644 crates/networking-round-based-extension/Cargo.toml create mode 100644 crates/networking-round-based-extension/src/lib.rs create mode 100644 crates/networking/src/service_handle.rs diff --git a/Cargo.lock b/Cargo.lock index 30e3ecbc9..99a12fe70 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3934,6 +3934,15 @@ dependencies = [ "chrono", ] +[[package]] +name = "crossbeam" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1137cd7e7fc0fb5d3c5a8678be38ec56e819125d8d7907411fe24ccb943faca8" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-channel" version = "0.5.14" @@ -7152,6 +7161,7 @@ dependencies = [ "auto_impl", "bincode", "blake3", + "crossbeam-channel", "dashmap", "futures", "gadget-crypto", @@ -7174,6 +7184,23 @@ dependencies = [ "tracing-subscriber 0.3.19", ] +[[package]] +name = "gadget-networking-round-based-extension" +version = "0.1.0" +dependencies = [ + "crossbeam", + "crossbeam-channel", + "dashmap", + "futures", + "gadget-networking", + "round-based", + "serde", + "serde_json", + "thiserror 2.0.11", + "tokio", + "tracing", +] + [[package]] name = "gadget-rpc-calls" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 2aa2e85ea..36aa1d505 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -108,7 +108,7 @@ gadget-std = { version = "0.1.0", path = "./crates/std", default-features = fals # P2P gadget-networking = { version = "0.1.0", path = "./crates/networking", default-features = false } -gadget-networking-behaviours = { version = "0.1.0", path = "./crates/networking/behaviours", default-features = false } +gadget-networking-round-based-extension = { version = "0.1.0", path = "./crates/networking-round-based-extension", default-features = false } # Utilities gadget-utils = { version = "0.1.0", path = "./crates/utils", default-features = false } @@ -143,6 +143,8 @@ sp-runtime = { version = "39.0.0", default-features = false } # Async & Runtime async-trait = { version = "0.1.86", default-features = false } +crossbeam = { version = "0.8", default-features = false } +crossbeam-channel = { version = "0.5", default-features = false } futures = { version = "0.3.30", default-features = false } futures-util = { version = "0.3.31", default-features = false } tokio = { version = "1.40", default-features = false } diff --git a/crates/macros/blueprint-proc-macro/src/shared.rs b/crates/macros/blueprint-proc-macro/src/shared.rs index 7eb8a86e0..0300e5871 100644 --- a/crates/macros/blueprint-proc-macro/src/shared.rs +++ b/crates/macros/blueprint-proc-macro/src/shared.rs @@ -237,7 +237,7 @@ pub(crate) trait MacroExt { ResultsKind::Types(types) => { let xs = types .iter() - .map(|ty| type_to_field_type(ty)) + .map(type_to_field_type) .collect::>>()?; Ok(xs) } diff --git a/crates/networking-round-based-extension/Cargo.toml b/crates/networking-round-based-extension/Cargo.toml new file mode 100644 index 000000000..1ca17c5d3 --- /dev/null +++ b/crates/networking-round-based-extension/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "gadget-networking-round-based-extension" +version = "0.1.0" +authors.workspace = true +edition.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true + +[dependencies] +gadget-networking = { path = "../networking" } +round-based = { workspace = true } +tokio = { workspace = true } +futures = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +tracing = { workspace = true } +dashmap = { workspace = true } +crossbeam = { workspace = true } +crossbeam-channel = { workspace = true } +thiserror = { workspace = true } + +[lints] +workspace = true diff --git a/crates/networking-round-based-extension/src/lib.rs b/crates/networking-round-based-extension/src/lib.rs new file mode 100644 index 000000000..ba82b50a5 --- /dev/null +++ b/crates/networking-round-based-extension/src/lib.rs @@ -0,0 +1,280 @@ +use crossbeam_channel::{self, Receiver, Sender}; +use dashmap::DashMap; +use futures::Future; +use futures::{Sink, Stream}; +use gadget_networking::{ + key_types::InstanceMsgPublicKey, + service_handle::NetworkServiceHandle, + types::{ParticipantInfo, ProtocolMessage}, +}; +use round_based::{ + Delivery, Incoming, MessageDestination, MessageType, MsgId, Outgoing, PartyIndex, +}; +use serde::{de::DeserializeOwned, Serialize}; +use std::{ + collections::HashMap, + pin::Pin, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }, + task::{Context, Poll}, +}; + +/// Wrapper to adapt NetworkServiceHandle to round-based protocols +pub struct RoundBasedNetworkAdapter { + /// The underlying network handle + handle: NetworkServiceHandle, + /// Current party's index + party_index: PartyIndex, + /// Mapping of party indices to their public keys + parties: Arc>, + /// Counter for message IDs + next_msg_id: Arc, + /// Channel for forwarding messages + forward_tx: Sender<(PartyIndex, M)>, + /// Channel for receiving forwarded messages + forward_rx: Receiver<(PartyIndex, M)>, + /// Protocol identifier + protocol_id: String, + _phantom: std::marker::PhantomData, +} + +impl RoundBasedNetworkAdapter +where + M: Clone + Send + Sync + Unpin + 'static, + M: Serialize + DeserializeOwned, + M: round_based::ProtocolMessage, +{ + pub fn new( + handle: NetworkServiceHandle, + party_index: PartyIndex, + parties: HashMap, + protocol_id: impl Into, + ) -> Self { + let (forward_tx, forward_rx) = crossbeam_channel::unbounded(); + + Self { + handle, + party_index, + parties: Arc::new(DashMap::from_iter(parties)), + next_msg_id: Arc::new(AtomicU64::new(0)), + forward_tx, + forward_rx, + protocol_id: protocol_id.into(), + _phantom: std::marker::PhantomData, + } + } +} + +impl Delivery for RoundBasedNetworkAdapter +where + M: Clone + Send + Sync + Unpin + 'static, + M: Serialize + DeserializeOwned, + M: round_based::ProtocolMessage, +{ + type Send = RoundBasedSender; + type Receive = RoundBasedReceiver; + type SendError = NetworkError; + type ReceiveError = NetworkError; + + fn split(self) -> (Self::Receive, Self::Send) { + let RoundBasedNetworkAdapter { + handle, + party_index, + parties, + next_msg_id, + forward_tx, + forward_rx, + protocol_id, + .. + } = self; + + let sender = RoundBasedSender { + handle: handle.clone(), + party_index, + parties: parties.clone(), + next_msg_id: next_msg_id.clone(), + protocol_id: protocol_id.clone(), + forward_tx, + _phantom: std::marker::PhantomData, + }; + + let receiver = RoundBasedReceiver::new(handle, party_index, forward_rx); + + (receiver, sender) + } +} + +pub struct RoundBasedSender { + handle: NetworkServiceHandle, + party_index: PartyIndex, + parties: Arc>, + next_msg_id: Arc, + protocol_id: String, + forward_tx: Sender<(PartyIndex, M)>, + _phantom: std::marker::PhantomData, +} + +impl Sink> for RoundBasedSender +where + M: Serialize + round_based::ProtocolMessage + Clone, +{ + type Error = NetworkError; + + fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn start_send(self: Pin<&mut Self>, outgoing: Outgoing) -> Result<(), Self::Error> { + let this = unsafe { self.get_unchecked_mut() }; + let msg_id = this.next_msg_id.fetch_add(1, Ordering::Relaxed); + let round = outgoing.msg.round(); + + // Handle local message forwarding for self-messages + if outgoing.recipient == MessageDestination::OneParty(this.party_index) { + return this + .forward_tx + .send((this.party_index, outgoing.msg.clone())) + .map_err(|_| NetworkError::Send("Failed to forward local message".into())); + } + + let (recipient, recipient_key) = match outgoing.recipient { + MessageDestination::AllParties => (None, None), + MessageDestination::OneParty(p) => { + let key = this.parties.get(&p).map(|k| k.clone()); + (Some(p), key) + } + }; + + let protocol_message = ProtocolMessage { + protocol: format!("{}/{}", this.protocol_id, round), + routing: gadget_networking::types::MessageRouting { + message_id: msg_id, + round_id: round, + sender: ParticipantInfo { + id: gadget_networking::types::ParticipantId(this.party_index), + public_key: this.parties.get(&this.party_index).map(|k| k.clone()), + }, + recipient: recipient.map(|p| ParticipantInfo { + id: gadget_networking::types::ParticipantId(p), + public_key: recipient_key, + }), + }, + payload: serde_json::to_vec(&outgoing.msg) + .map_err(|e| NetworkError::Serialization(e))?, + }; + + this.handle + .send_protocol_message(protocol_message) + .map_err(|e| NetworkError::Send(e)) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } +} + +pub struct RoundBasedReceiver { + handle: NetworkServiceHandle, + party_index: PartyIndex, + forward_rx: Receiver<(PartyIndex, M)>, + _phantom: std::marker::PhantomData, + next_message_future: Option> + Send>>>, +} + +impl RoundBasedReceiver { + fn new( + handle: NetworkServiceHandle, + party_index: PartyIndex, + forward_rx: Receiver<(PartyIndex, M)>, + ) -> Self { + Self { + handle, + party_index, + forward_rx, + _phantom: std::marker::PhantomData, + next_message_future: None, + } + } +} + +impl Stream for RoundBasedReceiver +where + M: DeserializeOwned + round_based::ProtocolMessage + Unpin, +{ + type Item = Result, NetworkError>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + // First check forwarded messages + if let Ok((sender, msg)) = self.forward_rx.try_recv() { + return Poll::Ready(Some(Ok(Incoming { + msg, + sender, + id: 0, + msg_type: MessageType::P2P, + }))); + } + + // Get a mutable reference to self + let this = self.get_mut(); + + // Create and store the future if we don't have one + if this.next_message_future.is_none() { + let mut handle = this.handle.clone(); + this.next_message_future = + Some(Box::pin( + async move { handle.next_protocol_message().await }, + )); + } + + // Poll the stored future + if let Some(future) = &mut this.next_message_future { + match future.as_mut().poll(cx) { + Poll::Ready(Some(msg)) => { + // Clear the future so we create a new one next time + this.next_message_future = None; + + let msg_type = if msg.routing.recipient.is_some() { + MessageType::P2P + } else { + MessageType::Broadcast + }; + + let sender = msg.routing.sender.id.0; + let id = msg.routing.message_id; + + match serde_json::from_slice(&msg.payload) { + Ok(msg) => Poll::Ready(Some(Ok(Incoming { + msg, + sender, + id, + msg_type, + }))), + Err(e) => Poll::Ready(Some(Err(NetworkError::Serialization(e)))), + } + } + Poll::Ready(None) => { + this.next_message_future = None; + Poll::Ready(None) + } + Poll::Pending => Poll::Pending, + } + } else { + // This shouldn't happen because we create the future above if it's None + Poll::Ready(None) + } + } +} + +#[derive(Debug, thiserror::Error)] +pub enum NetworkError { + #[error("Failed to serialize message: {0}")] + Serialization(#[from] serde_json::Error), + #[error("Network error: {0}")] + Send(String), +} diff --git a/crates/networking/Cargo.toml b/crates/networking/Cargo.toml index 5e245f343..24c466b3e 100644 --- a/crates/networking/Cargo.toml +++ b/crates/networking/Cargo.toml @@ -32,6 +32,7 @@ itertools = { workspace = true, features = ["use_alloc"] } parking_lot = { workspace = true } thiserror = { workspace = true } anyhow = { workspace = true } +crossbeam-channel = { workspace = true } # Crypto dependencies gadget-crypto = { workspace = true, features = ["k256", "hashing"] } diff --git a/crates/networking/src/behaviours.rs b/crates/networking/src/behaviours.rs index 3c843d0e1..a5fa02ba7 100644 --- a/crates/networking/src/behaviours.rs +++ b/crates/networking/src/behaviours.rs @@ -1,4 +1,5 @@ use crate::key_types::{InstanceMsgKeyPair, InstanceMsgPublicKey}; +use crate::types::ProtocolMessage; use crate::{ blueprint_protocol::{BlueprintProtocolBehaviour, BlueprintProtocolEvent}, discovery::{ @@ -7,6 +8,7 @@ use crate::{ PeerInfo, PeerManager, }, }; +use crossbeam_channel::Sender; use libp2p::{ connection_limits::{self, ConnectionLimits}, identity::Keypair, @@ -23,7 +25,7 @@ use std::{ const MAX_ESTABLISHED_PER_PEER: u32 = 4; -/// Events that can be emitted by the GadgetBehavior +/// Events that can be emitted by the `GadgetBehavior` #[derive(Debug)] pub enum GadgetEvent { /// Discovery-related events @@ -36,7 +38,7 @@ pub enum GadgetEvent { #[derive(NetworkBehaviour)] pub struct GadgetBehaviour { - /// Connection limits to prevent DoS + /// Connection limits to prevent `DoS` connection_limits: connection_limits::Behaviour, /// Discovery mechanisms (Kademlia, mDNS, etc) pub(super) discovery: DiscoveryBehaviour, @@ -47,6 +49,7 @@ pub struct GadgetBehaviour { } impl GadgetBehaviour { + #[must_use] pub fn new( network_name: &str, blueprint_protocol_name: &str, @@ -55,6 +58,7 @@ impl GadgetBehaviour { instance_public_key: &InstanceMsgPublicKey, target_peer_count: u64, peer_manager: Arc, + protocol_message_sender: Sender, ) -> Self { let connection_limits = connection_limits::Behaviour::new( ConnectionLimits::default() @@ -88,6 +92,7 @@ impl GadgetBehaviour { instance_public_key, peer_manager, blueprint_protocol_name, + protocol_message_sender, ); Self { @@ -104,15 +109,18 @@ impl GadgetBehaviour { } /// Returns a set of peer ids + #[must_use] pub fn peers(&self) -> &HashSet { self.discovery.get_peers() } /// Returns a map of peer ids and their multi-addresses + #[must_use] pub fn peer_addresses(&self) -> HashMap> { self.discovery.get_peer_addresses() } + #[must_use] pub fn peer_info(&self, peer_id: &PeerId) -> Option<&PeerInfo> { self.discovery.get_peer_info(peer_id) } diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index fa5e14056..fc7f10334 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -1,7 +1,9 @@ use crate::{ - Curve, InstanceMsgKeyPair, InstanceMsgPublicKey, InstanceSignedMsgSignature, KeySignExt, + types::ProtocolMessage, Curve, InstanceMsgKeyPair, InstanceMsgPublicKey, + InstanceSignedMsgSignature, KeySignExt, }; -use dashmap::{DashMap, DashSet}; +use crossbeam_channel::Sender; +use dashmap::DashMap; use gadget_crypto::{hashing::blake3_256, KeyType}; use gadget_logging::{debug, info, trace, warn}; use libp2p::{ @@ -34,7 +36,7 @@ pub struct DerivedBlueprintProtocolBehaviour { gossipsub: gossipsub::Behaviour, } -/// Events emitted by the BlueprintProtocolBehaviour +/// Events emitted by the `BlueprintProtocolBehaviour` #[derive(Debug)] pub enum BlueprintProtocolEvent { /// Request received from a peer @@ -67,9 +69,9 @@ pub struct BlueprintProtocolBehaviour { pub(crate) peer_manager: Arc, /// Libp2p peer ID pub(crate) local_peer_id: PeerId, - /// Instance public key for handshakes and blueprint_protocol + /// Instance public key for handshakes and `blueprint_protocol` pub(crate) instance_public_key: InstanceMsgPublicKey, - /// Instance secret key for handshakes and blueprint_protocol + /// Instance secret key for handshakes and `blueprint_protocol` pub(crate) instance_secret_key: InstanceMsgKeyPair, /// Peers with pending inbound handshakes pub(crate) inbound_handshakes: DashMap, @@ -78,16 +80,20 @@ pub struct BlueprintProtocolBehaviour { /// Active response channels pub(crate) response_channels: DashMap>, + /// Protocol message sender + pub(crate) protocol_message_sender: Sender, } impl BlueprintProtocolBehaviour { /// Create a new blueprint protocol behaviour + #[must_use] pub fn new( local_key: &Keypair, instance_secret_key: &InstanceMsgKeyPair, instance_public_key: &InstanceMsgPublicKey, peer_manager: Arc, blueprint_protocol_name: &str, + protocol_message_sender: Sender, ) -> Self { let blueprint_protocol_name = blueprint_protocol_name.to_string(); let protocols = vec![( @@ -125,11 +131,12 @@ impl BlueprintProtocolBehaviour { blueprint_protocol_name, peer_manager, local_peer_id, - instance_public_key: instance_public_key.clone(), + instance_public_key: *instance_public_key, instance_secret_key: instance_secret_key.clone(), inbound_handshakes: DashMap::new(), outbound_handshakes: DashMap::new(), response_channels: DashMap::new(), + protocol_message_sender, } } @@ -300,7 +307,7 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { event: THandlerOutEvent, ) { self.blueprint_protocol - .on_connection_handler_event(peer_id, connection_id, event) + .on_connection_handler_event(peer_id, connection_id, event); } fn on_swarm_event(&mut self, event: FromSwarm<'_>) { @@ -354,7 +361,7 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { _ => {} } - self.blueprint_protocol.on_swarm_event(event) + self.blueprint_protocol.on_swarm_event(event); } fn poll( @@ -368,7 +375,7 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { blueprint_protocol_event, ) => self.handle_request_response_event(blueprint_protocol_event), DerivedBlueprintProtocolBehaviourEvent::Gossipsub(gossip_event) => { - self.handle_gossipsub_event(gossip_event) + self.handle_gossipsub_event(gossip_event); } }, ToSwarm::Dial { opts } => { diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index b29ef8da2..a33a74cfe 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -1,9 +1,12 @@ use std::time::{Duration, Instant}; -use libp2p::{gossipsub, request_response, PeerId}; +use libp2p::{ + gossipsub::{self}, + request_response, PeerId, +}; use tracing::{debug, warn}; -use crate::key_types::InstanceMsgPublicKey; +use crate::{key_types::InstanceMsgPublicKey, types::ProtocolMessage}; use super::{BlueprintProtocolBehaviour, InstanceMessageRequest, InstanceMessageResponse}; @@ -49,7 +52,7 @@ impl BlueprintProtocolBehaviour { // Send handshake response let response = InstanceMessageResponse::Handshake { - public_key: self.instance_public_key.clone(), + public_key: self.instance_public_key, signature: self.sign_handshake(&peer), }; @@ -67,7 +70,6 @@ impl BlueprintProtocolBehaviour { if let Err(e) = self.send_response(channel, response) { warn!(%peer, "Failed to send error response: {:?}", e); } - self.handle_handshake_failure(&peer, "Invalid handshake"); } } } @@ -92,6 +94,12 @@ impl BlueprintProtocolBehaviour { return; } + if !self.peer_manager.is_key_whitelisted(&public_key) { + warn!(%peer, "Received handshake response from unwhitelisted peer"); + self.peer_manager.handle_nonwhitelisted_peer(&peer); + return; + } + // Verify the handshake match self.verify_handshake(&peer, &public_key, &signature) { Ok(()) => { @@ -99,25 +107,11 @@ impl BlueprintProtocolBehaviour { self.complete_handshake(&peer, &public_key); } Err(e) => { - warn!(%peer, "Invalid handshake response: {:?}", e); - self.handle_handshake_failure(&peer, "Invalid handshake response"); + warn!(%peer, "Invalid handshake verification: {:?}", e); + self.outbound_handshakes.remove(&peer); + self.handle_handshake_failure(&peer, "Invalid handshake verification"); } } - - // Remove the outbound handshake - self.outbound_handshakes.remove(&peer); - } - request_response::Event::Message { - peer, - message: - request_response::Message::Response { - response: InstanceMessageResponse::Error { code, message }, - .. - }, - .. - } => { - warn!(%peer, code, %message, "Received error response"); - self.handle_handshake_failure(&peer, &message); } request_response::Event::Message { peer, @@ -134,6 +128,11 @@ impl BlueprintProtocolBehaviour { }, .. } => { + // Reject messages from self + if peer == self.local_peer_id { + return; + } + // Only accept protocol messages from peers we've completed handshakes with if !self.peer_manager.is_peer_verified(&peer) { warn!(%peer, "Received protocol message from unverified peer"); @@ -147,8 +146,48 @@ impl BlueprintProtocolBehaviour { return; } - debug!(%peer, %protocol, "Received protocol request"); - // Handle protocol message... + let protocol_message: ProtocolMessage = match bincode::deserialize(&payload) { + Ok(message) => message, + Err(e) => { + warn!(%peer, "Failed to deserialize protocol message: {:?}", e); + let response = InstanceMessageResponse::Error { + code: 400, + message: format!("Invalid protocol message: {:?}", e), + }; + if let Err(e) = self.send_response(channel, response) { + warn!(%peer, "Failed to send error response: {:?}", e); + } + return; + } + }; + + debug!(%peer, %protocol, %protocol_message, "Received protocol request"); + self.protocol_message_sender.send(protocol_message); + } + request_response::Event::Message { + peer, + message: + request_response::Message::Response { + response: InstanceMessageResponse::Error { code, message }, + .. + }, + .. + } => { + if !self.peer_manager.is_peer_verified(&peer) { + warn!(%peer, code, %message, "Received error response from unverified peer"); + return; + } + } + request_response::Event::Message { + peer, + message: + request_response::Message::Response { + response: InstanceMessageResponse::Success { protocol, data }, + .. + }, + .. + } => { + debug!(%peer, %protocol, "Received successful protocol response"); } _ => {} } diff --git a/crates/networking/src/blueprint_protocol/mod.rs b/crates/networking/src/blueprint_protocol/mod.rs index 9f4efe078..8973e8fb5 100644 --- a/crates/networking/src/blueprint_protocol/mod.rs +++ b/crates/networking/src/blueprint_protocol/mod.rs @@ -39,6 +39,8 @@ pub enum InstanceMessageResponse { }, /// Success response with optional data Success { + /// Protocol identifier (e.g., "consensus/1.0.0", "sync/1.0.0") + protocol: String, /// Response data specific to the protocol data: Option>, }, @@ -49,9 +51,4 @@ pub enum InstanceMessageResponse { /// Error message message: String, }, - /// Protocol-specific response - Protocol { - /// Protocol-specific response data - data: Vec, - }, } diff --git a/crates/networking/src/discovery/behaviour.rs b/crates/networking/src/discovery/behaviour.rs index cabb5cdc7..decc3876c 100644 --- a/crates/networking/src/discovery/behaviour.rs +++ b/crates/networking/src/discovery/behaviour.rs @@ -35,7 +35,7 @@ pub struct DerivedDiscoveryBehaviour { pub identify: identify::Behaviour, /// NAT traversal pub autonat: autonat::Behaviour, - /// UPnP port mapping + /// `UPnP` port mapping pub upnp: Toggle, /// Circuit relay for NAT traversal pub relay: Toggle, @@ -89,18 +89,22 @@ impl DiscoveryBehaviour { } } + #[must_use] pub fn get_peers(&self) -> &HashSet { &self.peers } + #[must_use] pub fn get_peer_info(&self, peer_id: &PeerId) -> Option<&PeerInfo> { self.peer_info.get(peer_id) } + #[must_use] pub fn nat_status(&self) -> autonat::NatStatus { self.discovery.autonat.nat_status() } + #[must_use] pub fn get_peer_addresses(&self) -> HashMap> { self.peer_info .iter() @@ -188,8 +192,8 @@ impl NetworkBehaviour for DiscoveryBehaviour { } } _ => {} - }; - self.discovery.on_swarm_event(event) + } + self.discovery.on_swarm_event(event); } #[allow(clippy::type_complexity)] @@ -271,7 +275,7 @@ impl NetworkBehaviour for DiscoveryBehaviour { // Intentionally ignore } other => { - trace!("Libp2p => Unhandled Kademlia event: {:?}", other) + trace!("Libp2p => Unhandled Kademlia event: {:?}", other); } }, DerivedDiscoveryBehaviourEvent::Mdns(ev) => match ev { diff --git a/crates/networking/src/discovery/config.rs b/crates/networking/src/discovery/config.rs index 7eb869135..94f72ea52 100644 --- a/crates/networking/src/discovery/config.rs +++ b/crates/networking/src/discovery/config.rs @@ -26,7 +26,7 @@ pub struct DiscoveryConfig { enable_mdns: bool, /// Enable Kademlia discovery. enable_kademlia: bool, - /// Enable UPnP discovery. + /// Enable `UPnP` discovery. enable_upnp: bool, /// Enable relay nodes. enable_relay: bool, @@ -64,36 +64,43 @@ impl DiscoveryConfig { self } + #[must_use] pub fn with_bootstrap_peers(mut self, peers: Vec<(PeerId, Multiaddr)>) -> Self { self.bootstrap_peers = peers; self } + #[must_use] pub fn with_relay_nodes(mut self, nodes: Vec<(PeerId, Multiaddr)>) -> Self { self.relay_nodes = nodes; self } + #[must_use] pub fn with_target_peer_count(mut self, count: u64) -> Self { self.target_peer_count = count; self } + #[must_use] pub fn with_mdns(mut self, enable: bool) -> Self { self.enable_mdns = enable; self } + #[must_use] pub fn with_kademlia(mut self, enable: bool) -> Self { self.enable_kademlia = enable; self } + #[must_use] pub fn with_upnp(mut self, enable: bool) -> Self { self.enable_upnp = enable; self } + #[must_use] pub fn with_relay(mut self, enable: bool) -> Self { self.enable_relay = enable; self diff --git a/crates/networking/src/discovery/mod.rs b/crates/networking/src/discovery/mod.rs index c9aa05957..54b3e7f4e 100644 --- a/crates/networking/src/discovery/mod.rs +++ b/crates/networking/src/discovery/mod.rs @@ -13,6 +13,7 @@ pub use peers::{PeerEvent, PeerInfo, PeerManager}; const MAX_ESTABLISHED_PER_PEER: u32 = 4; +#[must_use] pub fn new_kademlia(peer_id: PeerId, protocol: StreamProtocol) -> kad::Behaviour { let store = kad::store::MemoryStore::new(peer_id); let mut config = kad::Config::new(protocol); diff --git a/crates/networking/src/discovery/peers.rs b/crates/networking/src/discovery/peers.rs index ed20b5539..0791e787c 100644 --- a/crates/networking/src/discovery/peers.rs +++ b/crates/networking/src/discovery/peers.rs @@ -1,5 +1,5 @@ use std::{ - collections::{BTreeMap, HashSet}, + collections::HashSet, sync::Arc, time::{Duration, Instant, SystemTime}, }; @@ -7,7 +7,7 @@ use std::{ use crate::InstanceMsgPublicKey; use dashmap::{DashMap, DashSet}; use libp2p::{core::Multiaddr, identify, PeerId}; -use tokio::sync::{broadcast, RwLock}; +use tokio::sync::broadcast; use tracing::debug; /// Information about a peer's connection and behavior @@ -65,9 +65,11 @@ pub struct PeerManager { /// Verified peers from completed handshakes verified_peers: DashSet, /// Handshake keys to peer ids - public_keys_to_peer_ids: DashMap, + public_keys_to_peer_ids: Arc>, /// Banned peers with optional expiration time banned_peers: DashMap>, + /// Allowed public keys + whitelisted_keys: DashSet, /// Event sender for peer updates event_tx: broadcast::Sender, } @@ -80,13 +82,45 @@ impl Default for PeerManager { banned_peers: Default::default(), verified_peers: Default::default(), public_keys_to_peer_ids: Default::default(), + whitelisted_keys: Default::default(), event_tx, } } } impl PeerManager { + #[must_use] + pub fn new(whitelisted_keys: HashSet) -> Self { + let (event_tx, _) = broadcast::channel(100); + Self { + peers: Default::default(), + banned_peers: Default::default(), + verified_peers: Default::default(), + public_keys_to_peer_ids: Default::default(), + whitelisted_keys: DashSet::from_iter(whitelisted_keys), + event_tx, + } + } + + pub fn update_whitelisted_keys(&self, keys: HashSet) { + self.whitelisted_keys.clear(); + for key in keys { + self.whitelisted_keys.insert(key); + } + } + + #[must_use] + pub fn is_key_whitelisted(&self, key: &InstanceMsgPublicKey) -> bool { + self.whitelisted_keys.contains(key) + } + + pub fn handle_nonwhitelisted_peer(&self, peer: &PeerId) { + self.remove_peer(peer, "non-whitelisted"); + self.ban_peer(*peer, "non-whitelisted", None); + } + /// Get a subscription to peer events + #[must_use] pub fn subscribe(&self) -> broadcast::Receiver { self.event_tx.subscribe() } @@ -121,6 +155,7 @@ impl PeerManager { } /// Check if a peer is verified + #[must_use] pub fn is_peer_verified(&self, peer_id: &PeerId) -> bool { self.verified_peers.contains(peer_id) } @@ -147,7 +182,7 @@ impl PeerManager { /// Bans a peer with the default duration(`1h`) pub async fn ban_peer_with_default_duration(&self, peer: PeerId, reason: impl Into) { const BAN_PEER_DURATION: Duration = Duration::from_secs(60 * 60); //1h - self.ban_peer(peer, reason, Some(BAN_PEER_DURATION)) + self.ban_peer(peer, reason, Some(BAN_PEER_DURATION)); } /// Unban a peer @@ -161,6 +196,7 @@ impl PeerManager { } /// Check if a peer is banned + #[must_use] pub fn is_banned(&self, peer_id: &PeerId) -> bool { self.banned_peers.contains_key(peer_id) } @@ -184,16 +220,19 @@ impl PeerManager { } /// Get peer information + #[must_use] pub fn get_peer_info(&self, peer_id: &PeerId) -> Option { self.peers.get(peer_id).map(|info| info.value().clone()) } /// Get all active peers + #[must_use] pub fn get_peers(&self) -> DashMap { self.peers.clone() } /// Get number of active peers + #[must_use] pub fn peer_count(&self) -> usize { self.peers.len() } @@ -225,14 +264,18 @@ impl PeerManager { /// Add a peer id to the public key to peer id map after verifying handshake pub fn add_peer_id_to_public_key(&self, peer_id: &PeerId, public_key: &InstanceMsgPublicKey) { - self.public_keys_to_peer_ids - .insert(public_key.clone(), *peer_id); + self.public_keys_to_peer_ids.insert(*public_key, *peer_id); } /// Remove a peer id from the public key to peer id map pub fn remove_peer_id_from_public_key(&self, peer_id: &PeerId) { self.public_keys_to_peer_ids.retain(|_, id| id != peer_id); } + + #[must_use] + pub fn get_peer_id_from_public_key(&self, public_key: &InstanceMsgPublicKey) -> Option { + self.public_keys_to_peer_ids.get(public_key).map(|id| *id) + } } /// Update the average response time for a peer diff --git a/crates/networking/src/error.rs b/crates/networking/src/error.rs index e2e4b056f..20990b925 100644 --- a/crates/networking/src/error.rs +++ b/crates/networking/src/error.rs @@ -64,4 +64,7 @@ pub enum Error { #[error(transparent)] TokioSendError(#[from] tokio::sync::mpsc::error::SendError), + + #[error(transparent)] + CrossbeamSendError(#[from] crossbeam_channel::SendError), } diff --git a/crates/networking/src/lib.rs b/crates/networking/src/lib.rs index a085802ea..74d3e89dc 100644 --- a/crates/networking/src/lib.rs +++ b/crates/networking/src/lib.rs @@ -6,13 +6,14 @@ pub mod discovery; pub mod error; pub mod handlers; pub mod service; +pub mod service_handle; pub mod types; #[cfg(test)] mod tests; pub use key_types::*; -pub use service::{NetworkConfig, NetworkEvent, NetworkMessage, NetworkService}; +pub use service::{NetworkConfig, NetworkEvent, NetworkService}; #[cfg(all( feature = "sp-core-ecdsa", @@ -99,4 +100,4 @@ compile_error!( pub(crate) trait KeySignExt { fn sign_prehash(&self, prehash: &[u8; 32]) -> InstanceSignedMsgSignature; -} \ No newline at end of file +} diff --git a/crates/networking/src/networking/tests.rs b/crates/networking/src/networking/tests.rs index fc1773976..8b1378917 100644 --- a/crates/networking/src/networking/tests.rs +++ b/crates/networking/src/networking/tests.rs @@ -1,769 +1 @@ -use self::std::time::Duration; -use super::*; -use crate::gossip::GossipHandle; -use futures::{stream, StreamExt}; -use gadget_crypto::hashing::blake3_256; -use gadget_crypto::KeyType; -use gadget_logging::setup_log; -use gadget_std::collections::BTreeMap; -use gadget_std::sync::LazyLock; -use serde::{Deserialize, Serialize}; -use tokio::time::sleep; -const TOPIC: &str = "/gadget/test/1.0.0"; - -fn deserialize<'a, T>(data: &'a [u8]) -> Result -where - T: Deserialize<'a>, -{ - bincode::deserialize(data).map_err(|err| Error::Other(err.to_string())) -} - -#[derive(Serialize, Deserialize, Debug)] -struct StressTestPayload { - value: u64, -} - -#[derive(Debug, Serialize, Deserialize, Clone)] -enum Msg { - Round1(Round1Msg), - Round2(Round2Msg), - Round3(Round3Msg), -} - -#[derive(Debug, Serialize, Deserialize, Clone)] -struct Round1Msg { - pub power: u16, - pub hitpoints: u16, - pub armor: u16, - pub name: String, -} - -#[derive(Debug, Serialize, Deserialize, Clone)] -struct Round2Msg { - pub x: u16, - pub y: u16, - pub z: u16, -} - -#[derive(Debug, Serialize, Deserialize, Clone)] -struct Round3Msg { - rotation: u16, - velocity: (u16, u16, u16), -} - -async fn wait_for_nodes_connected(nodes: &[GossipHandle]) { - let node_count = nodes.len(); - - // wait for the nodes to connect to each other - let max_retries = 10 * node_count; - let mut retry = 0; - loop { - gadget_logging::debug!(%node_count, %max_retries, %retry, "Checking if all nodes are connected to each other"); - let connected = nodes - .iter() - .map(super::super::gossip::GossipHandle::connected_peers) - .collect::>(); - - let all_connected = connected - .iter() - .enumerate() - .inspect(|(node, peers)| { - gadget_logging::debug!("Node {node} has {peers} connected peers"); - }) - .all(|(_, &peers)| peers >= node_count - 1); - if all_connected { - gadget_logging::debug!("All nodes are connected to each other"); - return; - } - sleep(Duration::from_millis(300)).await; - retry += 1; - assert!( - retry <= max_retries, - "Failed to connect all nodes to each other" - ); - } -} - -#[tokio::test(flavor = "multi_thread")] -#[allow(clippy::cast_possible_truncation)] -async fn test_p2p() { - setup_log(); - let nodes = stream::iter(0..*NODE_COUNT) - .map(|_| node()) - .collect::>() - .await; - - wait_for_nodes_connected(&nodes).await; - - let mut mapping = BTreeMap::new(); - for (i, node) in nodes.iter().enumerate() { - mapping.insert(i as u16, node.my_id); - } - - let mut tasks = Vec::new(); - for (i, node) in nodes.into_iter().enumerate() { - let task = tokio::spawn(run_protocol(node, i as u16, mapping.clone())); - tasks.push(task); - } - // Wait for all tasks to finish - let results = futures::future::try_join_all(tasks) - .await - .expect("Failed to run protocol"); - // Assert that all are okay. - assert!( - results.iter().all(std::result::Result::is_ok), - "Some nodes failed to run protocol" - ); -} - -#[allow(clippy::too_many_lines, clippy::cast_possible_truncation)] -async fn run_protocol( - node: N, - i: u16, - mapping: BTreeMap, -) -> Result<(), crate::error::Error> { - let task_hash = [0u8; 32]; - // Safety note: We should be passed a NetworkMultiplexer, and all uses of the N: Network - // used throughout the program must also use the multiplexer to prevent mixed messages. - let multiplexer = NetworkMultiplexer::new(node); - - let round1_network = multiplexer.multiplex(StreamKey { - task_hash, // To differentiate between different instances of a running program (i.e., a task) - round_id: 0, // To differentiate between different subsets of a running task - }); - - let round2_network = multiplexer.multiplex(StreamKey { - task_hash, // To differentiate between different instances of a running program (i.e., a task) - round_id: 1, // To differentiate between different subsets of a running task - }); - - let round3_network = multiplexer.multiplex(StreamKey { - task_hash, // To differentiate between different instances of a running program (i.e., a task) - round_id: 2, // To differentiate between different subsets of a running task - }); - - //let (round1_tx, round1_rx) = node. - // Round 1 (broadcast) - let msg = { - let round = Round1Msg { - power: i * 100, - hitpoints: (i + 1) * 50, - armor: i + 2, - name: format!("Player {}", i), - }; - round1_network.build_protocol_message( - IdentifierInfo { - message_id: 0, - round_id: 0, - }, - i, - None, - &Msg::Round1(round), - None, - ) - }; - - gadget_logging::debug!("Broadcast Message"); - round1_network - .send(msg) - .map_err(|_| crate::error::Error::Other("Failed to send message".into()))?; - - // Wait for all other nodes to send their messages - let mut msgs = BTreeMap::new(); - while let Some(msg) = round1_network.recv().await { - let m = deserialize::(&msg.payload).unwrap(); - gadget_logging::debug!(from = %msg.sender.user_id, ?m, "Received message"); - // Expecting Round1 message - assert!( - matches!(m, Msg::Round1(_)), - "Expected Round1 message but got {:?} from node {}", - m, - msg.sender.user_id, - ); - let old = msgs.insert(msg.sender.user_id, m); - assert!( - old.is_none(), - "Duplicate message from node {}", - msg.sender.user_id, - ); - // Break if all messages are received - if msgs.len() == *NODE_COUNT - 1 { - break; - } - } - gadget_logging::debug!("Done r1 w/ {i}"); - - // Round 2 (P2P) - let msgs = (0..*NODE_COUNT) - .map(|r| r as u16) - .filter(|&j| j != i) - .map(|j| { - let peer_pk = mapping.get(&j).copied().unwrap(); - round2_network.build_protocol_message( - IdentifierInfo { - message_id: 0, - round_id: 0, - }, - i, - Some(j), - &Msg::Round2(Round2Msg { - x: i * 10, - y: (i + 1) * 20, - z: i + 2, - }), - Some(peer_pk), - ) - }) - .collect::>(); - for msg in msgs { - let to = msg - .recipient - .map(|r| r.user_id) - .expect("Recipient should be present for P2P message. This is a bug in the test code"); - gadget_logging::debug!(%to, "Send P2P Message"); - round2_network.send(msg)?; - } - - // Wait for all other nodes to send their messages - let mut msgs = BTreeMap::new(); - while let Some(msg) = round2_network.recv().await { - let m = deserialize::(&msg.payload).unwrap(); - gadget_logging::info!( - "[Node {}] Received message from {} | Intended Recipient: {}", - i, - msg.sender.user_id, - msg.recipient - .as_ref() - .map_or_else(|| "Broadcast".into(), |r| r.user_id.to_string()) - ); - // Expecting Round2 message - assert!( - matches!(m, Msg::Round2(_)), - "Expected Round2 message but got {:?} from node {}", - m, - msg.sender.user_id, - ); - let old = msgs.insert(msg.sender.user_id, m); - assert!( - old.is_none(), - "Duplicate message from node {}", - msg.sender.user_id, - ); - // Break if all messages are received - if msgs.len() == *NODE_COUNT - 1 { - break; - } - } - gadget_logging::debug!("Done r2 w/ {i}"); - - // Round 3 (broadcast) - - let msg = { - let round = Round3Msg { - rotation: i * 30, - velocity: (i + 1, i + 2, i + 3), - }; - round3_network.build_protocol_message( - IdentifierInfo { - message_id: 0, - round_id: 0, - }, - i, - None, - &Msg::Round3(round), - None, - ) - }; - - gadget_logging::debug!("Broadcast Message"); - round3_network.send(msg)?; - - // Wait for all other nodes to send their messages - let mut msgs = BTreeMap::new(); - while let Some(msg) = round3_network.recv().await { - let m = deserialize::(&msg.payload).unwrap(); - gadget_logging::debug!(from = %msg.sender.user_id, ?m, "Received message"); - // Expecting Round3 message - assert!( - matches!(m, Msg::Round3(_)), - "Expected Round3 message but got {:?} from node {}", - m, - msg.sender.user_id, - ); - let old = msgs.insert(msg.sender.user_id, m); - assert!( - old.is_none(), - "Duplicate message from node {}", - msg.sender.user_id, - ); - // Break if all messages are received - if msgs.len() == *NODE_COUNT - 1 { - break; - } - } - gadget_logging::debug!("Done r3 w/ {i}"); - - gadget_logging::info!(node = i, "Protocol completed"); - - Ok(()) -} - -fn node_with_id() -> (GossipHandle, crate::key_types::GossipMsgKeyPair) { - let identity = libp2p::identity::Keypair::generate_ed25519(); - let crypto_key = crate::key_types::Curve::generate_with_seed(None).unwrap(); - let bind_port = 0; - let handle = crate::setup::start_p2p_network(crate::setup::NetworkConfig::new_service_network( - identity, - crypto_key.clone(), - Vec::default(), - bind_port, - TOPIC, - )) - .unwrap(); - - (handle, crypto_key) -} - -fn node() -> GossipHandle { - node_with_id().0 -} - -static NODE_COUNT: LazyLock = - LazyLock::new(|| std::env::var("IN_CI").map_or_else(|_| 10, |_| 2)); -#[allow(dead_code)] -static MESSAGE_COUNT: LazyLock = - LazyLock::new(|| std::env::var("IN_CI").map_or_else(|_| 10, |_| 100)); - -#[tokio::test(flavor = "multi_thread")] -#[allow(clippy::cast_possible_truncation)] -async fn test_stress_test_multiplexer() { - setup_log(); - gadget_logging::info!("Starting test_stress_test_multiplexer"); - - let (network0, network1) = get_networks().await; - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let stream_key = StreamKey { - task_hash: blake3_256(&[1]), - round_id: 0, - }; - - let _subnetwork0 = multiplexer0.multiplex(stream_key); - let _subnetwork1 = multiplexer1.multiplex(stream_key); - - // Create a channel for forwarding - let (forward_tx, mut forward_rx) = tokio::sync::mpsc::unbounded_channel(); - - // Create a subnetwork with forwarding - let subnetwork0 = multiplexer0.multiplex(stream_key); - let subnetwork1 = multiplexer1.multiplex_with_forwarding(stream_key, forward_tx); - - let payload = StressTestPayload { value: 42 }; - let msg = subnetwork0.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), - &payload, - Some(subnetwork1.public_id()), - ); - - gadget_logging::info!("Sending message from subnetwork0"); - subnetwork0.send(msg.clone()).unwrap(); - - // Message should be forwarded to the forward_rx channel - let forwarded_msg = forward_rx.recv().await.unwrap(); - let received: StressTestPayload = deserialize(&forwarded_msg.payload).unwrap(); - assert_eq!(received.value, payload.value); -} - -#[tokio::test(flavor = "multi_thread")] -async fn test_nested_multiplexer() { - setup_log(); - gadget_logging::info!("Starting test_nested_multiplexer"); - let (network0, network1) = get_networks().await; - - nested_multiplex(0, 10, network0, network1).await; -} - -async fn get_networks() -> (GossipHandle, GossipHandle) { - let network0 = node(); - let network1 = node(); - - let mut gossip_networks = vec![network0, network1]; - - wait_for_nodes_connected(&gossip_networks).await; - - (gossip_networks.remove(0), gossip_networks.remove(0)) -} - -async fn nested_multiplex( - cur_depth: usize, - max_depth: usize, - network0: N, - network1: N, -) { - gadget_logging::info!("At nested depth = {cur_depth}/{max_depth}"); - - if cur_depth == max_depth { - return; - } - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let stream_key = StreamKey { - #[allow(clippy::cast_possible_truncation)] - task_hash: blake3_256(&[(cur_depth % 255) as u8]), - round_id: 0, - }; - - let subnetwork0 = multiplexer0.multiplex(stream_key); - let subnetwork1 = multiplexer1.multiplex(stream_key); - let subnetwork1_id = subnetwork1.public_id(); - - // Send a message in the subnetwork0 to subnetwork1 and vice versa, assert values of message - let payload = StressTestPayload { value: 42 }; - let msg = subnetwork0.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), - &payload, - Some(subnetwork1_id), - ); - - gadget_logging::info!("Sending message from subnetwork0"); - subnetwork0.send(msg.clone()).unwrap(); - - // Receive message - let received_msg = subnetwork1.recv().await.unwrap(); - let received: StressTestPayload = deserialize(&received_msg.payload).unwrap(); - assert_eq!(received.value, payload.value); - - let msg = subnetwork1.build_protocol_message( - IdentifierInfo::default(), - 1, - Some(0), - &payload, - Some(subnetwork0.public_id()), - ); - - gadget_logging::info!("Sending message from subnetwork1"); - subnetwork1.send(msg.clone()).unwrap(); - - // Receive message - let received_msg = subnetwork0.recv().await.unwrap(); - let received: StressTestPayload = deserialize(&received_msg.payload).unwrap(); - assert_eq!(received.value, payload.value); - tracing::info!("Done nested depth = {cur_depth}/{max_depth}"); - - Box::pin(nested_multiplex( - cur_depth + 1, - max_depth, - subnetwork0, - subnetwork1, - )) - .await; -} - -#[tokio::test(flavor = "multi_thread")] -async fn test_closed_channel_handling() { - setup_log(); - let (network0, network1) = get_networks().await; - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let stream_key = StreamKey { - task_hash: blake3_256(&[1]), - round_id: 0, - }; - - let subnetwork0 = multiplexer0.multiplex(stream_key); - // Drop subnetwork1's receiver to simulate closed channel - let subnetwork1 = multiplexer1.multiplex(stream_key); - drop(subnetwork1); - - let payload = StressTestPayload { value: 42 }; - let msg = - subnetwork0.build_protocol_message(IdentifierInfo::default(), 0, None, &payload, None); - - // Sending to a closed channel should return an error - assert!(subnetwork0.send(msg).is_ok()); // Changed to ok() since the message will be sent but not received -} - -#[tokio::test(flavor = "multi_thread")] -async fn test_empty_payload() { - setup_log(); - let (network0, network1) = get_networks().await; - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let stream_key = StreamKey { - task_hash: blake3_256(&[1]), - round_id: 0, - }; - - let subnetwork0 = multiplexer0.multiplex(stream_key); - let subnetwork1 = multiplexer1.multiplex(stream_key); - - // Test empty payload - let empty_payload = StressTestPayload { value: 0 }; - let msg = subnetwork0.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), - &empty_payload, - Some(subnetwork1.public_id()), - ); - - gadget_logging::info!("Sending message from subnetwork0"); - subnetwork0.send(msg).unwrap(); - - // Receive message - let received_msg = subnetwork1.recv().await.unwrap(); - let received: StressTestPayload = deserialize(&received_msg.payload).unwrap(); - assert_eq!(received.value, empty_payload.value); -} - -#[tokio::test(flavor = "multi_thread")] -#[allow(clippy::cast_possible_truncation)] -async fn test_concurrent_messaging() { - setup_log(); - let (network0, network1) = get_networks().await; - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let mut send_handles = Vec::new(); - let mut receive_handles = Vec::new(); - - // Create multiple messages to send concurrently - let message_count = 10; - - // Spawn tasks to send messages - for i in 0..message_count { - let stream_key = StreamKey { - task_hash: blake3_256(&[i]), - round_id: 0, - }; - - let subnetwork0 = multiplexer0.multiplex(stream_key); - let subnetwork1 = multiplexer1.multiplex(stream_key); - let subnetwork1_id = subnetwork1.public_id(); - - let i_u64: u64 = i.into(); - let payload = StressTestPayload { value: i_u64 }; - let send_subnetwork0 = subnetwork0; - let handle = tokio::spawn(async move { - let msg = send_subnetwork0.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), - &payload, - Some(subnetwork1_id), - ); - send_subnetwork0.send(msg).unwrap(); - }); - - send_handles.push(handle); - - // Spawn tasks to receive messages - let handle = tokio::spawn(async move { - let msg = subnetwork1.recv().await.unwrap(); - let received: StressTestPayload = deserialize(&msg.payload).unwrap(); - received.value as u8 // Return the payload value for verification - }); - - receive_handles.push(handle); - } - - // Wait for all sends to complete - for handle in send_handles { - handle.await.unwrap(); - } - - // Wait for all receives and verify we got all messages - let mut received_values = Vec::new(); - for handle in receive_handles { - received_values.push(handle.await.unwrap()); - } - - received_values.sort_unstable(); - assert_eq!(received_values.len(), message_count as usize); - for i in 0..message_count { - assert_eq!(received_values[i as usize], i); - } -} - -#[tokio::test(flavor = "multi_thread")] -#[allow(clippy::cast_possible_truncation)] -async fn test_message_ordering() { - setup_log(); - let (network0, network1) = get_networks().await; - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let stream_key = StreamKey { - task_hash: blake3_256(&[1]), - round_id: 0, - }; - - let subnetwork0 = multiplexer0.multiplex(stream_key); - let subnetwork1 = multiplexer1.multiplex(stream_key); - - // Send messages with sequential sequence numbers - let message_count = 10; - for i in 0..message_count { - let payload = StressTestPayload { value: i }; - let msg = subnetwork0.build_protocol_message( - IdentifierInfo { - message_id: i, - ..Default::default() - }, - 0, - Some(1), - &payload, - Some(subnetwork1.public_id()), - ); - subnetwork0.send(msg).unwrap(); - } - - // Verify messages are received in order - let mut last_seq = 0; - for _ in 0..message_count { - let msg = subnetwork1.recv().await.unwrap(); - assert!( - msg.identifier_info.message_id >= last_seq, - "Messages should be received in order or equal to last sequence number" - ); - last_seq = msg.identifier_info.message_id; - } -} - -#[tokio::test(flavor = "multi_thread")] -async fn test_network_id_handling() { - setup_log(); - let (network0, network1) = get_networks().await; - let _network0_id = network0.public_id(); - let network1_id = network1.public_id(); - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - let stream_key = StreamKey { - task_hash: blake3_256(&[1]), - round_id: 0, - }; - - let subnetwork0 = multiplexer0.multiplex(stream_key); - let subnetwork1 = multiplexer1.multiplex(stream_key); - - // Test sending with correct network ID - let payload = StressTestPayload { value: 42 }; - let msg = subnetwork0.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), - &payload, - Some(network1_id), - ); - gadget_logging::info!("Sending message from subnetwork0"); - subnetwork0.send(msg.clone()).unwrap(); - - // Receive message - let received_msg = subnetwork1.recv().await.unwrap(); - let received: StressTestPayload = deserialize(&received_msg.payload).unwrap(); - assert_eq!(received.value, payload.value); - - // Test sending with wrong network ID - let wrong_key = crate::key_types::Curve::generate_with_seed(None).unwrap(); - let msg = subnetwork0.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), - &payload, - Some(wrong_key.public()), - ); - gadget_logging::info!("Sending message from subnetwork0"); - subnetwork0.send(msg).unwrap(); - - // Message with wrong network ID should not be received - let timeout = tokio::time::sleep(tokio::time::Duration::from_millis(100)); - tokio::select! { - () = timeout => (), - _ = subnetwork1.recv() => panic!("Should not receive message with wrong network ID"), - } -} - -#[tokio::test(flavor = "multi_thread")] -async fn test_stream_isolation() { - setup_log(); - let (network0, network1) = get_networks().await; - - let multiplexer0 = NetworkMultiplexer::new(network0); - let multiplexer1 = NetworkMultiplexer::new(network1); - - // Create two different stream keys - let stream_key1 = StreamKey { - task_hash: blake3_256(&[1]), - round_id: 0, - }; - let stream_key2 = StreamKey { - task_hash: blake3_256(&[2]), - round_id: 0, - }; - - let subnetwork0_stream1 = multiplexer0.multiplex(stream_key1); - let subnetwork0_stream2 = multiplexer0.multiplex(stream_key2); - let subnetwork1_stream1 = multiplexer1.multiplex(stream_key1); - let subnetwork1_stream2 = multiplexer1.multiplex(stream_key2); - - // Send messages on both streams - let payload1 = StressTestPayload { value: 1 }; - let payload2 = StressTestPayload { value: 2 }; - - let msg1 = subnetwork0_stream1.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), // Send to node 1 - &payload1, - Some(subnetwork1_stream1.public_id()), - ); - let msg2 = subnetwork0_stream2.build_protocol_message( - IdentifierInfo::default(), - 0, - Some(1), // Send to node 1 - &payload2, - Some(subnetwork1_stream2.public_id()), - ); - - gadget_logging::info!("Sending message from subnetwork0_stream1"); - subnetwork0_stream1.send(msg1.clone()).unwrap(); - gadget_logging::info!("Sending message from subnetwork0_stream2"); - subnetwork0_stream2.send(msg2.clone()).unwrap(); - - // Verify messages are received on correct streams - gadget_logging::info!("Waiting for message on subnetwork1_stream1"); - let received_msg1 = subnetwork1_stream1.recv().await.unwrap(); - gadget_logging::info!("Waiting for message on subnetwork1_stream2"); - let received_msg2 = subnetwork1_stream2.recv().await.unwrap(); - - let received1: StressTestPayload = deserialize(&received_msg1.payload).unwrap(); - let received2: StressTestPayload = deserialize(&received_msg2.payload).unwrap(); - - assert_eq!(received1.value, payload1.value); - assert_eq!(received2.value, payload2.value); - - // Verify no cross-stream message leakage - let timeout = tokio::time::sleep(tokio::time::Duration::from_millis(100)); - tokio::select! { - () = timeout => (), - _ = subnetwork1_stream1.recv() => panic!("Should not receive more messages on stream 1"), - _ = subnetwork1_stream2.recv() => panic!("Should not receive more messages on stream 2"), - } -} diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index 87ea5b787..3939ca299 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -8,26 +8,23 @@ use crate::{ behaviours::{GadgetBehaviour, GadgetBehaviourEvent}, blueprint_protocol::{BlueprintProtocolEvent, InstanceMessageRequest, InstanceMessageResponse}, discovery::{ - behaviour::{DerivedDiscoveryBehaviour, DerivedDiscoveryBehaviourEvent, DiscoveryEvent}, + behaviour::{DerivedDiscoveryBehaviourEvent, DiscoveryEvent}, PeerInfo, PeerManager, }, error::Error, key_types::{InstanceMsgKeyPair, InstanceMsgPublicKey, InstanceSignedMsgSignature}, + service_handle::NetworkServiceHandle, + types::ProtocolMessage, }; -use futures::{Stream, StreamExt}; +use crossbeam_channel::{self, Receiver, Sender}; +use dashmap::DashMap; +use futures::StreamExt; use gadget_logging::trace; use libp2p::{ - core::{transport::Boxed, upgrade}, - gossipsub::{IdentTopic, Topic}, - identify, - identity::Keypair, - noise, ping, - swarm::{ConnectionId, SwarmEvent}, - tcp, yamux, Multiaddr, PeerId, Swarm, SwarmBuilder, Transport, + identify, identity::Keypair, ping, swarm::SwarmEvent, Multiaddr, PeerId, Swarm, SwarmBuilder, + Transport, }; -use tokio::sync::mpsc; -use tokio_stream::wrappers::{IntervalStream, UnboundedReceiverStream}; -use tracing::{debug, error, info, warn}; +use tracing::{debug, info, warn}; /// Events emitted by the network service #[derive(Debug)] @@ -70,6 +67,7 @@ pub enum NetworkEvent { HandshakeFailed { peer: PeerId, reason: String }, } +/// Network message types #[derive(Debug)] pub enum NetworkMessage { InstanceRequest { @@ -128,13 +126,17 @@ pub struct NetworkService { /// Peer manager for tracking peer states peer_manager: Arc, /// Channel for sending messages to the network service - network_sender: mpsc::UnboundedSender, + network_sender: Sender, /// Channel for receiving messages from the network service - network_receiver: UnboundedReceiverStream, + network_receiver: Receiver, + /// Channel for sending messages to the network service + protocol_message_sender: Sender, + /// Channel for receiving messages from the network service + protocol_message_receiver: Receiver, /// Channel for sending events to the network service - event_sender: mpsc::UnboundedSender, + event_sender: Sender, /// Channel for receiving events from the network service - event_receiver: UnboundedReceiverStream, + event_receiver: Receiver, /// Network name/namespace network_name: String, /// Bootstrap peers @@ -143,7 +145,11 @@ pub struct NetworkService { impl NetworkService { /// Create a new network service - pub async fn new(config: NetworkConfig) -> Result { + pub async fn new( + config: NetworkConfig, + allowed_keys: HashSet, + allowed_keys_rx: Receiver>, + ) -> Result { let NetworkConfig { network_name, instance_id, @@ -157,10 +163,13 @@ impl NetworkService { enable_kademlia, } = config; - let peer_manager = Arc::new(PeerManager::default()); - + let peer_manager = Arc::new(PeerManager::new(allowed_keys)); let blueprint_protocol_name = format!("/blueprint_protocol/{}/1.0.0", instance_id); + let (network_sender, network_receiver) = crossbeam_channel::unbounded(); + let (protocol_message_sender, protocol_message_receiver) = crossbeam_channel::unbounded(); + let (event_sender, event_receiver) = crossbeam_channel::unbounded(); + // Create the swarm let behaviour = GadgetBehaviour::new( &network_name, @@ -170,6 +179,7 @@ impl NetworkService { &instance_public_key, target_peer_count, peer_manager.clone(), + protocol_message_sender.clone(), ); let mut swarm = SwarmBuilder::with_existing_identity(local_key) @@ -190,31 +200,49 @@ impl NetworkService { // Start listening swarm.listen_on(listen_addr)?; - - let (network_sender, network_receiver) = mpsc::unbounded_channel(); - let (event_sender, event_receiver) = mpsc::unbounded_channel(); - let bootstrap_peers = bootstrap_peers.into_iter().collect(); - let service = Self { + Ok(Self { swarm, peer_manager, network_sender, - network_receiver: UnboundedReceiverStream::new(network_receiver), - event_sender: event_sender.clone(), - event_receiver: UnboundedReceiverStream::new(event_receiver), + network_receiver, + protocol_message_sender, + protocol_message_receiver, + event_sender, + event_receiver, network_name, bootstrap_peers, - }; - - Ok(service) + }) } /// Get a sender to send messages to the network service - pub fn message_sender(&self) -> mpsc::UnboundedSender { + pub fn network_sender(&self) -> Sender { self.network_sender.clone() } + pub fn start(self) -> NetworkServiceHandle { + let local_peer_id = *self.swarm.local_peer_id(); + let public_keys_to_peer_ids = Arc::new(DashMap::new()); + let network_sender = self.network_sender.clone(); + let protocol_message_receiver = self.protocol_message_receiver.clone(); + + // Create handle with new interface + let handle = NetworkServiceHandle::new( + local_peer_id, + public_keys_to_peer_ids.clone(), + network_sender, + protocol_message_receiver, + ); + + // Spawn background task + tokio::spawn(async move { + self.run().await; + }); + + handle + } + /// Run the network service async fn run(mut self) { info!("Starting network service"); @@ -232,49 +260,35 @@ impl NetworkService { } } - let mut swarm_stream = self.swarm.fuse(); - let mut network_stream = self.network_receiver.fuse(); loop { tokio::select! { - message = network_stream.next() => { - match message { - Some(msg) => match handle_network_message( - swarm_stream.get_mut(), - msg, - &self.peer_manager, - &self.event_sender, - ) - .await - { - Ok(_) => {} - Err(e) => { - warn!("Failed to handle network message: {}", e); - } - }, - None => break, - } - } - swarm_event = swarm_stream.next() => match swarm_event { - // outbound events - Some(SwarmEvent::Behaviour(event)) => { - match handle_behaviour_event( - swarm_stream.get_mut(), + swarm_event = self.swarm.select_next_some() => { + if let SwarmEvent::Behaviour(event) = swarm_event { + if let Err(e) = handle_behaviour_event( + &mut self.swarm, &self.peer_manager, event, &self.event_sender, - &self.network_sender, ) .await { - Ok(_) => {} - Err(e) => { - warn!("Failed to handle swarm event: {}", e); - } + warn!("Failed to handle swarm event: {}", e); } - }, - None => { break; }, - _ => { }, - }, + } + } + Ok(msg) = async { self.network_receiver.try_recv() } => { + if let Err(e) = handle_network_message( + &mut self.swarm, + msg, + &self.peer_manager, + &self.event_sender, + ) + .await + { + warn!("Failed to handle network message: {}", e); + } + } + else => break, } } @@ -287,21 +301,10 @@ async fn handle_swarm_event( swarm: &mut Swarm, peer_manager: &Arc, event: SwarmEvent, - event_sender: &mpsc::UnboundedSender, - network_sender: &mpsc::UnboundedSender, + event_sender: &Sender, ) -> Result<(), Error> { - match event { - SwarmEvent::Behaviour(behaviour_event) => { - handle_behaviour_event( - swarm, - peer_manager, - behaviour_event, - event_sender, - network_sender, - ) - .await? - } - _ => {} + if let SwarmEvent::Behaviour(behaviour_event) = event { + handle_behaviour_event(swarm, peer_manager, behaviour_event, event_sender).await?; } Ok(()) @@ -312,8 +315,7 @@ async fn handle_behaviour_event( swarm: &mut Swarm, peer_manager: &Arc, event: GadgetBehaviourEvent, - event_sender: &mpsc::UnboundedSender, - network_sender: &mpsc::UnboundedSender, + event_sender: &Sender, ) -> Result<(), Error> { match event { GadgetBehaviourEvent::ConnectionLimits(_) => {} @@ -323,30 +325,16 @@ async fn handle_behaviour_event( peer_manager, discovery_event, event_sender, - network_sender, &swarm.behaviour().blueprint_protocol.blueprint_protocol_name, ) - .await? + .await?; } GadgetBehaviourEvent::BlueprintProtocol(blueprint_event) => { - handle_blueprint_protocol_event( - swarm, - peer_manager, - blueprint_event, - event_sender, - network_sender, - ) - .await? + handle_blueprint_protocol_event(swarm, peer_manager, blueprint_event, event_sender) + .await?; } GadgetBehaviourEvent::Ping(ping_event) => { - handle_ping_event( - swarm, - peer_manager, - ping_event, - event_sender, - network_sender, - ) - .await? + handle_ping_event(swarm, peer_manager, ping_event, event_sender).await?; } } @@ -358,8 +346,7 @@ async fn handle_discovery_event( peer_info_map: &HashMap, peer_manager: &Arc, event: DiscoveryEvent, - event_sender: &mpsc::UnboundedSender, - network_sender: &mpsc::UnboundedSender, + event_sender: &Sender, blueprint_protocol_name: &str, ) -> Result<(), Error> { match event { @@ -378,7 +365,7 @@ async fn handle_discovery_event( .. }) => { let protocols: HashSet = - HashSet::from_iter(info.protocols.iter().map(|p| p.to_string())); + HashSet::from_iter(info.protocols.iter().map(std::string::ToString::to_string)); if !protocols.contains(blueprint_protocol_name) { peer_manager .ban_peer_with_default_duration(*peer_id, "hello protocol unsupported"); @@ -387,7 +374,7 @@ async fn handle_discovery_event( DerivedDiscoveryBehaviourEvent::Identify(_) => {} _ => {} }, - }; + } Ok(()) } @@ -397,8 +384,7 @@ async fn handle_blueprint_protocol_event( swarm: &mut Swarm, peer_manager: &Arc, event: BlueprintProtocolEvent, - event_sender: &mpsc::UnboundedSender, - network_sender: &mpsc::UnboundedSender, + event_sender: &Sender, ) -> Result<(), Error> { match event { BlueprintProtocolEvent::Request { @@ -430,8 +416,7 @@ async fn handle_ping_event( swarm: &mut Swarm, peer_manager: &Arc, event: ping::Event, - event_sender: &mpsc::UnboundedSender, - network_sender: &mpsc::UnboundedSender, + event_sender: &Sender, ) -> Result<(), Error> { match event.result { Ok(rtt) => { @@ -460,31 +445,7 @@ async fn handle_network_message( swarm: &mut Swarm, msg: NetworkMessage, peer_manager: &Arc, - event_sender: &mpsc::UnboundedSender, + event_sender: &Sender, ) -> Result<(), Error> { - match msg { - NetworkMessage::InstanceRequest { peer, request } => { - event_sender.send(NetworkEvent::InstanceRequestOutbound { peer, request })? - } - NetworkMessage::InstanceResponse { peer, response } => { - event_sender.send(NetworkEvent::InstanceResponseOutbound { peer, response })? - } - NetworkMessage::GossipMessage { - source, - topic, - message, - } => event_sender.send(NetworkEvent::GossipSent { topic, message })?, - NetworkMessage::HandshakeRequest { - peer, - public_key, - signature, - } => event_sender.send(NetworkEvent::HandshakeCompleted { peer })?, - NetworkMessage::HandshakeResponse { - peer, - public_key, - signature, - } => event_sender.send(NetworkEvent::HandshakeCompleted { peer })?, - } - Ok(()) } diff --git a/crates/networking/src/service_handle.rs b/crates/networking/src/service_handle.rs new file mode 100644 index 000000000..0dd153767 --- /dev/null +++ b/crates/networking/src/service_handle.rs @@ -0,0 +1,134 @@ +use crate::{ + blueprint_protocol::InstanceMessageRequest, key_types::InstanceMsgPublicKey, + service::NetworkMessage, types::ProtocolMessage, +}; +use crossbeam_channel::{self, Receiver, Sender}; +use dashmap::DashMap; +use gadget_logging::info; +use libp2p::PeerId; +use std::sync::Arc; +use tokio::task::JoinHandle; + +/// Handle for sending outgoing messages to the network +#[derive(Clone)] +pub struct NetworkSender { + network_message_sender: Sender, +} + +impl NetworkSender { + #[must_use] + pub fn new(network_message_sender: Sender) -> Self { + Self { + network_message_sender, + } + } + + /// Send a protocol message over the network + pub fn send_message(&self, message: NetworkMessage) -> Result<(), String> { + self.network_message_sender + .send(message) + .map_err(|e| e.to_string()) + } +} + +/// Handle for receiving incoming messages from the network +pub struct NetworkReceiver { + protocol_message_receiver: Receiver, +} + +impl NetworkReceiver { + #[must_use] + pub fn new(protocol_message_receiver: Receiver) -> Self { + Self { + protocol_message_receiver, + } + } + + /// Get the next protocol message + pub fn try_recv(&self) -> Result { + self.protocol_message_receiver.try_recv() + } +} + +/// Combined handle for the network service +pub struct NetworkServiceHandle { + pub local_peer_id: PeerId, + pub sender: NetworkSender, + pub receiver: NetworkReceiver, + pub public_keys_to_peer_ids: Arc>, +} + +impl Clone for NetworkServiceHandle { + fn clone(&self) -> Self { + Self { + local_peer_id: self.local_peer_id, + sender: self.sender.clone(), + receiver: NetworkReceiver::new(self.receiver.protocol_message_receiver.clone()), + public_keys_to_peer_ids: self.public_keys_to_peer_ids.clone(), + } + } +} + +impl NetworkServiceHandle { + #[must_use] + pub fn new( + local_peer_id: PeerId, + public_keys_to_peer_ids: Arc>, + network_message_sender: Sender, + protocol_message_receiver: Receiver, + ) -> Self { + Self { + local_peer_id, + sender: NetworkSender::new(network_message_sender), + receiver: NetworkReceiver::new(protocol_message_receiver), + public_keys_to_peer_ids, + } + } + + pub async fn next_protocol_message(&mut self) -> Option { + self.receiver.try_recv().ok() + } + + pub fn send_protocol_message(&self, message: ProtocolMessage) -> Result<(), String> { + let raw_payload = bincode::serialize(&message).map_err(|err| err.to_string())?; + if message.routing.recipient.is_some() { + let instance_message_request = InstanceMessageRequest::Protocol { + protocol: message.protocol, + payload: raw_payload, + metadata: None, + }; + let recipient = message.routing.recipient.unwrap(); + if let Some(public_key) = recipient.public_key { + if let Some(peer_id) = self.public_keys_to_peer_ids.get(&public_key) { + self.sender.send_message(NetworkMessage::InstanceRequest { + peer: *peer_id, + request: instance_message_request, + })?; + info!("Sent outbound p2p `NetworkMessage` to {:?}", peer_id); + } + } + } else { + let gossip_message = NetworkMessage::GossipMessage { + source: self.local_peer_id, + topic: message.protocol, + message: raw_payload, + }; + self.sender.send_message(gossip_message)?; + info!("Sent outbound gossip `NetworkMessage`"); + } + + Ok(()) + } + + /// Split the handle into separate sender and receiver + #[must_use] + pub fn split(self) -> (NetworkSender, NetworkReceiver) { + (self.sender, self.receiver) + } +} + +/// We might also bundle a `JoinHandle` so the user can await its completion if needed. +pub struct NetworkServiceTaskHandle { + /// The join handle for the background service task. + pub service_task: JoinHandle<()>, +} diff --git a/crates/networking/src/types.rs b/crates/networking/src/types.rs index 63b7e5e07..6c1e4f4f9 100644 --- a/crates/networking/src/types.rs +++ b/crates/networking/src/types.rs @@ -10,9 +10,9 @@ pub const MAX_MESSAGE_SIZE: usize = 16 * 1024 * 1024; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct ParticipantId(pub u16); -impl Into for ParticipantId { - fn into(self) -> u16 { - self.0 +impl From for u16 { + fn from(val: ParticipantId) -> Self { + val.0 } } @@ -56,6 +56,8 @@ pub struct ParticipantInfo { /// A protocol message that can be sent over the network #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ProtocolMessage { + /// The protocol name + pub protocol: String, /// Routing information for the message pub routing: MessageRouting, /// The actual message payload From 8d9bbcb826358e6515b7ad1c2a5a47061f31eae9 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 00:37:25 -0700 Subject: [PATCH 12/52] chore: get 2 tests passing peers connecting, not identifying --- Cargo.lock | 3 + crates/networking/Cargo.toml | 5 + crates/networking/src/behaviours.rs | 10 + .../src/blueprint_protocol/behaviour.rs | 2 +- crates/networking/src/discovery/behaviour.rs | 77 ++++++-- crates/networking/src/discovery/config.rs | 2 +- crates/networking/src/service.rs | 88 +++++++-- crates/networking/src/service_handle.rs | 47 +++-- .../networking/src/tests/discovery_tests.rs | 181 ++++++++++++++++++ crates/networking/src/tests/mod.rs | 4 + crates/networking/src/tests/test_helpers.rs | 150 +++++++++++++++ 11 files changed, 525 insertions(+), 44 deletions(-) create mode 100644 crates/networking/src/tests/discovery_tests.rs create mode 100644 crates/networking/src/tests/test_helpers.rs diff --git a/Cargo.lock b/Cargo.lock index 99a12fe70..1a0d4918a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7163,9 +7163,12 @@ dependencies = [ "blake3", "crossbeam-channel", "dashmap", + "fastrand", "futures", "gadget-crypto", + "gadget-crypto-core", "gadget-logging", + "gadget-networking", "gadget-std", "hex", "itertools 0.14.0", diff --git a/crates/networking/Cargo.toml b/crates/networking/Cargo.toml index 24c466b3e..27efe4b66 100644 --- a/crates/networking/Cargo.toml +++ b/crates/networking/Cargo.toml @@ -36,6 +36,7 @@ crossbeam-channel = { workspace = true } # Crypto dependencies gadget-crypto = { workspace = true, features = ["k256", "hashing"] } +gadget-crypto-core = { workspace = true, features = ["k256"] } k256 = { workspace = true } # Round-based protocol support @@ -65,8 +66,12 @@ features = [ ] [dev-dependencies] +gadget-networking = { workspace = true, features = ["sp-core-ecdsa"] } +gadget-crypto = { workspace = true, features = ["sp-core"] } +gadget-crypto-core = { workspace = true, features = ["tangle"] } tracing-subscriber = { workspace = true } lazy_static = { workspace = true } +fastrand = "2.0" [features] default = ["std"] diff --git a/crates/networking/src/behaviours.rs b/crates/networking/src/behaviours.rs index a5fa02ba7..236940eb1 100644 --- a/crates/networking/src/behaviours.rs +++ b/crates/networking/src/behaviours.rs @@ -22,6 +22,7 @@ use std::{ sync::Arc, time::Duration, }; +use tracing::{debug, info}; const MAX_ESTABLISHED_PER_PEER: u32 = 4; @@ -79,6 +80,10 @@ impl GadgetBehaviour { let ping = ping::Behaviour::new(ping::Config::new().with_interval(Duration::from_secs(30))); + info!( + "Setting up discovery behavior with network name: {}", + network_name + ); let discovery = DiscoveryConfig::new(local_key.public(), network_name) .with_mdns(true) .with_kademlia(true) @@ -86,6 +91,10 @@ impl GadgetBehaviour { .build() .unwrap(); + info!( + "Setting up blueprint protocol with name: {}", + blueprint_protocol_name + ); let blueprint_protocol = BlueprintProtocolBehaviour::new( local_key, instance_secret_key, @@ -95,6 +104,7 @@ impl GadgetBehaviour { protocol_message_sender, ); + debug!("Created GadgetBehaviour with all components initialized"); Self { connection_limits, discovery, diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index fc7f10334..1a0f4e318 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -5,7 +5,6 @@ use crate::{ use crossbeam_channel::Sender; use dashmap::DashMap; use gadget_crypto::{hashing::blake3_256, KeyType}; -use gadget_logging::{debug, info, trace, warn}; use libp2p::{ core::transport::PortUse, gossipsub::{self, IdentTopic, MessageAuthenticity, MessageId, Sha256Topic}, @@ -22,6 +21,7 @@ use std::{ task::Poll, time::{Duration, Instant}, }; +use tracing::{debug, info, trace, warn}; use crate::discovery::PeerManager; diff --git a/crates/networking/src/discovery/behaviour.rs b/crates/networking/src/discovery/behaviour.rs index decc3876c..10c0670f1 100644 --- a/crates/networking/src/discovery/behaviour.rs +++ b/crates/networking/src/discovery/behaviour.rs @@ -5,14 +5,13 @@ use std::{ time::Duration, }; -use gadget_logging::trace; use libp2p::{ autonat, core::Multiaddr, identify, identity::PeerId, - kad::{self, store::MemoryStore}, - mdns::{tokio::Behaviour as Mdns, Event as MdnsEvent}, + kad::{self, store::MemoryStore, Event as KademliaEvent}, + mdns::{self, Event as MdnsEvent}, relay, swarm::{ behaviour::toggle::Toggle, derive_prelude::*, dial_opts::DialOpts, NetworkBehaviour, @@ -21,6 +20,7 @@ use libp2p::{ upnp, }; use tokio::time::Interval; +use tracing::trace; use tracing::{debug, info}; use super::PeerInfo; @@ -30,7 +30,7 @@ pub struct DerivedDiscoveryBehaviour { /// Kademlia discovery pub kademlia: Toggle>, /// Local network discovery via mDNS - pub mdns: Toggle, + pub mdns: Toggle, /// Identify protocol for peer information exchange pub identify: identify::Behaviour, /// NAT traversal @@ -124,6 +124,7 @@ impl NetworkBehaviour for DiscoveryBehaviour { local_addr: &Multiaddr, remote_addr: &Multiaddr, ) -> Result, ConnectionDenied> { + debug!(%peer, "Handling inbound connection"); self.discovery.handle_established_inbound_connection( connection_id, peer, @@ -140,16 +141,20 @@ impl NetworkBehaviour for DiscoveryBehaviour { role_override: libp2p::core::Endpoint, port_use: PortUse, ) -> Result, ConnectionDenied> { + debug!(%peer, "Handling outbound connection"); self.peer_info .entry(peer) - .or_insert_with(|| PeerInfo { - addresses: HashSet::new(), - identify_info: None, - last_seen: std::time::SystemTime::now(), - ping_latency: None, - successes: 0, - failures: 0, - average_response_time: None, + .or_insert_with(|| { + debug!(%peer, "Creating new peer info for outbound connection"); + PeerInfo { + addresses: HashSet::new(), + identify_info: None, + last_seen: std::time::SystemTime::now(), + ping_latency: None, + successes: 0, + failures: 0, + average_response_time: None, + } }) .addresses .insert(addr.clone()); @@ -168,6 +173,7 @@ impl NetworkBehaviour for DiscoveryBehaviour { connection: ConnectionId, event: THandlerOutEvent, ) { + debug!(%peer_id, "Handling connection handler event"); self.discovery .on_connection_handler_event(peer_id, connection, event); } @@ -176,6 +182,7 @@ impl NetworkBehaviour for DiscoveryBehaviour { match &event { FromSwarm::ConnectionEstablished(e) => { if e.other_established == 0 { + debug!(%e.peer_id, "First connection established with peer"); self.n_node_connected += 1; self.peers.insert(e.peer_id); self.pending_events @@ -184,6 +191,7 @@ impl NetworkBehaviour for DiscoveryBehaviour { } FromSwarm::ConnectionClosed(e) => { if e.remaining_established == 0 { + debug!(%e.peer_id, "Last connection closed with peer"); self.n_node_connected -= 1; self.peers.remove(&e.peer_id); self.peer_info.remove(&e.peer_id); @@ -240,17 +248,46 @@ impl NetworkBehaviour for DiscoveryBehaviour { match ev { ToSwarm::GenerateEvent(ev) => { match &ev { - DerivedDiscoveryBehaviourEvent::Identify(ev) => { - if let identify::Event::Received { peer_id, info, .. } = ev { - self.peer_info.entry(*peer_id).or_default().identify_info = - Some(info.clone()); - if let Some(kademlia) = self.discovery.kademlia.as_mut() { - for address in &info.listen_addrs { - kademlia.add_address(peer_id, address.clone()); - } + DerivedDiscoveryBehaviourEvent::Identify(identify::Event::Received { + peer_id, + info, + connection_id, + }) => { + debug!(%peer_id, "Received identify event in discovery behaviour"); + self.peer_info.entry(*peer_id).or_default().identify_info = + Some(info.clone()); + if let Some(kademlia) = self.discovery.kademlia.as_mut() { + for address in &info.listen_addrs { + kademlia.add_address(peer_id, address.clone()); } } + self.pending_events + .push_back(DiscoveryEvent::Discovery(Box::new( + DerivedDiscoveryBehaviourEvent::Identify( + identify::Event::Received { + peer_id: *peer_id, + info: info.clone(), + connection_id: *connection_id, + }, + ), + ))); + } + DerivedDiscoveryBehaviourEvent::Identify(identify::Event::Sent { + .. + }) => { + debug!("Identify event sent"); } + DerivedDiscoveryBehaviourEvent::Identify(identify::Event::Pushed { + .. + }) => { + debug!("Identify event pushed"); + } + DerivedDiscoveryBehaviourEvent::Identify(identify::Event::Error { + .. + }) => { + debug!("Identify event error"); + } + DerivedDiscoveryBehaviourEvent::Autonat(_) => {} DerivedDiscoveryBehaviourEvent::Upnp(ev) => match ev { upnp::Event::NewExternalAddr(addr) => { diff --git a/crates/networking/src/discovery/config.rs b/crates/networking/src/discovery/config.rs index 94f72ea52..727311bb1 100644 --- a/crates/networking/src/discovery/config.rs +++ b/crates/networking/src/discovery/config.rs @@ -2,7 +2,6 @@ use super::{ behaviour::{DerivedDiscoveryBehaviour, DiscoveryBehaviour}, new_kademlia, }; -use gadget_logging::warn; use libp2p::{ autonat, identify, identity::PublicKey, mdns, relay, upnp, Multiaddr, PeerId, StreamProtocol, }; @@ -10,6 +9,7 @@ use std::{ collections::{HashMap, HashSet, VecDeque}, time::Duration, }; +use tracing::warn; pub struct DiscoveryConfig { /// The local peer ID. diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index 3939ca299..b3f5143eb 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -17,13 +17,12 @@ use crate::{ types::ProtocolMessage, }; use crossbeam_channel::{self, Receiver, Sender}; -use dashmap::DashMap; use futures::StreamExt; -use gadget_logging::trace; use libp2p::{ - identify, identity::Keypair, ping, swarm::SwarmEvent, Multiaddr, PeerId, Swarm, SwarmBuilder, - Transport, + identify, identity::Keypair, kad, mdns, ping, swarm::SwarmEvent, Multiaddr, PeerId, Swarm, + SwarmBuilder, }; +use tracing::trace; use tracing::{debug, info, warn}; /// Events emitted by the network service @@ -124,7 +123,7 @@ pub struct NetworkService { /// The libp2p swarm swarm: Swarm, /// Peer manager for tracking peer states - peer_manager: Arc, + pub(crate) peer_manager: Arc, /// Channel for sending messages to the network service network_sender: Sender, /// Channel for receiving messages from the network service @@ -223,18 +222,24 @@ impl NetworkService { pub fn start(self) -> NetworkServiceHandle { let local_peer_id = *self.swarm.local_peer_id(); - let public_keys_to_peer_ids = Arc::new(DashMap::new()); let network_sender = self.network_sender.clone(); let protocol_message_receiver = self.protocol_message_receiver.clone(); // Create handle with new interface let handle = NetworkServiceHandle::new( local_peer_id, - public_keys_to_peer_ids.clone(), + self.peer_manager.clone(), network_sender, protocol_message_receiver, ); + // Add our own peer ID to the peer manager with all listening addresses + let mut info = PeerInfo::default(); + for addr in self.swarm.listeners() { + info.addresses.insert(addr.clone()); + } + self.peer_manager.update_peer(local_peer_id, info); + // Spawn background task tokio::spawn(async move { self.run().await; @@ -294,6 +299,15 @@ impl NetworkService { info!("Network service stopped"); } + + /// Get the current listening address + pub async fn get_listen_addr(&self) -> Option { + if let Some(addr) = self.swarm.listeners().next() { + Some(addr.clone()) + } else { + None + } + } } /// Handle a swarm event @@ -351,11 +365,16 @@ async fn handle_discovery_event( ) -> Result<(), Error> { match event { DiscoveryEvent::PeerConnected(peer_id) => { - trace!("Peer connected, {peer_id}"); + info!("Peer connected, {peer_id}"); + // Update peer info when connected + if let Some(info) = peer_info_map.get(&peer_id) { + peer_manager.update_peer(peer_id, info.clone()); + } event_sender.send(NetworkEvent::PeerConnected(peer_id))?; } DiscoveryEvent::PeerDisconnected(peer_id) => { - trace!("Peer disconnected, {peer_id}"); + info!("Peer disconnected, {peer_id}"); + peer_manager.remove_peer(&peer_id, "disconnected"); event_sender.send(NetworkEvent::PeerDisconnected(peer_id))?; } DiscoveryEvent::Discovery(discovery_event) => match &*discovery_event { @@ -364,14 +383,61 @@ async fn handle_discovery_event( info, .. }) => { + info!(%peer_id, "Received identify event"); let protocols: HashSet = HashSet::from_iter(info.protocols.iter().map(std::string::ToString::to_string)); + + debug!(%peer_id, ?protocols, "Supported protocols"); + if !protocols.contains(blueprint_protocol_name) { + warn!(%peer_id, %blueprint_protocol_name, "Peer does not support required protocol"); peer_manager - .ban_peer_with_default_duration(*peer_id, "hello protocol unsupported"); + .ban_peer_with_default_duration(*peer_id, "protocol unsupported") + .await; + return Ok(()); + } + + // Get existing peer info or create new one + let mut peer_info = peer_manager.get_peer_info(peer_id).unwrap_or_default(); + + // Update identify info + peer_info.identify_info = Some(info.clone()); + + debug!(%peer_id, listen_addrs=?info.listen_addrs, "Adding identify addresses"); + // Add all addresses from identify info + for addr in &info.listen_addrs { + peer_info.addresses.insert(addr.clone()); + } + + debug!(%peer_id, "Updating peer info with identify information"); + peer_manager.update_peer(*peer_id, peer_info); + info!(%peer_id, "Successfully processed identify information"); + } + DerivedDiscoveryBehaviourEvent::Identify(_) => { + // Ignore other identify events + } + DerivedDiscoveryBehaviourEvent::Kademlia(kad::Event::OutboundQueryProgressed { + result: kad::QueryResult::GetClosestPeers(Ok(ok)), + .. + }) => { + // Process newly discovered peers + for peer_info in ok.peers.iter() { + if !peer_manager.get_peers().contains_key(&peer_info.peer_id) { + let info = PeerInfo::default(); + peer_manager.update_peer(peer_info.peer_id, info); + } + } + } + DerivedDiscoveryBehaviourEvent::Mdns(mdns::Event::Discovered(list)) => { + // Add newly discovered peers from mDNS + for (peer_id, addr) in list { + if !peer_manager.get_peers().contains_key(peer_id) { + let mut info = PeerInfo::default(); + info.addresses.insert(addr.clone()); + peer_manager.update_peer(*peer_id, info); + } } } - DerivedDiscoveryBehaviourEvent::Identify(_) => {} _ => {} }, } diff --git a/crates/networking/src/service_handle.rs b/crates/networking/src/service_handle.rs index 0dd153767..04b33043a 100644 --- a/crates/networking/src/service_handle.rs +++ b/crates/networking/src/service_handle.rs @@ -1,13 +1,14 @@ use crate::{ - blueprint_protocol::InstanceMessageRequest, key_types::InstanceMsgPublicKey, - service::NetworkMessage, types::ProtocolMessage, + blueprint_protocol::InstanceMessageRequest, + discovery::{PeerInfo, PeerManager}, + service::NetworkMessage, + types::ProtocolMessage, }; use crossbeam_channel::{self, Receiver, Sender}; -use dashmap::DashMap; -use gadget_logging::info; -use libp2p::PeerId; +use libp2p::{Multiaddr, PeerId}; use std::sync::Arc; use tokio::task::JoinHandle; +use tracing::info; /// Handle for sending outgoing messages to the network #[derive(Clone)] @@ -55,7 +56,7 @@ pub struct NetworkServiceHandle { pub local_peer_id: PeerId, pub sender: NetworkSender, pub receiver: NetworkReceiver, - pub public_keys_to_peer_ids: Arc>, + pub peer_manager: Arc, } impl Clone for NetworkServiceHandle { @@ -64,7 +65,7 @@ impl Clone for NetworkServiceHandle { local_peer_id: self.local_peer_id, sender: self.sender.clone(), receiver: NetworkReceiver::new(self.receiver.protocol_message_receiver.clone()), - public_keys_to_peer_ids: self.public_keys_to_peer_ids.clone(), + peer_manager: self.peer_manager.clone(), } } } @@ -73,7 +74,7 @@ impl NetworkServiceHandle { #[must_use] pub fn new( local_peer_id: PeerId, - public_keys_to_peer_ids: Arc>, + peer_manager: Arc, network_message_sender: Sender, protocol_message_receiver: Receiver, ) -> Self { @@ -81,7 +82,7 @@ impl NetworkServiceHandle { local_peer_id, sender: NetworkSender::new(network_message_sender), receiver: NetworkReceiver::new(protocol_message_receiver), - public_keys_to_peer_ids, + peer_manager, } } @@ -89,6 +90,20 @@ impl NetworkServiceHandle { self.receiver.try_recv().ok() } + pub fn peers(&self) -> Vec { + self.peer_manager + .get_peers() + .clone() + .into_read_only() + .iter() + .map(|(peer_id, _)| *peer_id) + .collect() + } + + pub fn peer_info(&self, peer_id: &PeerId) -> Option { + self.peer_manager.get_peer_info(peer_id) + } + pub fn send_protocol_message(&self, message: ProtocolMessage) -> Result<(), String> { let raw_payload = bincode::serialize(&message).map_err(|err| err.to_string())?; if message.routing.recipient.is_some() { @@ -99,9 +114,9 @@ impl NetworkServiceHandle { }; let recipient = message.routing.recipient.unwrap(); if let Some(public_key) = recipient.public_key { - if let Some(peer_id) = self.public_keys_to_peer_ids.get(&public_key) { + if let Some(peer_id) = self.peer_manager.get_peer_id_from_public_key(&public_key) { self.sender.send_message(NetworkMessage::InstanceRequest { - peer: *peer_id, + peer: peer_id, request: instance_message_request, })?; info!("Sent outbound p2p `NetworkMessage` to {:?}", peer_id); @@ -120,6 +135,16 @@ impl NetworkServiceHandle { Ok(()) } + pub async fn get_listen_addr(&self) -> Option { + // Get the first peer info for our local peer ID + if let Some(peer_info) = self.peer_manager.get_peer_info(&self.local_peer_id) { + // Return the first address from our peer info + peer_info.addresses.iter().next().cloned() + } else { + None + } + } + /// Split the handle into separate sender and receiver #[must_use] pub fn split(self) -> (NetworkSender, NetworkReceiver) { diff --git a/crates/networking/src/tests/discovery_tests.rs b/crates/networking/src/tests/discovery_tests.rs new file mode 100644 index 000000000..50f8dd5c9 --- /dev/null +++ b/crates/networking/src/tests/discovery_tests.rs @@ -0,0 +1,181 @@ +use super::test_helpers::TestNode; +use std::{collections::HashSet, time::Duration}; +use tokio::time::timeout; +use tracing_subscriber::{fmt, EnvFilter}; + +#[cfg(test)] +mod tests { + use tracing::{debug, info}; + + use super::*; + + fn init_tracing() { + let _ = fmt() + .with_env_filter(EnvFilter::from_default_env()) + .with_target(true) + .with_thread_ids(true) + .with_file(true) + .with_line_number(true) + .try_init(); + } + + #[tokio::test] + async fn test_peer_discovery_mdns() { + let network_name = "test-network"; + let instance_id = "test-instance"; + let allowed_keys = HashSet::new(); + + // Create two nodes + let mut node1 = + TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; + let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; + + // Start both nodes and wait for them to be listening + let handle1 = node1.start().await.expect("Failed to start node1"); + let handle2 = node2.start().await.expect("Failed to start node2"); + + // Wait for peer discovery (with timeout) + let discovery_timeout = Duration::from_secs(10); + match timeout(discovery_timeout, async { + loop { + let peers1 = handle1.peers(); + let peers2 = handle2.peers(); + + if peers1.contains(&node2.peer_id) && peers2.contains(&node1.peer_id) { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + { + Ok(_) => println!("Peers discovered each other successfully"), + Err(_) => panic!("Peer discovery timed out"), + } + } + + #[tokio::test] + async fn test_peer_discovery_kademlia() { + let network_name = "test-network"; + let instance_id = "test-instance"; + let allowed_keys = HashSet::new(); + + // Create the first node (bootstrap node) + let mut node1 = + TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; + + // Start node1 and get its listening address + let handle1 = node1.start().await.expect("Failed to start node1"); + let node1_addr = node1.get_listen_addr().expect("Node1 should be listening"); + + // Create two more nodes that will bootstrap from node1 + let bootstrap_peers = vec![(node1.peer_id, node1_addr.clone())]; + let mut node2 = TestNode::new( + network_name, + instance_id, + allowed_keys.clone(), + bootstrap_peers.clone(), + ) + .await; + let mut node3 = + TestNode::new(network_name, instance_id, allowed_keys, bootstrap_peers).await; + + // Start the remaining nodes + let handle2 = node2.start().await.expect("Failed to start node2"); + let handle3 = node3.start().await.expect("Failed to start node3"); + + // Wait for peer discovery through Kademlia DHT + let discovery_timeout = Duration::from_secs(20); + match timeout(discovery_timeout, async { + loop { + let peers1 = handle1.peers(); + let peers2 = handle2.peers(); + let peers3 = handle3.peers(); + + if peers1.contains(&node2.peer_id) + && peers1.contains(&node3.peer_id) + && peers2.contains(&node1.peer_id) + && peers2.contains(&node3.peer_id) + && peers3.contains(&node1.peer_id) + && peers3.contains(&node2.peer_id) + { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + { + Ok(_) => println!("All peers discovered each other through Kademlia"), + Err(_) => panic!("Kademlia peer discovery timed out"), + } + } + + #[tokio::test] + async fn test_peer_info_updates() { + init_tracing(); + + let network_name = "test-network"; + let instance_id = "test-instance"; + let allowed_keys = HashSet::new(); + + info!("Creating test nodes..."); + // Create two nodes + let mut node1 = + TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; + let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; + + info!("Starting nodes..."); + // Start both nodes + let handle1 = node1.start().await.expect("Failed to start node1"); + let handle2 = node2.start().await.expect("Failed to start node2"); + + info!("Waiting for peer discovery..."); + // First wait for basic peer discovery (they see each other) + let discovery_timeout = Duration::from_secs(20); + timeout(discovery_timeout, async { + loop { + let peers1 = handle1.peers(); + let peers2 = handle2.peers(); + + debug!("Node1 peers: {:?}, Node2 peers: {:?}", peers1, peers2); + + if peers1.contains(&node2.peer_id) && peers2.contains(&node1.peer_id) { + info!("Basic peer discovery successful"); + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Basic peer discovery timed out"); + + info!("Waiting for identify info..."); + // Now wait for identify info to be populated + let identify_timeout = Duration::from_secs(20); + match timeout(identify_timeout, async { + loop { + let peer_info1 = handle1.peer_info(&node2.peer_id); + let peer_info2 = handle2.peer_info(&node1.peer_id); + + if let Some(peer_info) = peer_info1 { + if peer_info.identify_info.is_some() { + // Also verify reverse direction + if let Some(peer_info) = peer_info2 { + if peer_info.identify_info.is_some() { + info!("Identify info populated in both directions"); + break; + } + } + } + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + { + Ok(_) => info!("Peer info updated successfully in both directions"), + Err(_) => panic!("Peer info update timed out"), + } + } +} diff --git a/crates/networking/src/tests/mod.rs b/crates/networking/src/tests/mod.rs index 8b1378917..993863136 100644 --- a/crates/networking/src/tests/mod.rs +++ b/crates/networking/src/tests/mod.rs @@ -1 +1,5 @@ +mod discovery_tests; +mod test_helpers; +// Re-export test helpers for use in other test modules +pub(crate) use test_helpers::*; diff --git a/crates/networking/src/tests/test_helpers.rs b/crates/networking/src/tests/test_helpers.rs new file mode 100644 index 000000000..a0cadf8f1 --- /dev/null +++ b/crates/networking/src/tests/test_helpers.rs @@ -0,0 +1,150 @@ +use std::{collections::HashSet, mem, time::Duration}; + +use gadget_crypto::{sp_core::SpEcdsa, KeyType}; +use libp2p::{identity, Multiaddr, PeerId}; +use tokio::time::timeout; +use tracing::info; + +use crate::{ + key_types::InstanceMsgPublicKey, service_handle::NetworkServiceHandle, NetworkConfig, + NetworkService, +}; + +/// Test node configuration for network tests +pub struct TestNode { + pub service: Option, + pub peer_id: PeerId, + pub listen_addr: Option, +} + +impl TestNode { + pub async fn new( + network_name: &str, + instance_id: &str, + allowed_keys: HashSet, + bootstrap_peers: Vec<(PeerId, Multiaddr)>, + ) -> Self { + let local_key = identity::Keypair::generate_ed25519(); + let peer_id = local_key.public().to_peer_id(); + + let listen_addr: Multiaddr = format!("/ip4/127.0.0.1/tcp/0").parse().unwrap(); + info!( + "Creating test node {} with TCP address: {}", + peer_id, listen_addr + ); + + let instance_secret_key = SpEcdsa::generate_with_seed(None).unwrap(); + let instance_public_key = instance_secret_key.public(); + + let config = NetworkConfig { + network_name: network_name.to_string(), + instance_id: instance_id.to_string(), + instance_secret_key, + instance_public_key, + local_key, + listen_addr: listen_addr.clone(), + target_peer_count: 10, + bootstrap_peers, + enable_mdns: true, + enable_kademlia: true, + }; + + let (allowed_keys_tx, allowed_keys_rx) = crossbeam_channel::unbounded(); + allowed_keys_tx.send(allowed_keys.clone()).unwrap(); + let service = NetworkService::new(config, allowed_keys, allowed_keys_rx) + .await + .expect("Failed to create network service"); + + Self { + service: Some(service), + peer_id, + listen_addr: Some(listen_addr), + } + } + + /// Start the node + pub async fn start(&mut self) -> Result { + // Take ownership of the service + let service = self.service.take().ok_or("Service already started")?; + let handle = service.start(); + + // Wait for the actual listening address + let timeout_duration = Duration::from_secs(5); + match timeout(timeout_duration, async { + while self.listen_addr.is_none() { + if let Some(addr) = handle.get_listen_addr().await { + info!("Node {} listening on {}", self.peer_id, addr); + self.listen_addr = Some(addr); + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + { + Ok(_) => Ok(handle), + Err(_) => Err("Timeout waiting for node to start listening"), + } + } + + /// Get the actual listening address + pub fn get_listen_addr(&self) -> Option { + self.listen_addr.clone() + } +} + +/// Wait for a condition with timeout +pub async fn wait_for_condition(timeout: Duration, mut condition: F) -> Result<(), &'static str> +where + F: FnMut() -> bool, +{ + let start = std::time::Instant::now(); + while !condition() { + if start.elapsed() > timeout { + return Err("Timeout waiting for condition"); + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + Ok(()) +} + +/// Wait for peers to discover each other +pub async fn wait_for_peer_discovery( + nodes: &[&TestNode], + timeout: Duration, +) -> Result<(), &'static str> { + wait_for_condition(timeout, || { + for (i, node1) in nodes.iter().enumerate() { + for (j, node2) in nodes.iter().enumerate() { + if i != j + && !node1 + .service + .as_ref() + .map(|s| s.peer_manager.get_peers()) + .unwrap_or_default() + .contains_key(&node2.peer_id) + { + return false; + } + } + } + true + }) + .await +} + +/// Wait for peer info to be updated +pub async fn wait_for_peer_info( + node: &TestNode, + peer_id: &PeerId, + timeout: Duration, +) -> Result<(), &'static str> { + wait_for_condition(timeout, || { + node.service + .as_ref() + .map(|s| s.peer_manager.get_peer_info(peer_id)) + .unwrap_or(None) + .map_or(false, |info| info.identify_info.is_some()) + }) + .await +} From d29a40da593b043178f082ea6332dbc78eb85038 Mon Sep 17 00:00:00 2001 From: Shady Khalifa Date: Wed, 19 Feb 2025 16:24:04 +0200 Subject: [PATCH 13/52] fix: it builds --- Cargo.lock | 18 --- Cargo.toml | 11 +- crates/clients/Cargo.toml | 7 +- crates/clients/networking/Cargo.toml | 34 ----- crates/clients/networking/src/error.rs | 22 ---- crates/clients/networking/src/lib.rs | 4 - crates/clients/networking/src/p2p.rs | 117 ------------------ crates/config/src/context_config.rs | 29 +++++ crates/config/src/lib.rs | 73 ++++++++--- crates/contexts/Cargo.toml | 16 +-- .../src/lib.rs | 61 +++------ crates/networking/src/service.rs | 31 +++-- crates/networking/src/service_handle.rs | 8 +- crates/sdk/Cargo.toml | 34 ++--- flake.lock | 18 +-- flake.nix | 28 ++++- 16 files changed, 172 insertions(+), 339 deletions(-) delete mode 100644 crates/clients/networking/Cargo.toml delete mode 100644 crates/clients/networking/src/error.rs delete mode 100644 crates/clients/networking/src/lib.rs delete mode 100644 crates/clients/networking/src/p2p.rs diff --git a/Cargo.lock b/Cargo.lock index 1a0d4918a..c6156e988 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2627,7 +2627,6 @@ dependencies = [ "blueprint-build-utils", "blueprint-metadata", "eigensdk", - "gadget-client-networking", "gadget-clients", "gadget-config", "gadget-context-derive", @@ -6647,22 +6646,6 @@ dependencies = [ "url", ] -[[package]] -name = "gadget-client-networking" -version = "0.1.0" -dependencies = [ - "gadget-client-core", - "gadget-config", - "gadget-crypto", - "gadget-logging", - "gadget-networking", - "gadget-std", - "libp2p", - "serde", - "serde_json", - "thiserror 2.0.11", -] - [[package]] name = "gadget-client-tangle" version = "0.1.0" @@ -6692,7 +6675,6 @@ dependencies = [ "gadget-client-core", "gadget-client-eigenlayer", "gadget-client-evm", - "gadget-client-networking", "gadget-client-tangle", "gadget-std", "thiserror 2.0.11", diff --git a/Cargo.toml b/Cargo.toml index 36aa1d505..adc070995 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,13 +1,7 @@ [workspace] resolver = "2" -members = [ - "cli", - "blueprints/*", - "crates/*", -] -exclude = [ - "blueprints/incredible-squaring-symbiotic", -] +members = ["cli", "blueprints/*", "crates/*"] +exclude = ["blueprints/incredible-squaring-symbiotic"] [workspace.package] authors = ["Tangle Network"] @@ -68,7 +62,6 @@ gadget-clients = { version = "0.1.0", path = "./crates/clients", default-feature gadget-client-core = { version = "0.1.0", path = "./crates/clients/core", default-features = false } gadget-client-eigenlayer = { version = "0.1.0", path = "./crates/clients/eigenlayer", default-features = false } gadget-client-evm = { version = "0.1.0", path = "./crates/clients/evm", default-features = false } -gadget-client-networking = { version = "0.1.0", path = "./crates/clients/networking", default-features = false } gadget-client-tangle = { version = "0.1.0", path = "./crates/clients/tangle", default-features = false } gadget-contexts = { version = "0.1.0", path = "./crates/contexts", default-features = false } diff --git a/crates/clients/Cargo.toml b/crates/clients/Cargo.toml index b080eb982..dc31ed001 100644 --- a/crates/clients/Cargo.toml +++ b/crates/clients/Cargo.toml @@ -12,7 +12,6 @@ repository.workspace = true gadget-client-eigenlayer = { workspace = true, optional = true } gadget-client-evm = { workspace = true, optional = true } gadget-client-tangle = { workspace = true, optional = true } -gadget-client-networking = { workspace = true, optional = true } gadget-client-core = { workspace = true } gadget-std.workspace = true @@ -23,14 +22,12 @@ default = ["std"] std = [ "gadget-client-eigenlayer?/std", "gadget-client-evm?/std", - "gadget-client-networking?/std", "gadget-client-tangle?/std", "gadget-std/std", - "thiserror/std" + "thiserror/std", ] web = ["gadget-client-tangle?/web"] eigenlayer = ["dep:gadget-client-eigenlayer"] evm = ["dep:gadget-client-evm"] -networking = ["dep:gadget-client-networking"] -tangle = ["dep:gadget-client-tangle"] \ No newline at end of file +tangle = ["dep:gadget-client-tangle"] diff --git a/crates/clients/networking/Cargo.toml b/crates/clients/networking/Cargo.toml deleted file mode 100644 index 1b37a9963..000000000 --- a/crates/clients/networking/Cargo.toml +++ /dev/null @@ -1,34 +0,0 @@ -[package] -name = "gadget-client-networking" -version = "0.1.0" -description = "Networking client for Tangle Blueprints" -authors.workspace = true -edition.workspace = true -license.workspace = true -homepage.workspace = true -repository.workspace = true - -[dependencies] -gadget-config = { workspace = true, features = ["networking"] } -gadget-crypto = { workspace = true, features = ["k256"] } -gadget-logging = { workspace = true } -gadget-networking = { workspace = true, features = ["round-based-compat"] } -gadget-client-core = { workspace = true } -gadget-std = { workspace = true } -libp2p = { workspace = true } -serde = { workspace = true } -serde_json = { workspace = true, features = ["alloc"] } -thiserror = { workspace = true } - -[features] -default = ["std"] -std = [ - "gadget-config/std", - "gadget-crypto/std", - "gadget-logging/std", - "gadget-client-core/std", - "gadget-networking/std", - "gadget-std/std", - "serde/std", - "serde_json/std", -] \ No newline at end of file diff --git a/crates/clients/networking/src/error.rs b/crates/clients/networking/src/error.rs deleted file mode 100644 index 8d1e18031..000000000 --- a/crates/clients/networking/src/error.rs +++ /dev/null @@ -1,22 +0,0 @@ -use gadget_std::string::String; -use thiserror::Error; - -#[derive(Debug, Error)] -pub enum Error { - #[error("P2P error: {0}")] - P2p(String), - #[error("Transport error: {0}")] - Transport(String), - #[error("Protocol error: {0}")] - Protocol(String), - #[error("Configuration error: {0}")] - Configuration(String), -} - -impl From for gadget_client_core::error::Error { - fn from(value: Error) -> Self { - gadget_client_core::error::Error::Network(value.to_string()) - } -} - -pub type Result = gadget_std::result::Result; diff --git a/crates/clients/networking/src/lib.rs b/crates/clients/networking/src/lib.rs deleted file mode 100644 index 893f5304b..000000000 --- a/crates/clients/networking/src/lib.rs +++ /dev/null @@ -1,4 +0,0 @@ -#![cfg_attr(not(feature = "std"), no_std)] - -pub mod error; -pub mod p2p; diff --git a/crates/clients/networking/src/p2p.rs b/crates/clients/networking/src/p2p.rs deleted file mode 100644 index 6bec7bec7..000000000 --- a/crates/clients/networking/src/p2p.rs +++ /dev/null @@ -1,117 +0,0 @@ -use crate::error::{Error, Result}; -use gadget_config::GadgetConfiguration; -use gadget_networking::gossip::GossipHandle; -use gadget_networking::round_based_compat::NetworkDeliveryWrapper; -use gadget_networking::setup::NetworkConfig; -use gadget_networking::{networking::NetworkMultiplexer, round_based}; -use gadget_networking::{GossipMsgKeyPair, GossipMsgPublicKey}; -use gadget_std::collections::BTreeMap; -use gadget_std::sync::Arc; -use round_based::PartyIndex; - -pub struct P2PClient { - name: String, - config: GadgetConfiguration, - target_port: u16, - gossip_msg_keypair: GossipMsgKeyPair, -} - -impl P2PClient { - pub fn new( - name: String, - config: GadgetConfiguration, - target_port: u16, - gossip_msg_keypair: GossipMsgKeyPair, - ) -> Self { - Self { - name, - config, - target_port, - gossip_msg_keypair, - } - } - - pub fn config(&self) -> &GadgetConfiguration { - &self.config - } - - /// Returns the network protocol identifier - pub fn network_protocol(&self, version: Option) -> String { - let name = self.name.to_lowercase(); - match version { - Some(v) => format!("/{}/{}", name, v), - None => format!("/{}/1.0.0", name), - } - } - - pub fn libp2p_identity(&self, ed25519_seed: Vec) -> Result { - let mut seed_bytes = ed25519_seed; - let keypair = libp2p::identity::Keypair::ed25519_from_bytes(&mut seed_bytes) - .map_err(|err| Error::Configuration(err.to_string()))?; - Ok(keypair) - } - - /// Returns a new `NetworkConfig` for the current environment. - pub fn libp2p_network_config>( - &self, - network_name: T, - ed25519_seed: Vec, - ) -> Result { - let network_identity = self.libp2p_identity(ed25519_seed)?; - let network_config = NetworkConfig::new_service_network( - network_identity, - self.gossip_msg_keypair.clone(), - self.config.bootnodes.clone(), - self.target_port, - network_name, - ); - - Ok(network_config) - } - - /// Starts the P2P network and returns the gossip handle - pub fn start_p2p_network>( - &self, - network_name: T, - ed25519_seed: Vec, - ) -> Result { - let network_config = self.libp2p_network_config(network_name, ed25519_seed)?; - match gadget_networking::setup::start_p2p_network(network_config) { - Ok(handle) => Ok(handle), - Err(err) => { - gadget_logging::error!("Failed to start network: {}", err.to_string()); - Err(Error::Protocol(format!("Failed to start network: {err}"))) - } - } - } - - /// Creates a network multiplexer backend - pub fn create_network_multiplexer>( - &self, - network_name: T, - ed25519_seed: Vec, - ) -> Result> { - let handle = self.start_p2p_network(network_name, ed25519_seed)?; - Ok(Arc::new(NetworkMultiplexer::new(handle))) - } - - /// Creates a network delivery wrapper - pub fn create_network_delivery_wrapper( - &self, - mux: Arc, - party_index: PartyIndex, - task_hash: [u8; 32], - parties: BTreeMap, - ) -> NetworkDeliveryWrapper - where - M: Clone - + Send - + Unpin - + 'static - + serde::Serialize - + serde::de::DeserializeOwned - + round_based::ProtocolMessage, - { - NetworkDeliveryWrapper::new::(mux, party_index, task_hash, parties) - } -} diff --git a/crates/config/src/context_config.rs b/crates/config/src/context_config.rs index d751eb7bc..2ec4e66ba 100644 --- a/crates/config/src/context_config.rs +++ b/crates/config/src/context_config.rs @@ -37,6 +37,18 @@ pub enum GadgetCLICoreSettings { #[arg(long, env)] #[serde(default)] network_bind_port: Option, + #[cfg(feature = "networking")] + #[arg(long, env)] + #[serde(default)] + enable_mdns: bool, + #[cfg(feature = "networking")] + #[arg(long, env)] + #[serde(default)] + enable_kademlia: bool, + #[cfg(feature = "networking")] + #[arg(long, env)] + #[serde(default)] + target_peer_count: Option, #[arg(long, short = 'd', env)] keystore_uri: String, #[arg(long, value_enum, env)] @@ -218,6 +230,12 @@ impl Default for GadgetCLICoreSettings { bootnodes: None, #[cfg(feature = "networking")] network_bind_port: None, + #[cfg(feature = "networking")] + enable_mdns: false, + #[cfg(feature = "networking")] + enable_kademlia: false, + #[cfg(feature = "networking")] + target_peer_count: None, keystore_uri: String::new(), chain: SupportedChains::default(), verbose: 0, @@ -342,6 +360,11 @@ impl ContextConfig { #[cfg(feature = "tangle")] let service_id = tangle_settings.and_then(|s| s.service_id); + #[cfg(feature = "networking")] + let enable_mdns = cfg!(debug_assertions); + #[cfg(feature = "networking")] + let enable_kademlia = !cfg!(debug_assertions); + ContextConfig { gadget_core_settings: GadgetCLICoreSettings::Run { test_mode: false, @@ -350,6 +373,12 @@ impl ContextConfig { bootnodes: None, #[cfg(feature = "networking")] network_bind_port: None, + #[cfg(feature = "networking")] + enable_mdns, + #[cfg(feature = "networking")] + enable_kademlia, + #[cfg(feature = "networking")] + target_peer_count: None, keystore_uri, chain, verbose: 3, diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index e7840107d..f73980135 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -69,6 +69,11 @@ pub enum Error { /// Missing `SymbioticContractAddresses` #[error("Missing SymbioticContractAddresses")] MissingSymbioticContractAddresses, + + #[cfg(feature = "networking")] + #[error(transparent)] + Networking(#[from] gadget_networking::error::Error), + #[error("Bad RPC Connection: {0}")] BadRpcConnection(String), #[error("Configuration error: {0}")] @@ -79,8 +84,9 @@ pub enum Error { pub use networking_imports::*; #[cfg(feature = "networking")] mod networking_imports { - pub use gadget_networking::networking::NetworkMultiplexer; - pub use gadget_networking::setup::start_p2p_network; + // pub use gadget_networking::networking::NetworkMultiplexer; + // pub use gadget_networking::start_p2p_network; + pub use gadget_networking::NetworkConfig; pub use libp2p::Multiaddr; pub use std::sync::Arc; } @@ -111,6 +117,15 @@ pub struct GadgetConfiguration { pub protocol_settings: ProtocolSettings, /// Whether the gadget is in test mode pub test_mode: bool, + /// Whether to enable mDNS + #[cfg(feature = "networking")] + pub enable_mdns: bool, + /// Whether to enable Kademlia + #[cfg(feature = "networking")] + pub enable_kademlia: bool, + /// The target number of peers to connect to + #[cfg(feature = "networking")] + pub target_peer_count: u64, } impl GadgetConfiguration { @@ -118,16 +133,16 @@ impl GadgetConfiguration { pub fn libp2p_start_network( &self, network_name: impl Into, - ) -> Result, Error> { - tracing::info!(target: "gadget", "AB0"); - let network_config = self - .libp2p_network_config(network_name) - .map_err(|err| Error::ConfigurationError(err.to_string()))?; + ) -> Result { + let network_config = self.libp2p_network_config(network_name)?; + // TODO: Add allowed keys + let allowed_keys = Default::default(); + let networking_service = + gadget_networking::NetworkService::new(network_config, allowed_keys)?; + + let handle = networking_service.start(); - tracing::info!(target: "gadget", "AB1"); - start_p2p_network(network_config) - .map_err(|err| Error::ConfigurationError(err.to_string())) - .map(|net| Arc::new(NetworkMultiplexer::new(net))) + Ok(handle) } /// Returns a new `NetworkConfig` for the current environment. @@ -135,7 +150,7 @@ impl GadgetConfiguration { pub fn libp2p_network_config( &self, network_name: impl Into, - ) -> Result { + ) -> Result { use gadget_keystore::backends::Backend; use gadget_keystore::crypto::sp_core::SpEd25519 as LibP2PKeyType; use gadget_networking::key_types::Curve as GossipMsgKeyPair; @@ -159,13 +174,23 @@ impl GadgetConfiguration { .get_secret::(&ecdsa_pub_key) .map_err(|err| Error::ConfigurationError(err.to_string()))?; - let network_config = gadget_networking::setup::NetworkConfig::new_service_network( - network_identity, - ecdsa_pair, - self.bootnodes.clone(), - self.network_bind_port, + let listen_addr: Multiaddr = format!("/ip4/0.0.0.0/tcp/{}", self.network_bind_port) + .parse() + .expect("valid multiaddr; qed"); + + let network_name: String = network_name.into(); + let network_config = gadget_networking::NetworkConfig { + instance_id: network_name.clone(), network_name, - ); + instance_secret_key: ecdsa_pair, + instance_public_key: ecdsa_pub_key, + local_key: network_identity, + listen_addr, + target_peer_count: self.target_peer_count, + bootstrap_peers: self.bootnodes.clone(), + enable_mdns: self.enable_mdns, + enable_kademlia: self.enable_kademlia, + }; Ok(network_config) } @@ -191,6 +216,12 @@ fn load_inner(config: ContextConfig) -> Result { bootnodes, #[cfg(feature = "networking")] network_bind_port, + #[cfg(feature = "networking")] + enable_mdns, + #[cfg(feature = "networking")] + enable_kademlia, + #[cfg(feature = "networking")] + target_peer_count, keystore_uri, protocol, #[cfg(feature = "tangle")] @@ -310,6 +341,12 @@ fn load_inner(config: ContextConfig) -> Result { bootnodes: bootnodes.unwrap_or_default(), #[cfg(feature = "networking")] network_bind_port: network_bind_port.unwrap_or_default(), + #[cfg(feature = "networking")] + enable_mdns, + #[cfg(feature = "networking")] + enable_kademlia, + #[cfg(feature = "networking")] + target_peer_count: target_peer_count.unwrap_or(24), protocol, protocol_settings, }) diff --git a/crates/contexts/Cargo.toml b/crates/contexts/Cargo.toml index 11356861e..9c7e88b84 100644 --- a/crates/contexts/Cargo.toml +++ b/crates/contexts/Cargo.toml @@ -22,19 +22,11 @@ tangle-subxt = { workspace = true, optional = true } [features] default = ["std", "keystore"] -std = [ - "gadget-std/std", - "gadget-clients/std", - "gadget-keystore?/std", - "gadget-networking?/std", - "tangle-subxt?/std", -] -web = [ - "tangle-subxt?/web", -] +std = ["gadget-std/std", "gadget-clients/std", "gadget-keystore?/std", "gadget-networking?/std", "tangle-subxt?/std"] +web = ["tangle-subxt?/web"] evm = ["gadget-clients/evm"] eigenlayer = ["gadget-clients/eigenlayer"] -networking = ["gadget-clients/networking", "dep:gadget-networking", "dep:proc-macro2"] +networking = ["dep:gadget-networking", "dep:proc-macro2"] keystore = ["dep:gadget-config", "dep:gadget-keystore"] -tangle = ["dep:tangle-subxt", "gadget-clients/tangle"] \ No newline at end of file +tangle = ["dep:tangle-subxt", "gadget-clients/tangle"] diff --git a/crates/networking-round-based-extension/src/lib.rs b/crates/networking-round-based-extension/src/lib.rs index ba82b50a5..892a0ba08 100644 --- a/crates/networking-round-based-extension/src/lib.rs +++ b/crates/networking-round-based-extension/src/lib.rs @@ -184,7 +184,6 @@ pub struct RoundBasedReceiver { party_index: PartyIndex, forward_rx: Receiver<(PartyIndex, M)>, _phantom: std::marker::PhantomData, - next_message_future: Option> + Send>>>, } impl RoundBasedReceiver { @@ -198,7 +197,6 @@ impl RoundBasedReceiver { party_index, forward_rx, _phantom: std::marker::PhantomData, - next_message_future: None, } } } @@ -223,50 +221,29 @@ where // Get a mutable reference to self let this = self.get_mut(); - // Create and store the future if we don't have one - if this.next_message_future.is_none() { - let mut handle = this.handle.clone(); - this.next_message_future = - Some(Box::pin( - async move { handle.next_protocol_message().await }, - )); - } - - // Poll the stored future - if let Some(future) = &mut this.next_message_future { - match future.as_mut().poll(cx) { - Poll::Ready(Some(msg)) => { - // Clear the future so we create a new one next time - this.next_message_future = None; + let next_protocol_message = this.handle.next_protocol_message(); + match next_protocol_message { + Some(protocol_message) => { + let msg_type = if protocol_message.routing.recipient.is_some() { + MessageType::P2P + } else { + MessageType::Broadcast + }; - let msg_type = if msg.routing.recipient.is_some() { - MessageType::P2P - } else { - MessageType::Broadcast - }; + let sender = protocol_message.routing.sender.id.0; + let id = protocol_message.routing.message_id; - let sender = msg.routing.sender.id.0; - let id = msg.routing.message_id; - - match serde_json::from_slice(&msg.payload) { - Ok(msg) => Poll::Ready(Some(Ok(Incoming { - msg, - sender, - id, - msg_type, - }))), - Err(e) => Poll::Ready(Some(Err(NetworkError::Serialization(e)))), - } - } - Poll::Ready(None) => { - this.next_message_future = None; - Poll::Ready(None) + match serde_json::from_slice(&protocol_message.payload) { + Ok(msg) => Poll::Ready(Some(Ok(Incoming { + msg, + sender, + id, + msg_type, + }))), + Err(e) => Poll::Ready(Some(Err(NetworkError::Serialization(e)))), } - Poll::Pending => Poll::Pending, } - } else { - // This shouldn't happen because we create the future above if it's None - Poll::Ready(None) + None => Poll::Pending, } } } diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index b3f5143eb..e778eebf0 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -112,7 +112,7 @@ pub struct NetworkConfig { /// Target number of peers to maintain pub target_peer_count: u64, /// Bootstrap peers to connect to - pub bootstrap_peers: Vec<(PeerId, Multiaddr)>, + pub bootstrap_peers: Vec, /// Whether to enable mDNS discovery pub enable_mdns: bool, /// Whether to enable Kademlia DHT @@ -139,15 +139,15 @@ pub struct NetworkService { /// Network name/namespace network_name: String, /// Bootstrap peers - bootstrap_peers: HashMap, + bootstrap_peers: HashSet, } impl NetworkService { /// Create a new network service - pub async fn new( + pub fn new( config: NetworkConfig, allowed_keys: HashSet, - allowed_keys_rx: Receiver>, + // allowed_keys_rx: Receiver>, ) -> Result { let NetworkConfig { network_name, @@ -242,7 +242,7 @@ impl NetworkService { // Spawn background task tokio::spawn(async move { - self.run().await; + Box::pin(self.run()).await; }); handle @@ -258,8 +258,8 @@ impl NetworkService { } // Connect to bootstrap peers - for (peer_id, addr) in &self.bootstrap_peers { - debug!("Dialing bootstrap peer {} at {}", peer_id, addr); + for addr in &self.bootstrap_peers { + debug!("Dialing bootstrap peer at {}", addr); if let Err(e) = self.swarm.dial(addr.clone()) { warn!("Failed to dial bootstrap peer: {}", e); } @@ -301,12 +301,8 @@ impl NetworkService { } /// Get the current listening address - pub async fn get_listen_addr(&self) -> Option { - if let Some(addr) = self.swarm.listeners().next() { - Some(addr.clone()) - } else { - None - } + pub fn get_listen_addr(&self) -> Option { + self.swarm.listeners().next().cloned() } } @@ -384,8 +380,11 @@ async fn handle_discovery_event( .. }) => { info!(%peer_id, "Received identify event"); - let protocols: HashSet = - HashSet::from_iter(info.protocols.iter().map(std::string::ToString::to_string)); + let protocols: HashSet = info + .protocols + .iter() + .map(std::string::ToString::to_string) + .collect(); debug!(%peer_id, ?protocols, "Supported protocols"); @@ -421,7 +420,7 @@ async fn handle_discovery_event( .. }) => { // Process newly discovered peers - for peer_info in ok.peers.iter() { + for peer_info in &ok.peers { if !peer_manager.get_peers().contains_key(&peer_info.peer_id) { let info = PeerInfo::default(); peer_manager.update_peer(peer_info.peer_id, info); diff --git a/crates/networking/src/service_handle.rs b/crates/networking/src/service_handle.rs index 04b33043a..02bf9222e 100644 --- a/crates/networking/src/service_handle.rs +++ b/crates/networking/src/service_handle.rs @@ -86,10 +86,11 @@ impl NetworkServiceHandle { } } - pub async fn next_protocol_message(&mut self) -> Option { + pub fn next_protocol_message(&mut self) -> Option { self.receiver.try_recv().ok() } + #[must_use] pub fn peers(&self) -> Vec { self.peer_manager .get_peers() @@ -100,10 +101,12 @@ impl NetworkServiceHandle { .collect() } + #[must_use] pub fn peer_info(&self, peer_id: &PeerId) -> Option { self.peer_manager.get_peer_info(peer_id) } + #[must_use] pub fn send_protocol_message(&self, message: ProtocolMessage) -> Result<(), String> { let raw_payload = bincode::serialize(&message).map_err(|err| err.to_string())?; if message.routing.recipient.is_some() { @@ -135,7 +138,8 @@ impl NetworkServiceHandle { Ok(()) } - pub async fn get_listen_addr(&self) -> Option { + #[must_use] + pub fn get_listen_addr(&self) -> Option { // Get the first peer info for our local peer ID if let Some(peer_info) = self.peer_manager.get_peer_info(&self.local_peer_id) { // Return the first address from our peer info diff --git a/crates/sdk/Cargo.toml b/crates/sdk/Cargo.toml index 64807d16a..4f0a1efb2 100644 --- a/crates/sdk/Cargo.toml +++ b/crates/sdk/Cargo.toml @@ -44,7 +44,6 @@ tokio = { workspace = true, default-features = false } # Networking-related dependencies gadget-networking = { workspace = true, optional = true } -gadget-client-networking = { workspace = true, optional = true } gadget-contexts = { workspace = true } gadget-context-derive = { workspace = true, optional = true } @@ -82,10 +81,7 @@ web = [ "gadget-macros?/web", ] -macros = [ - "dep:gadget-macros", - "dep:gadget-context-derive", -] +macros = ["dep:gadget-macros", "dep:gadget-context-derive"] build = ["dep:blueprint-metadata", "dep:blueprint-build-utils"] @@ -97,7 +93,7 @@ tangle = [ "gadget-macros?/tangle", "gadget-testing-utils?/tangle", "gadget-utils/tangle", - "gadget-event-listeners/tangle" + "gadget-event-listeners/tangle", ] evm = [ @@ -108,7 +104,7 @@ evm = [ "gadget-utils/evm", "gadget-testing-utils?/anvil", "gadget-macros?/evm", - "gadget-event-listeners/evm" + "gadget-event-listeners/evm", ] eigenlayer = [ @@ -120,40 +116,28 @@ eigenlayer = [ "gadget-testing-utils?/eigenlayer", "gadget-utils/eigenlayer", "gadget-macros?/eigenlayer", - "gadget-event-listeners/evm" + "gadget-event-listeners/evm", ] -testing = [ - "dep:gadget-testing-utils", - "dep:tempfile", - "std", -] +testing = ["dep:gadget-testing-utils", "dep:tempfile", "std"] networking = [ "dep:gadget-networking", - "dep:gadget-client-networking", "gadget-contexts/networking", - "gadget-clients/networking", "gadget-keystore/std", "gadget-config/networking", "gadget-macros?/networking", "gadget-context-derive?/networking", ] -networking-sp-core-ecdsa = [ - "gadget-networking/sp-core-ecdsa" -] +networking-sp-core-ecdsa = ["gadget-networking/sp-core-ecdsa"] -networking-sr25519 = [ - "gadget-networking/sp-core-sr25519" -] +networking-sr25519 = ["gadget-networking/sp-core-sr25519"] -networking-ed25519 = [ - "gadget-networking/sp-core-ed25519" -] +networking-ed25519 = ["gadget-networking/sp-core-ed25519"] local-store = ["gadget-stores/local"] round-based-compat = ["gadget-networking/round-based-compat"] -cronjob = ["gadget-event-listeners/cronjob"] \ No newline at end of file +cronjob = ["gadget-event-listeners/cronjob"] diff --git a/flake.lock b/flake.lock index 63fc2a2c2..0bfa5f257 100644 --- a/flake.lock +++ b/flake.lock @@ -28,11 +28,11 @@ ] }, "locked": { - "lastModified": 1733130721, - "narHash": "sha256-FeL2dez6gE3/u+Dq5Ot4TA0EbonBRo2/7pTYNmBBTNY=", + "lastModified": 1738660302, + "narHash": "sha256-aLWyhJx2cO/M3/QLoDBpsObFfjC9e/VEN6HtaI0U6IA=", "owner": "shazow", "repo": "foundry.nix", - "rev": "277722b2fc52af44c0827c47f0f170c7a9b17b0e", + "rev": "33a209625b9e31227a5f11417e95a3ac7264d811", "type": "github" }, "original": { @@ -44,11 +44,11 @@ }, "nixpkgs": { "locked": { - "lastModified": 1733064805, - "narHash": "sha256-7NbtSLfZO0q7MXPl5hzA0sbVJt6pWxxtGWbaVUDDmjs=", + "lastModified": 1739863612, + "narHash": "sha256-UbtgxplOhFcyjBcNbTVO8+HUHAl/WXFDOb6LvqShiZo=", "owner": "NixOS", "repo": "nixpkgs", - "rev": "31d66ae40417bb13765b0ad75dd200400e98de84", + "rev": "632f04521e847173c54fa72973ec6c39a371211c", "type": "github" }, "original": { @@ -73,11 +73,11 @@ ] }, "locked": { - "lastModified": 1733106880, - "narHash": "sha256-aJmAIjZfWfPSWSExwrYBLRgXVvgF5LP1vaeUGOOIQ98=", + "lastModified": 1739932111, + "narHash": "sha256-WkayjH0vuGw0hx2gmjTUGFRvMKpM17gKcpL/U8EUUw0=", "owner": "oxalica", "repo": "rust-overlay", - "rev": "e66c0d43abf5bdefb664c3583ca8994983c332ae", + "rev": "75b2271c5c087d830684cd5462d4410219acc367", "type": "github" }, "original": { diff --git a/flake.nix b/flake.nix index 1568804b4..59ec63992 100644 --- a/flake.nix +++ b/flake.nix @@ -19,10 +19,21 @@ }; }; - outputs = { self, nixpkgs, rust-overlay, foundry, flake-utils }: - flake-utils.lib.eachDefaultSystem (system: + outputs = + { + self, + nixpkgs, + rust-overlay, + foundry, + flake-utils, + }: + flake-utils.lib.eachDefaultSystem ( + system: let - overlays = [ (import rust-overlay) foundry.overlay ]; + overlays = [ + (import rust-overlay) + foundry.overlay + ]; pkgs = import nixpkgs { inherit system overlays; }; @@ -41,7 +52,6 @@ # Protocol Buffers pkgs.protobuf # Mold Linker for faster builds (only on Linux) - (lib.optionals pkgs.stdenv.isLinux pkgs.om4) (lib.optionals pkgs.stdenv.isLinux pkgs.mold) (lib.optionals pkgs.stdenv.isDarwin pkgs.darwin.apple_sdk.frameworks.Security) (lib.optionals pkgs.stdenv.isDarwin pkgs.darwin.apple_sdk.frameworks.SystemConfiguration) @@ -66,7 +76,13 @@ ]; # Environment variables RUST_SRC_PATH = "${toolchain}/lib/rustlib/src/rust/library"; - LD_LIBRARY_PATH = lib.makeLibraryPath [ pkgs.gmp pkgs.libclang pkgs.openssl.dev pkgs.stdenv.cc.cc ]; + LD_LIBRARY_PATH = lib.makeLibraryPath [ + pkgs.gmp + pkgs.libclang + pkgs.openssl.dev + pkgs.stdenv.cc.cc + ]; }; - }); + } + ); } From f71313a1d2c66141a073b28cf63f4cdd961c83d9 Mon Sep 17 00:00:00 2001 From: Shady Khalifa Date: Wed, 19 Feb 2025 18:28:22 +0200 Subject: [PATCH 14/52] fix: get it to compile with blueprints --- Cargo.lock | 2 +- crates/clients/Cargo.toml | 1 + crates/clients/tangle/src/client.rs | 9 +++++---- crates/config/src/lib.rs | 5 ++--- crates/contexts/src/p2p.rs | 13 ------------- crates/networking/Cargo.toml | 4 ---- crates/networking/src/lib.rs | 3 +++ crates/sdk/Cargo.toml | 5 ++++- crates/sdk/src/lib.rs | 2 ++ 9 files changed, 18 insertions(+), 26 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c6156e988..517005e0c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2638,6 +2638,7 @@ dependencies = [ "gadget-logging", "gadget-macros", "gadget-networking", + "gadget-networking-round-based-extension", "gadget-runners", "gadget-std", "gadget-stores", @@ -7159,7 +7160,6 @@ dependencies = [ "libp2p", "lru-mem", "parking_lot 0.12.3", - "round-based", "serde", "serde_json", "thiserror 2.0.11", diff --git a/crates/clients/Cargo.toml b/crates/clients/Cargo.toml index dc31ed001..eae5fe236 100644 --- a/crates/clients/Cargo.toml +++ b/crates/clients/Cargo.toml @@ -31,3 +31,4 @@ web = ["gadget-client-tangle?/web"] eigenlayer = ["dep:gadget-client-eigenlayer"] evm = ["dep:gadget-client-evm"] tangle = ["dep:gadget-client-tangle"] +networking = [] diff --git a/crates/clients/tangle/src/client.rs b/crates/clients/tangle/src/client.rs index d790568fa..135e00dda 100644 --- a/crates/clients/tangle/src/client.rs +++ b/crates/clients/tangle/src/client.rs @@ -45,8 +45,6 @@ pub struct TangleClient { services_client: TangleServicesClient, } -const KEY_ID: &str = "tangle-default"; - impl TangleClient { /// Create a new Tangle runtime client from an existing [`GadgetConfiguration`]. pub async fn new(config: GadgetConfiguration) -> std::result::Result { @@ -154,7 +152,10 @@ impl TangleClient { Error, > { let parties = self.get_operators().await?; - let my_id = self.keystore.get_public_key_local::(KEY_ID)?; + let my_id = self + .keystore + .first_local::() + .map_err(Error::Keystore)?; gadget_logging::trace!( "Looking for {my_id:?} in parties: {:?}", @@ -303,7 +304,7 @@ impl GadgetServicesClient for TangleClient { async fn operator_id( &self, ) -> std::result::Result { - Ok(self.keystore.get_public_key_local::(KEY_ID)?.0) + Ok(self.keystore.first_local::()?.0) } /// Retrieves the current blueprint ID from the configuration diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index f73980135..89945afd2 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -132,10 +132,9 @@ impl GadgetConfiguration { #[cfg(feature = "networking")] pub fn libp2p_start_network( &self, - network_name: impl Into, + network_config: gadget_networking::NetworkConfig, ) -> Result { - let network_config = self.libp2p_network_config(network_name)?; - // TODO: Add allowed keys + // TODO: Add allowed keys or pass them. let allowed_keys = Default::default(); let networking_service = gadget_networking::NetworkService::new(network_config, allowed_keys)?; diff --git a/crates/contexts/src/p2p.rs b/crates/contexts/src/p2p.rs index 6461f9001..8b1378917 100644 --- a/crates/contexts/src/p2p.rs +++ b/crates/contexts/src/p2p.rs @@ -1,14 +1 @@ -pub use gadget_clients::networking::p2p::P2PClient; -pub use gadget_networking::GossipMsgKeyPair; -pub use gadget_std::net::IpAddr; -pub use proc_macro2; -/// `P2pContext` trait provides access to a peer to peer networking client. -pub trait P2pContext { - fn p2p_client( - &self, - name: gadget_std::string::String, - target_port: u16, - my_ecdsa_key: GossipMsgKeyPair, - ) -> P2PClient; -} diff --git a/crates/networking/Cargo.toml b/crates/networking/Cargo.toml index 27efe4b66..28b1b0e9e 100644 --- a/crates/networking/Cargo.toml +++ b/crates/networking/Cargo.toml @@ -39,8 +39,6 @@ gadget-crypto = { workspace = true, features = ["k256", "hashing"] } gadget-crypto-core = { workspace = true, features = ["k256"] } k256 = { workspace = true } -# Round-based protocol support -round-based = { workspace = true, optional = true } [target.'cfg(not(target_family = "wasm"))'.dependencies.libp2p] workspace = true @@ -86,8 +84,6 @@ std = [ "serde_json/std", ] -round-based-compat = ["dep:round-based"] - # Only one of these features should be enabled at a time. # If none are enabled, k256 ECDSA will be used by default. sp-core-ecdsa = ["gadget-crypto/sp-core"] diff --git a/crates/networking/src/lib.rs b/crates/networking/src/lib.rs index 74d3e89dc..3e921f017 100644 --- a/crates/networking/src/lib.rs +++ b/crates/networking/src/lib.rs @@ -9,6 +9,9 @@ pub mod service; pub mod service_handle; pub mod types; +#[cfg(feature = "round-based-compat")] +pub use gadget_networking_round_based_extension as round_based_compat; + #[cfg(test)] mod tests; diff --git a/crates/sdk/Cargo.toml b/crates/sdk/Cargo.toml index 4f0a1efb2..c574ff830 100644 --- a/crates/sdk/Cargo.toml +++ b/crates/sdk/Cargo.toml @@ -46,6 +46,9 @@ tokio = { workspace = true, default-features = false } gadget-networking = { workspace = true, optional = true } gadget-contexts = { workspace = true } gadget-context-derive = { workspace = true, optional = true } +# Round-based protocol support +gadget-networking-round-based-extension = { workspace = true, optional = true } + # Optional dependencies for testing gadget-testing-utils = { workspace = true, optional = true } @@ -138,6 +141,6 @@ networking-ed25519 = ["gadget-networking/sp-core-ed25519"] local-store = ["gadget-stores/local"] -round-based-compat = ["gadget-networking/round-based-compat"] +round-based-compat = ["dep:gadget-networking-round-based-extension"] cronjob = ["gadget-event-listeners/cronjob"] diff --git a/crates/sdk/src/lib.rs b/crates/sdk/src/lib.rs index d157535ba..9c31c4f19 100644 --- a/crates/sdk/src/lib.rs +++ b/crates/sdk/src/lib.rs @@ -33,6 +33,8 @@ pub mod build { pub mod networking { /// Networking utilities for blueprints pub use gadget_networking::*; + #[cfg(feature = "round-based-compat")] + pub use gadget_networking_round_based_extension as round_based_compat; } /// Event listener infrastructure for handling blueprint events From c2bb2f197edfd975c9675187404a89b5b57cd8ad Mon Sep 17 00:00:00 2001 From: Shady Khalifa Date: Wed, 19 Feb 2025 18:47:29 +0200 Subject: [PATCH 15/52] fix: use blake2 vs blake3 for hashing --- crates/networking/src/blueprint_protocol/behaviour.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index 1a0f4e318..6f681de61 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -4,7 +4,8 @@ use crate::{ }; use crossbeam_channel::Sender; use dashmap::DashMap; -use gadget_crypto::{hashing::blake3_256, KeyType}; +use gadget_crypto::tangle_pair_signer::sp_core::blake2_256; +use gadget_crypto::KeyType; use libp2p::{ core::transport::PortUse, gossipsub::{self, IdentTopic, MessageAuthenticity, MessageId, Sha256Topic}, @@ -143,7 +144,7 @@ impl BlueprintProtocolBehaviour { /// Sign a handshake message for a peer pub(crate) fn sign_handshake(&self, peer: &PeerId) -> InstanceSignedMsgSignature { let msg = peer.to_bytes(); - let msg_hash = blake3_256(&msg); + let msg_hash = blake2_256(&msg); self.instance_secret_key.sign_prehash(&msg_hash) } From 42836672c9c8dd5e86567ea152d95575b79f71b7 Mon Sep 17 00:00:00 2001 From: Alex <69764315+Serial-ATA@users.noreply.github.com> Date: Wed, 19 Feb 2025 11:58:14 -0500 Subject: [PATCH 16/52] New networking fixes (#667) * fix: handle `NewListenAddr` events * refactor(networking): stop storing public key separately * refactor(networking): remove duplicated code * refactor(networking): remove `with_` prefix on builder methods * fix: test_peer_info_updates passes * refactor: remove `P2pContext` * refactor: move tests to shared directory --- crates/config/src/lib.rs | 9 +- crates/contexts/src/lib.rs | 2 - crates/contexts/src/p2p.rs | 1 - crates/macros/context-derive/src/lib.rs | 18 -- crates/macros/context-derive/src/p2p.rs | 41 ---- .../macros/context-derive/tests/ui/basic.rs | 43 +--- crates/networking/src/behaviours.rs | 17 +- .../src/blueprint_protocol/behaviour.rs | 16 +- .../src/blueprint_protocol/handler.rs | 2 +- crates/networking/src/discovery/behaviour.rs | 28 +-- crates/networking/src/discovery/config.rs | 23 +-- crates/networking/src/discovery/peers.rs | 10 +- crates/networking/src/error.rs | 6 + crates/networking/src/lib.rs | 8 +- crates/networking/src/networking/tests.rs | 1 - crates/networking/src/service.rs | 65 ++++--- crates/networking/src/service_handle.rs | 8 +- .../src/{tests => }/test_helpers.rs | 22 +-- crates/networking/src/tests/discovery.rs | 184 ++++++++++++++++++ .../networking/src/tests/discovery_tests.rs | 181 ----------------- crates/networking/src/tests/mod.rs | 6 +- 21 files changed, 297 insertions(+), 394 deletions(-) delete mode 100644 crates/macros/context-derive/src/p2p.rs delete mode 100644 crates/networking/src/networking/tests.rs rename crates/networking/src/{tests => }/test_helpers.rs (85%) create mode 100644 crates/networking/src/tests/discovery.rs delete mode 100644 crates/networking/src/tests/discovery_tests.rs diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 89945afd2..10cdc5878 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -78,6 +78,10 @@ pub enum Error { BadRpcConnection(String), #[error("Configuration error: {0}")] ConfigurationError(String), + + #[cfg(feature = "networking")] + #[error("Failed to parse Multiaddr: {0}")] + Multiaddr(#[from] libp2p::multiaddr::Error), } #[cfg(feature = "networking")] @@ -178,11 +182,10 @@ impl GadgetConfiguration { .expect("valid multiaddr; qed"); let network_name: String = network_name.into(); - let network_config = gadget_networking::NetworkConfig { + let network_config = NetworkConfig { instance_id: network_name.clone(), network_name, - instance_secret_key: ecdsa_pair, - instance_public_key: ecdsa_pub_key, + instance_key_pair: ecdsa_pair, local_key: network_identity, listen_addr, target_peer_count: self.target_peer_count, diff --git a/crates/contexts/src/lib.rs b/crates/contexts/src/lib.rs index 6a4c038a7..f3de4e538 100644 --- a/crates/contexts/src/lib.rs +++ b/crates/contexts/src/lib.rs @@ -4,8 +4,6 @@ pub mod eigenlayer; pub mod instrumented_evm_client; #[cfg(feature = "keystore")] pub mod keystore; -#[cfg(feature = "networking")] -pub mod p2p; #[cfg(feature = "tangle")] pub mod services; #[cfg(feature = "tangle")] diff --git a/crates/contexts/src/p2p.rs b/crates/contexts/src/p2p.rs index 8b1378917..e69de29bb 100644 --- a/crates/contexts/src/p2p.rs +++ b/crates/contexts/src/p2p.rs @@ -1 +0,0 @@ - diff --git a/crates/macros/context-derive/src/lib.rs b/crates/macros/context-derive/src/lib.rs index 620e7bdbc..75dd22a8d 100644 --- a/crates/macros/context-derive/src/lib.rs +++ b/crates/macros/context-derive/src/lib.rs @@ -20,9 +20,6 @@ mod eigenlayer; mod evm; /// Keystore context extension implementation. mod keystore; -/// P2P context extension implementation. -#[cfg(feature = "networking")] -mod p2p; /// Tangle context extensions. #[cfg(feature = "tangle")] mod tangle; @@ -116,18 +113,3 @@ pub fn derive_eigenlayer_context(input: TokenStream) -> TokenStream { Err(err) => TokenStream::from(err.to_compile_error()), } } - -/// Derive macro for generating Context Extensions trait implementation for `P2pContext`. -#[proc_macro_derive(P2pContext, attributes(config))] -#[cfg(feature = "networking")] -pub fn derive_p2p_context(input: TokenStream) -> TokenStream { - let input = syn::parse_macro_input!(input as syn::DeriveInput); - let result = - cfg::find_config_field(&input.ident, &input.data, CONFIG_TAG_NAME, CONFIG_TAG_TYPE) - .map(|config_field| p2p::generate_context_impl(input, config_field)); - - match result { - Ok(expanded) => TokenStream::from(expanded), - Err(err) => TokenStream::from(err.to_compile_error()), - } -} diff --git a/crates/macros/context-derive/src/p2p.rs b/crates/macros/context-derive/src/p2p.rs deleted file mode 100644 index c2fb1756f..000000000 --- a/crates/macros/context-derive/src/p2p.rs +++ /dev/null @@ -1,41 +0,0 @@ -use quote::quote; -use syn::DeriveInput; - -use crate::cfg::FieldInfo; - -/// Generate the `MPCContext` implementation for the given struct. -#[allow(clippy::too_many_lines)] -pub fn generate_context_impl( - DeriveInput { - ident: name, - generics, - .. - }: DeriveInput, - config_field: FieldInfo, -) -> proc_macro2::TokenStream { - let field_access = match config_field { - FieldInfo::Named(ident) => quote! { self.#ident }, - FieldInfo::Unnamed(index) => quote! { self.#index }, - }; - - let (impl_generics, ty_generics, where_clause) = generics.split_for_impl(); - - quote! { - #[::blueprint_sdk::macros::ext::async_trait::async_trait] - impl #impl_generics ::blueprint_sdk::macros::ext::contexts::p2p::P2pContext for #name #ty_generics #where_clause { - fn p2p_client( - &self, - name: ::blueprint_sdk::macros::ext::std::string::String, - target_port: u16, - my_ecdsa_key: ::blueprint_sdk::macros::ext::contexts::p2p::GossipMsgKeyPair, - ) -> ::blueprint_sdk::macros::ext::contexts::p2p::P2PClient { - ::blueprint_sdk::macros::ext::contexts::p2p::P2PClient::new( - name, - #field_access.clone(), - target_port, - my_ecdsa_key.clone() - ) - } - } - } -} diff --git a/crates/macros/context-derive/tests/ui/basic.rs b/crates/macros/context-derive/tests/ui/basic.rs index 10575315f..c12138446 100644 --- a/crates/macros/context-derive/tests/ui/basic.rs +++ b/crates/macros/context-derive/tests/ui/basic.rs @@ -2,7 +2,6 @@ use async_trait::async_trait; use blueprint_sdk::config::GadgetConfiguration; use blueprint_sdk::contexts::instrumented_evm_client::EvmInstrumentedClientContext as _; use blueprint_sdk::contexts::keystore::KeystoreContext as _; -use blueprint_sdk::contexts::p2p::P2pContext as _; use blueprint_sdk::contexts::services::ServicesContext as _; use blueprint_sdk::contexts::tangle::TangleClientContext as _; use blueprint_sdk::macros::ext::clients::GadgetServicesClient as _; @@ -14,12 +13,12 @@ use blueprint_sdk::std::collections::BTreeMap; use blueprint_sdk::std::sync::Arc; use blueprint_sdk::stores::local_database::LocalDatabase; use gadget_context_derive::{ - EVMProviderContext, KeystoreContext, P2pContext, ServicesContext, TangleClientContext, + EVMProviderContext, KeystoreContext, ServicesContext, TangleClientContext, }; use round_based::ProtocolMessage as RoundBasedProtocolMessage; use serde::{Deserialize, Serialize}; -#[derive(KeystoreContext, EVMProviderContext, TangleClientContext, ServicesContext, P2pContext)] +#[derive(KeystoreContext, EVMProviderContext, TangleClientContext, ServicesContext)] #[allow(dead_code)] struct MyContext { foo: String, @@ -41,7 +40,7 @@ fn main() { }; // Test existing context functions - let keystore = ctx.keystore(); + let _keystore = ctx.keystore(); let _evm_provider = ctx.evm_client(); let tangle_client = ctx.tangle_client().await.unwrap(); let _services_client = ctx.services_client().await; @@ -50,48 +49,12 @@ fn main() { .current_service_operators([0; 32], 0) .await .unwrap(); - let pub_key = keystore.generate::(None).unwrap(); - let pair = keystore.get_secret::(&pub_key).unwrap(); - let p2p_client = - ctx.p2p_client(String::from("Foo"), 1337, GossipMsgKeyPair(pair.0.clone())); - - // Test MPC context utility functions - let _config = p2p_client.config(); - let _protocol = p2p_client.network_protocol(None); - - // Test MPC context functions - - let mux = Arc::new(NetworkMultiplexer::new(StubNetwork)); - let party_index = 0; - let task_hash = [0u8; 32]; - let mut parties = BTreeMap::::new(); - parties.insert(0, pair.public()); - - // Test network delivery wrapper creation - let _network_wrapper = p2p_client.create_network_delivery_wrapper::( - mux.clone(), - party_index, - task_hash, - parties.clone(), - ); - - // TODO: Test party index retrieval - // let _party_idx = ctx.get_party_index().await; - - // TODO: Test participants retrieval - // let _participants = ctx.get_participants(&tangle_client).await; // Test blueprint ID retrieval let _blueprint_id = tangle_client.blueprint_id(); // Test party index and operators retrieval let _party_idx_ops = tangle_client.get_party_index_and_operators().await; - - // TODO: Test service operators ECDSA keys retrieval - // let _operator_keys = ctx.current_service_operators_ecdsa_keys().await; - - // TODO: Test current call ID retrieval - // let _call_id = tangle_client.current_call_id().await; }; drop(body); diff --git a/crates/networking/src/behaviours.rs b/crates/networking/src/behaviours.rs index 236940eb1..ab3b6a347 100644 --- a/crates/networking/src/behaviours.rs +++ b/crates/networking/src/behaviours.rs @@ -1,4 +1,5 @@ -use crate::key_types::{InstanceMsgKeyPair, InstanceMsgPublicKey}; +use crate::error::Result as NetworkingResult; +use crate::key_types::InstanceMsgKeyPair; use crate::types::ProtocolMessage; use crate::{ blueprint_protocol::{BlueprintProtocolBehaviour, BlueprintProtocolEvent}, @@ -55,8 +56,7 @@ impl GadgetBehaviour { network_name: &str, blueprint_protocol_name: &str, local_key: &Keypair, - instance_secret_key: &InstanceMsgKeyPair, - instance_public_key: &InstanceMsgPublicKey, + instance_key_pair: &InstanceMsgKeyPair, target_peer_count: u64, peer_manager: Arc, protocol_message_sender: Sender, @@ -85,9 +85,9 @@ impl GadgetBehaviour { network_name ); let discovery = DiscoveryConfig::new(local_key.public(), network_name) - .with_mdns(true) - .with_kademlia(true) - .with_target_peer_count(target_peer_count) + .mdns(true) + .kademlia(true) + .target_peer_count(target_peer_count) .build() .unwrap(); @@ -97,8 +97,7 @@ impl GadgetBehaviour { ); let blueprint_protocol = BlueprintProtocolBehaviour::new( local_key, - instance_secret_key, - instance_public_key, + instance_key_pair, peer_manager, blueprint_protocol_name, protocol_message_sender, @@ -114,7 +113,7 @@ impl GadgetBehaviour { } /// Bootstrap Kademlia network - pub fn bootstrap(&mut self) -> Result { + pub fn bootstrap(&mut self) -> NetworkingResult { self.discovery.bootstrap() } diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index 6f681de61..06ffb581b 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -70,10 +70,8 @@ pub struct BlueprintProtocolBehaviour { pub(crate) peer_manager: Arc, /// Libp2p peer ID pub(crate) local_peer_id: PeerId, - /// Instance public key for handshakes and `blueprint_protocol` - pub(crate) instance_public_key: InstanceMsgPublicKey, - /// Instance secret key for handshakes and `blueprint_protocol` - pub(crate) instance_secret_key: InstanceMsgKeyPair, + /// Instance key pair for handshakes and blueprint_protocol + pub(crate) instance_key_pair: InstanceMsgKeyPair, /// Peers with pending inbound handshakes pub(crate) inbound_handshakes: DashMap, /// Peers with pending outbound handshakes @@ -90,8 +88,7 @@ impl BlueprintProtocolBehaviour { #[must_use] pub fn new( local_key: &Keypair, - instance_secret_key: &InstanceMsgKeyPair, - instance_public_key: &InstanceMsgPublicKey, + instance_key_pair: &InstanceMsgKeyPair, peer_manager: Arc, blueprint_protocol_name: &str, protocol_message_sender: Sender, @@ -132,8 +129,7 @@ impl BlueprintProtocolBehaviour { blueprint_protocol_name, peer_manager, local_peer_id, - instance_public_key: *instance_public_key, - instance_secret_key: instance_secret_key.clone(), + instance_key_pair: instance_key_pair.clone(), inbound_handshakes: DashMap::new(), outbound_handshakes: DashMap::new(), response_channels: DashMap::new(), @@ -145,7 +141,7 @@ impl BlueprintProtocolBehaviour { pub(crate) fn sign_handshake(&self, peer: &PeerId) -> InstanceSignedMsgSignature { let msg = peer.to_bytes(); let msg_hash = blake2_256(&msg); - self.instance_secret_key.sign_prehash(&msg_hash) + self.instance_key_pair.sign_prehash(&msg_hash) } /// Send a request to a peer @@ -323,7 +319,7 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { self.send_request( &e.peer_id, InstanceMessageRequest::Handshake { - public_key: self.instance_public_key, + public_key: self.instance_key_pair.public(), signature: self.sign_handshake(&e.peer_id), }, ); diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index a33a74cfe..952c57ab1 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -52,7 +52,7 @@ impl BlueprintProtocolBehaviour { // Send handshake response let response = InstanceMessageResponse::Handshake { - public_key: self.instance_public_key, + public_key: self.instance_key_pair.public().clone(), signature: self.sign_handshake(&peer), }; diff --git a/crates/networking/src/discovery/behaviour.rs b/crates/networking/src/discovery/behaviour.rs index 10c0670f1..515d2b345 100644 --- a/crates/networking/src/discovery/behaviour.rs +++ b/crates/networking/src/discovery/behaviour.rs @@ -5,6 +5,7 @@ use std::{ time::Duration, }; +use crate::error::{Error, Result as NetworkingResult}; use libp2p::{ autonat, core::Multiaddr, @@ -81,11 +82,11 @@ pub struct DiscoveryBehaviour { impl DiscoveryBehaviour { /// Bootstrap Kademlia network - pub fn bootstrap(&mut self) -> Result { + pub fn bootstrap(&mut self) -> NetworkingResult { if let Some(active_kad) = self.discovery.kademlia.as_mut() { - active_kad.bootstrap().map_err(|e| e.to_string()) + active_kad.bootstrap().map_err(Into::into) } else { - Err("Kademlia is not activated".to_string()) + Err(Error::KademliaNotActivated) } } @@ -167,17 +168,6 @@ impl NetworkBehaviour for DiscoveryBehaviour { ) } - fn on_connection_handler_event( - &mut self, - peer_id: PeerId, - connection: ConnectionId, - event: THandlerOutEvent, - ) { - debug!(%peer_id, "Handling connection handler event"); - self.discovery - .on_connection_handler_event(peer_id, connection, event); - } - fn on_swarm_event(&mut self, event: FromSwarm<'_>) { match &event { FromSwarm::ConnectionEstablished(e) => { @@ -204,6 +194,16 @@ impl NetworkBehaviour for DiscoveryBehaviour { self.discovery.on_swarm_event(event); } + fn on_connection_handler_event( + &mut self, + peer_id: PeerId, + connection: ConnectionId, + event: THandlerOutEvent, + ) { + self.discovery + .on_connection_handler_event(peer_id, connection, event); + } + #[allow(clippy::type_complexity)] fn poll( &mut self, diff --git a/crates/networking/src/discovery/config.rs b/crates/networking/src/discovery/config.rs index 727311bb1..3a5ff1c74 100644 --- a/crates/networking/src/discovery/config.rs +++ b/crates/networking/src/discovery/config.rs @@ -59,49 +59,42 @@ impl DiscoveryConfig { /// Set the protocol version that uniquely identifies your P2P service. /// This should be unique to your application to avoid conflicts with other P2P networks. /// Format recommendation: "/" - pub fn with_protocol_version(mut self, version: impl Into) -> Self { + pub fn protocol_version(mut self, version: impl Into) -> Self { self.protocol_version = version.into(); self } - #[must_use] - pub fn with_bootstrap_peers(mut self, peers: Vec<(PeerId, Multiaddr)>) -> Self { + pub fn bootstrap_peers(mut self, peers: Vec<(PeerId, Multiaddr)>) -> Self { self.bootstrap_peers = peers; self } - #[must_use] - pub fn with_relay_nodes(mut self, nodes: Vec<(PeerId, Multiaddr)>) -> Self { + pub fn relay_nodes(mut self, nodes: Vec<(PeerId, Multiaddr)>) -> Self { self.relay_nodes = nodes; self } - #[must_use] - pub fn with_target_peer_count(mut self, count: u64) -> Self { + pub fn target_peer_count(mut self, count: u64) -> Self { self.target_peer_count = count; self } - #[must_use] - pub fn with_mdns(mut self, enable: bool) -> Self { + pub fn mdns(mut self, enable: bool) -> Self { self.enable_mdns = enable; self } - #[must_use] - pub fn with_kademlia(mut self, enable: bool) -> Self { + pub fn kademlia(mut self, enable: bool) -> Self { self.enable_kademlia = enable; self } - #[must_use] - pub fn with_upnp(mut self, enable: bool) -> Self { + pub fn upnp(mut self, enable: bool) -> Self { self.enable_upnp = enable; self } - #[must_use] - pub fn with_relay(mut self, enable: bool) -> Self { + pub fn relay(mut self, enable: bool) -> Self { self.enable_relay = enable; self } diff --git a/crates/networking/src/discovery/peers.rs b/crates/networking/src/discovery/peers.rs index 0791e787c..b07d78b35 100644 --- a/crates/networking/src/discovery/peers.rs +++ b/crates/networking/src/discovery/peers.rs @@ -76,15 +76,7 @@ pub struct PeerManager { impl Default for PeerManager { fn default() -> Self { - let (event_tx, _) = broadcast::channel(100); - Self { - peers: Default::default(), - banned_peers: Default::default(), - verified_peers: Default::default(), - public_keys_to_peer_ids: Default::default(), - whitelisted_keys: Default::default(), - event_tx, - } + Self::new(Default::default()) } } diff --git a/crates/networking/src/error.rs b/crates/networking/src/error.rs index 20990b925..005673143 100644 --- a/crates/networking/src/error.rs +++ b/crates/networking/src/error.rs @@ -36,6 +36,9 @@ pub enum Error { #[error("No network found")] NoNetworkFound, + #[error("Kademlia is not activated")] + KademliaNotActivated, + #[error("Other error: {0}")] Other(String), @@ -44,6 +47,9 @@ pub enum Error { Io(#[from] std::io::Error), // libp2p compat + #[error(transparent)] + NoKnownPeers(#[from] libp2p::kad::NoKnownPeers), + #[error(transparent)] Dial(#[from] libp2p::swarm::DialError), diff --git a/crates/networking/src/lib.rs b/crates/networking/src/lib.rs index 3e921f017..77a27caec 100644 --- a/crates/networking/src/lib.rs +++ b/crates/networking/src/lib.rs @@ -9,12 +9,14 @@ pub mod service; pub mod service_handle; pub mod types; -#[cfg(feature = "round-based-compat")] -pub use gadget_networking_round_based_extension as round_based_compat; - +#[cfg(test)] +pub(crate) mod test_helpers; #[cfg(test)] mod tests; +#[cfg(feature = "round-based-compat")] +pub use gadget_networking_round_based_extension as round_based_compat; + pub use key_types::*; pub use service::{NetworkConfig, NetworkEvent, NetworkService}; diff --git a/crates/networking/src/networking/tests.rs b/crates/networking/src/networking/tests.rs deleted file mode 100644 index 8b1378917..000000000 --- a/crates/networking/src/networking/tests.rs +++ /dev/null @@ -1 +0,0 @@ - diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index e778eebf0..beeee7161 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -69,6 +69,7 @@ pub enum NetworkEvent { /// Network message types #[derive(Debug)] pub enum NetworkMessage { + Dial(Multiaddr), InstanceRequest { peer: PeerId, request: InstanceMessageRequest, @@ -102,9 +103,7 @@ pub struct NetworkConfig { /// Instance id for blueprint protocol pub instance_id: String, /// Instance secret key for blueprint protocol - pub instance_secret_key: InstanceMsgKeyPair, - /// Instance public key for blueprint protocol - pub instance_public_key: InstanceMsgPublicKey, + pub instance_key_pair: InstanceMsgKeyPair, /// Local keypair for authentication pub local_key: Keypair, /// Address to listen on @@ -152,8 +151,7 @@ impl NetworkService { let NetworkConfig { network_name, instance_id, - instance_secret_key, - instance_public_key, + instance_key_pair, local_key, listen_addr, target_peer_count, @@ -174,8 +172,7 @@ impl NetworkService { &network_name, &blueprint_protocol_name, &local_key, - &instance_secret_key, - &instance_public_key, + &instance_key_pair, target_peer_count, peer_manager.clone(), protocol_message_sender.clone(), @@ -268,17 +265,28 @@ impl NetworkService { loop { tokio::select! { swarm_event = self.swarm.select_next_some() => { - if let SwarmEvent::Behaviour(event) = swarm_event { - if let Err(e) = handle_behaviour_event( - &mut self.swarm, - &self.peer_manager, - event, - &self.event_sender, - ) - .await - { - warn!("Failed to handle swarm event: {}", e); - } + match swarm_event { + SwarmEvent::NewListenAddr { address, .. } => { + info!("New listen address: {}", address); + let local_peer_id = *self.swarm.local_peer_id(); + let mut info = self.peer_manager.get_peer_info(&local_peer_id) + .unwrap_or_default(); + info.addresses.insert(address.clone()); + self.peer_manager.update_peer(local_peer_id, info); + }, + SwarmEvent::Behaviour(event) => { + if let Err(e) = handle_behaviour_event( + &mut self.swarm, + &self.peer_manager, + event, + &self.event_sender, + ) + .await + { + warn!("Failed to handle swarm event: {}", e); + } + }, + _ => {} } } Ok(msg) = async { self.network_receiver.try_recv() } => { @@ -446,8 +454,8 @@ async fn handle_discovery_event( /// Handle a blueprint event async fn handle_blueprint_protocol_event( - swarm: &mut Swarm, - peer_manager: &Arc, + _swarm: &mut Swarm, + _peer_manager: &Arc, event: BlueprintProtocolEvent, event_sender: &Sender, ) -> Result<(), Error> { @@ -455,12 +463,12 @@ async fn handle_blueprint_protocol_event( BlueprintProtocolEvent::Request { peer, request, - channel, + channel: _, } => event_sender.send(NetworkEvent::InstanceRequestInbound { peer, request })?, BlueprintProtocolEvent::Response { peer, response, - request_id, + request_id: _, } => event_sender.send(NetworkEvent::InstanceResponseInbound { peer, response })?, BlueprintProtocolEvent::GossipMessage { source, @@ -478,8 +486,8 @@ async fn handle_blueprint_protocol_event( /// Handle a ping event async fn handle_ping_event( - swarm: &mut Swarm, - peer_manager: &Arc, + _swarm: &mut Swarm, + _peer_manager: &Arc, event: ping::Event, event_sender: &Sender, ) -> Result<(), Error> { @@ -509,8 +517,15 @@ async fn handle_ping_event( async fn handle_network_message( swarm: &mut Swarm, msg: NetworkMessage, - peer_manager: &Arc, + _peer_manager: &Arc, event_sender: &Sender, ) -> Result<(), Error> { + match msg { + NetworkMessage::Dial(addr) => { + swarm.dial(addr)?; + } + _ => {} + } + Ok(()) } diff --git a/crates/networking/src/service_handle.rs b/crates/networking/src/service_handle.rs index 02bf9222e..a307b6168 100644 --- a/crates/networking/src/service_handle.rs +++ b/crates/networking/src/service_handle.rs @@ -118,7 +118,7 @@ impl NetworkServiceHandle { let recipient = message.routing.recipient.unwrap(); if let Some(public_key) = recipient.public_key { if let Some(peer_id) = self.peer_manager.get_peer_id_from_public_key(&public_key) { - self.sender.send_message(NetworkMessage::InstanceRequest { + self.send_network_message(NetworkMessage::InstanceRequest { peer: peer_id, request: instance_message_request, })?; @@ -131,13 +131,17 @@ impl NetworkServiceHandle { topic: message.protocol, message: raw_payload, }; - self.sender.send_message(gossip_message)?; + self.send_network_message(gossip_message)?; info!("Sent outbound gossip `NetworkMessage`"); } Ok(()) } + pub(crate) fn send_network_message(&self, message: NetworkMessage) -> Result<(), String> { + self.sender.send_message(message) + } + #[must_use] pub fn get_listen_addr(&self) -> Option { // Get the first peer info for our local peer ID diff --git a/crates/networking/src/tests/test_helpers.rs b/crates/networking/src/test_helpers.rs similarity index 85% rename from crates/networking/src/tests/test_helpers.rs rename to crates/networking/src/test_helpers.rs index a0cadf8f1..008156415 100644 --- a/crates/networking/src/tests/test_helpers.rs +++ b/crates/networking/src/test_helpers.rs @@ -22,25 +22,20 @@ impl TestNode { network_name: &str, instance_id: &str, allowed_keys: HashSet, - bootstrap_peers: Vec<(PeerId, Multiaddr)>, + bootstrap_peers: Vec, ) -> Self { let local_key = identity::Keypair::generate_ed25519(); let peer_id = local_key.public().to_peer_id(); - let listen_addr: Multiaddr = format!("/ip4/127.0.0.1/tcp/0").parse().unwrap(); - info!( - "Creating test node {} with TCP address: {}", - peer_id, listen_addr - ); + let listen_addr: Multiaddr = "/ip4/127.0.0.1/tcp/0".parse().unwrap(); + info!("Creating test node {peer_id} with TCP address: {listen_addr}"); - let instance_secret_key = SpEcdsa::generate_with_seed(None).unwrap(); - let instance_public_key = instance_secret_key.public(); + let instance_key_pair = SpEcdsa::generate_with_seed(None).unwrap(); let config = NetworkConfig { network_name: network_name.to_string(), instance_id: instance_id.to_string(), - instance_secret_key, - instance_public_key, + instance_key_pair, local_key, listen_addr: listen_addr.clone(), target_peer_count: 10, @@ -51,14 +46,13 @@ impl TestNode { let (allowed_keys_tx, allowed_keys_rx) = crossbeam_channel::unbounded(); allowed_keys_tx.send(allowed_keys.clone()).unwrap(); - let service = NetworkService::new(config, allowed_keys, allowed_keys_rx) - .await + let service = NetworkService::new(config, allowed_keys) // TODO: allowed_keys_rx .expect("Failed to create network service"); Self { service: Some(service), peer_id, - listen_addr: Some(listen_addr), + listen_addr: None, // To be set later } } @@ -72,7 +66,7 @@ impl TestNode { let timeout_duration = Duration::from_secs(5); match timeout(timeout_duration, async { while self.listen_addr.is_none() { - if let Some(addr) = handle.get_listen_addr().await { + if let Some(addr) = handle.get_listen_addr() { info!("Node {} listening on {}", self.peer_id, addr); self.listen_addr = Some(addr); break; diff --git a/crates/networking/src/tests/discovery.rs b/crates/networking/src/tests/discovery.rs new file mode 100644 index 000000000..49fbdc4cd --- /dev/null +++ b/crates/networking/src/tests/discovery.rs @@ -0,0 +1,184 @@ +use crate::service::NetworkMessage; +use crate::test_helpers::TestNode; +use std::{collections::HashSet, time::Duration}; +use tokio::time::timeout; +use tracing::{debug, info}; +use tracing_subscriber::{fmt, EnvFilter}; + +fn init_tracing() { + let _ = fmt() + .with_env_filter(EnvFilter::from_default_env()) + .with_target(true) + .with_thread_ids(true) + .with_file(true) + .with_line_number(true) + .try_init(); +} + +#[tokio::test] +async fn test_peer_discovery_mdns() { + init_tracing(); + + let network_name = "test-network"; + let instance_id = "test-instance"; + let allowed_keys = HashSet::new(); + + // Create two nodes + let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; + let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; + + // Start both nodes and wait for them to be listening + let handle1 = node1.start().await.expect("Failed to start node1"); + let handle2 = node2.start().await.expect("Failed to start node2"); + + // Wait for peer discovery (with timeout) + let discovery_timeout = Duration::from_secs(10); + match timeout(discovery_timeout, async { + loop { + let peers1 = handle1.peers(); + let peers2 = handle2.peers(); + + if peers1.contains(&node2.peer_id) && peers2.contains(&node1.peer_id) { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + { + Ok(_) => println!("Peers discovered each other successfully"), + Err(_) => panic!("Peer discovery timed out"), + } +} + +#[tokio::test] +async fn test_peer_discovery_kademlia() { + init_tracing(); + + let network_name = "test-network"; + let instance_id = "test-instance"; + let allowed_keys = HashSet::new(); + + // Create the first node (bootstrap node) + let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; + + // Start node1 and get its listening address + let handle1 = node1.start().await.expect("Failed to start node1"); + let node1_addr = node1.get_listen_addr().expect("Node1 should be listening"); + + // Create two more nodes that will bootstrap from node1 + let bootstrap_peers = vec![node1_addr.clone()]; + let mut node2 = TestNode::new( + network_name, + instance_id, + allowed_keys.clone(), + bootstrap_peers.clone(), + ) + .await; + let mut node3 = TestNode::new(network_name, instance_id, allowed_keys, bootstrap_peers).await; + + // Start the remaining nodes + let handle2 = node2.start().await.expect("Failed to start node2"); + let handle3 = node3.start().await.expect("Failed to start node3"); + + // Wait for peer discovery through Kademlia DHT + let discovery_timeout = Duration::from_secs(20); + match timeout(discovery_timeout, async { + loop { + let peers1 = handle1.peers(); + let peers2 = handle2.peers(); + let peers3 = handle3.peers(); + + if peers1.contains(&node2.peer_id) + && peers1.contains(&node3.peer_id) + && peers2.contains(&node1.peer_id) + && peers2.contains(&node3.peer_id) + && peers3.contains(&node1.peer_id) + && peers3.contains(&node2.peer_id) + { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + { + Ok(_) => println!("All peers discovered each other through Kademlia"), + Err(_) => panic!("Kademlia peer discovery timed out"), + } +} + +#[tokio::test] +async fn test_peer_info_updates() { + init_tracing(); + + let network_name = "test-network"; + let instance_id = "test-instance"; + let allowed_keys = HashSet::new(); + + info!("Creating test nodes..."); + // Create two nodes + let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; + let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; + + info!("Starting nodes..."); + // Start both nodes + let mut handle1 = node1.start().await.expect("Failed to start node1"); + let handle2 = node2.start().await.expect("Failed to start node2"); + + info!("Waiting for peer discovery..."); + // First wait for basic peer discovery (they see each other) + let discovery_timeout = Duration::from_secs(20); + timeout(discovery_timeout, async { + loop { + let peers1 = handle1.peers(); + let peers2 = handle2.peers(); + + debug!("Node1 peers: {:?}, Node2 peers: {:?}", peers1, peers2); + + if peers1.contains(&node2.peer_id) && peers2.contains(&node1.peer_id) { + info!("Basic peer discovery successful"); + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Basic peer discovery timed out"); + + let node2_listen_addr = node2 + .listen_addr + .expect("Node2 did not return a listening address"); + + handle1 + .send_network_message(NetworkMessage::Dial(node2_listen_addr.clone())) + .expect("Node1 failed to dial Node2"); + + info!("Waiting for identify info..."); + // Now wait for identify info to be populated + let identify_timeout = Duration::from_secs(20); + match timeout(identify_timeout, async { + loop { + let peer_info1 = handle1.peer_info(&node2.peer_id); + let peer_info2 = handle2.peer_info(&node1.peer_id); + + if let Some(peer_info) = peer_info1 { + if peer_info.identify_info.is_some() { + // Also verify reverse direction + if let Some(peer_info) = peer_info2 { + if peer_info.identify_info.is_some() { + info!("Identify info populated in both directions"); + break; + } + } + } + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + { + Ok(_) => info!("Peer info updated successfully in both directions"), + Err(_) => panic!("Peer info update timed out"), + } +} diff --git a/crates/networking/src/tests/discovery_tests.rs b/crates/networking/src/tests/discovery_tests.rs deleted file mode 100644 index 50f8dd5c9..000000000 --- a/crates/networking/src/tests/discovery_tests.rs +++ /dev/null @@ -1,181 +0,0 @@ -use super::test_helpers::TestNode; -use std::{collections::HashSet, time::Duration}; -use tokio::time::timeout; -use tracing_subscriber::{fmt, EnvFilter}; - -#[cfg(test)] -mod tests { - use tracing::{debug, info}; - - use super::*; - - fn init_tracing() { - let _ = fmt() - .with_env_filter(EnvFilter::from_default_env()) - .with_target(true) - .with_thread_ids(true) - .with_file(true) - .with_line_number(true) - .try_init(); - } - - #[tokio::test] - async fn test_peer_discovery_mdns() { - let network_name = "test-network"; - let instance_id = "test-instance"; - let allowed_keys = HashSet::new(); - - // Create two nodes - let mut node1 = - TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; - let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; - - // Start both nodes and wait for them to be listening - let handle1 = node1.start().await.expect("Failed to start node1"); - let handle2 = node2.start().await.expect("Failed to start node2"); - - // Wait for peer discovery (with timeout) - let discovery_timeout = Duration::from_secs(10); - match timeout(discovery_timeout, async { - loop { - let peers1 = handle1.peers(); - let peers2 = handle2.peers(); - - if peers1.contains(&node2.peer_id) && peers2.contains(&node1.peer_id) { - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - { - Ok(_) => println!("Peers discovered each other successfully"), - Err(_) => panic!("Peer discovery timed out"), - } - } - - #[tokio::test] - async fn test_peer_discovery_kademlia() { - let network_name = "test-network"; - let instance_id = "test-instance"; - let allowed_keys = HashSet::new(); - - // Create the first node (bootstrap node) - let mut node1 = - TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; - - // Start node1 and get its listening address - let handle1 = node1.start().await.expect("Failed to start node1"); - let node1_addr = node1.get_listen_addr().expect("Node1 should be listening"); - - // Create two more nodes that will bootstrap from node1 - let bootstrap_peers = vec![(node1.peer_id, node1_addr.clone())]; - let mut node2 = TestNode::new( - network_name, - instance_id, - allowed_keys.clone(), - bootstrap_peers.clone(), - ) - .await; - let mut node3 = - TestNode::new(network_name, instance_id, allowed_keys, bootstrap_peers).await; - - // Start the remaining nodes - let handle2 = node2.start().await.expect("Failed to start node2"); - let handle3 = node3.start().await.expect("Failed to start node3"); - - // Wait for peer discovery through Kademlia DHT - let discovery_timeout = Duration::from_secs(20); - match timeout(discovery_timeout, async { - loop { - let peers1 = handle1.peers(); - let peers2 = handle2.peers(); - let peers3 = handle3.peers(); - - if peers1.contains(&node2.peer_id) - && peers1.contains(&node3.peer_id) - && peers2.contains(&node1.peer_id) - && peers2.contains(&node3.peer_id) - && peers3.contains(&node1.peer_id) - && peers3.contains(&node2.peer_id) - { - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - { - Ok(_) => println!("All peers discovered each other through Kademlia"), - Err(_) => panic!("Kademlia peer discovery timed out"), - } - } - - #[tokio::test] - async fn test_peer_info_updates() { - init_tracing(); - - let network_name = "test-network"; - let instance_id = "test-instance"; - let allowed_keys = HashSet::new(); - - info!("Creating test nodes..."); - // Create two nodes - let mut node1 = - TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; - let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; - - info!("Starting nodes..."); - // Start both nodes - let handle1 = node1.start().await.expect("Failed to start node1"); - let handle2 = node2.start().await.expect("Failed to start node2"); - - info!("Waiting for peer discovery..."); - // First wait for basic peer discovery (they see each other) - let discovery_timeout = Duration::from_secs(20); - timeout(discovery_timeout, async { - loop { - let peers1 = handle1.peers(); - let peers2 = handle2.peers(); - - debug!("Node1 peers: {:?}, Node2 peers: {:?}", peers1, peers2); - - if peers1.contains(&node2.peer_id) && peers2.contains(&node1.peer_id) { - info!("Basic peer discovery successful"); - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Basic peer discovery timed out"); - - info!("Waiting for identify info..."); - // Now wait for identify info to be populated - let identify_timeout = Duration::from_secs(20); - match timeout(identify_timeout, async { - loop { - let peer_info1 = handle1.peer_info(&node2.peer_id); - let peer_info2 = handle2.peer_info(&node1.peer_id); - - if let Some(peer_info) = peer_info1 { - if peer_info.identify_info.is_some() { - // Also verify reverse direction - if let Some(peer_info) = peer_info2 { - if peer_info.identify_info.is_some() { - info!("Identify info populated in both directions"); - break; - } - } - } - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - { - Ok(_) => info!("Peer info updated successfully in both directions"), - Err(_) => panic!("Peer info update timed out"), - } - } -} diff --git a/crates/networking/src/tests/mod.rs b/crates/networking/src/tests/mod.rs index 993863136..58184b5c9 100644 --- a/crates/networking/src/tests/mod.rs +++ b/crates/networking/src/tests/mod.rs @@ -1,5 +1 @@ -mod discovery_tests; -mod test_helpers; - -// Re-export test helpers for use in other test modules -pub(crate) use test_helpers::*; +mod discovery; \ No newline at end of file From 9e61009f9318010c535d3b5192dc8b44c0c9ed43 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 10:11:24 -0700 Subject: [PATCH 17/52] chore: fix hashing --- crates/crypto/sp-core/src/lib.rs | 2 +- crates/networking/README.md | 77 ++++++++++++------- .../src/blueprint_protocol/behaviour.rs | 29 ++++--- .../src/blueprint_protocol/handler.rs | 5 +- crates/networking/src/lib.rs | 22 ------ 5 files changed, 73 insertions(+), 62 deletions(-) diff --git a/crates/crypto/sp-core/src/lib.rs b/crates/crypto/sp-core/src/lib.rs index b78610066..6fb82e4c4 100644 --- a/crates/crypto/sp-core/src/lib.rs +++ b/crates/crypto/sp-core/src/lib.rs @@ -126,7 +126,7 @@ macro_rules! impl_sp_core_pair_public { macro_rules! impl_sp_core_signature { ($key_type:ident, $pair_type:ty) => { paste::paste! { - #[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] + #[derive(Default, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub struct [](pub <$pair_type as sp_core::Pair>::Signature); impl PartialOrd for [] { diff --git a/crates/networking/README.md b/crates/networking/README.md index 71c8d1fb1..99b4ea8ae 100644 --- a/crates/networking/README.md +++ b/crates/networking/README.md @@ -11,73 +11,84 @@ sequenceDiagram participant A as Peer A participant B as Peer B - Note over A,B: Connection Established + Note over A,B: Initial TCP/QUIC Connection Established - A->>B: HandshakeRequest { + Note over A: Generate signature of B's peer ID + A->>+B: HandshakeRequest { public_key: A_pub, signature: sign(B_peer_id) } - Note over B: Verify A's signature - Note over B: Store A's public key + Note over B: 1. Verify A's signature
2. Store A's public key - B->>A: HandshakeResponse { + Note over B: Generate signature of A's peer ID + B-->>-A: HandshakeResponse { public_key: B_pub, signature: sign(A_peer_id) } - Note over A: Verify B's signature - Note over A: Store B's public key + Note over A: 1. Verify B's signature
2. Store B's public key - Note over A,B: Both peers verified - Note over A,B: Protocol messages allowed + Note over A,B: ✓ Handshake Complete + Note over A,B: ✓ Protocol Messages Allowed ``` ### Handshake States ```mermaid stateDiagram-v2 - [*] --> Connected: Connection Established + direction LR + [*] --> Connected: New Connection + Connected --> OutboundPending: Send Handshake Connected --> InboundPending: Receive Handshake OutboundPending --> Verified: Valid Response OutboundPending --> Failed: Invalid/Timeout - InboundPending --> Verified: Valid Request + InboundPending --> Verified: Valid Request & Response InboundPending --> Failed: Invalid/Timeout Verified --> [*]: Connection Closed Failed --> [*]: Connection Closed + + note right of Connected + Initial TCP/QUIC connection established + end note + + note right of Verified + Both peers authenticated + Protocol messages allowed + end note ``` -## Blueprint Protocol Instance Request/Response +## Protocol Message Exchange -After handshake completion, peers can exchange protocol messages. +After handshake completion, peers can exchange protocol messages through direct P2P or broadcast channels. ```mermaid sequenceDiagram participant A as Peer A (Verified) participant B as Peer B (Verified) - Note over A,B: Handshake Completed + Note over A,B: ✓ Handshake Completed - A->>B: InstanceMessageRequest::Protocol { + A->>+B: InstanceMessageRequest { protocol: String, payload: Vec, metadata: Option> } - alt Success Response - B->>A: InstanceMessageResponse::Success { + alt Success Case + B-->>-A: InstanceMessageResponse::Success { data: Option> } else Protocol Response - B->>A: InstanceMessageResponse::Protocol { + B-->>-A: InstanceMessageResponse::Protocol { data: Vec } - else Error Response - B->>A: InstanceMessageResponse::Error { + else Error Case + B-->>-A: InstanceMessageResponse::Error { code: u16, message: String } @@ -88,6 +99,7 @@ sequenceDiagram ```mermaid stateDiagram-v2 + direction LR [*] --> Handshaked: Peers Verified Handshaked --> RequestPending: Send Request @@ -100,6 +112,16 @@ stateDiagram-v2 ErrorSent --> Handshaked: Complete Handshaked --> [*]: Connection Closed + + note right of Handshaked + Peers authenticated + Ready for messages + end note + + note right of Processing + Validating request + Processing payload + end note ``` ## Protocol Details @@ -112,13 +134,16 @@ stateDiagram-v2 - Timeouts after 30 seconds - Handles concurrent handshakes gracefully -### Blueprint Protocol +### Protocol Message Types -- Requires completed handshake -- Supports protocol-specific messages -- Includes metadata for routing/handling -- Error responses for protocol violations -- Supports both direct and broadcast messaging +- Direct P2P messages: + - Targeted to specific peer + - Requires peer verification + - Guaranteed delivery attempt +- Broadcast messages: + - Sent to all peers + - Uses gossipsub protocol + - Best-effort delivery ### Security Features diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index 06ffb581b..cf7b6d0b0 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -1,10 +1,11 @@ +use super::{InstanceMessageRequest, InstanceMessageResponse}; +use crate::discovery::PeerManager; use crate::{ types::ProtocolMessage, Curve, InstanceMsgKeyPair, InstanceMsgPublicKey, - InstanceSignedMsgSignature, KeySignExt, + InstanceSignedMsgSignature, }; use crossbeam_channel::Sender; use dashmap::DashMap; -use gadget_crypto::tangle_pair_signer::sp_core::blake2_256; use gadget_crypto::KeyType; use libp2p::{ core::transport::PortUse, @@ -24,10 +25,6 @@ use std::{ }; use tracing::{debug, info, trace, warn}; -use crate::discovery::PeerManager; - -use super::{InstanceMessageRequest, InstanceMessageResponse}; - #[derive(NetworkBehaviour)] pub struct DerivedBlueprintProtocolBehaviour { /// Request/response protocol for p2p messaging @@ -138,10 +135,19 @@ impl BlueprintProtocolBehaviour { } /// Sign a handshake message for a peer - pub(crate) fn sign_handshake(&self, peer: &PeerId) -> InstanceSignedMsgSignature { + pub(crate) fn sign_handshake( + &self, + key_pair: &mut InstanceMsgKeyPair, + peer: &PeerId, + ) -> InstanceSignedMsgSignature { let msg = peer.to_bytes(); - let msg_hash = blake2_256(&msg); - self.instance_key_pair.sign_prehash(&msg_hash) + match ::sign_with_secret(key_pair, &msg) { + Ok(signature) => signature, + Err(e) => { + warn!("Failed to sign handshake message: {e}"); + InstanceSignedMsgSignature::default() + } + } } /// Send a request to a peer @@ -316,11 +322,12 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { "Established connection with unverified peer {:?}, sending handshake", e.peer_id ); + let mut key_pair = self.instance_key_pair.clone(); self.send_request( &e.peer_id, InstanceMessageRequest::Handshake { - public_key: self.instance_key_pair.public(), - signature: self.sign_handshake(&e.peer_id), + public_key: key_pair.public(), + signature: self.sign_handshake(&mut key_pair, &e.peer_id), }, ); self.outbound_handshakes.insert(e.peer_id, Instant::now()); diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index 952c57ab1..4baeaf0d0 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -51,9 +51,10 @@ impl BlueprintProtocolBehaviour { .add_peer_id_to_public_key(&peer, &public_key); // Send handshake response + let mut key_pair = self.instance_key_pair.clone(); let response = InstanceMessageResponse::Handshake { - public_key: self.instance_key_pair.public().clone(), - signature: self.sign_handshake(&peer), + public_key: key_pair.public().clone(), + signature: self.sign_handshake(&mut key_pair, &peer), }; if let Err(e) = self.send_response(channel, response) { diff --git a/crates/networking/src/lib.rs b/crates/networking/src/lib.rs index 77a27caec..7aa838009 100644 --- a/crates/networking/src/lib.rs +++ b/crates/networking/src/lib.rs @@ -30,12 +30,6 @@ pub mod key_types { SpEcdsa as Curve, SpEcdsaPair as InstanceMsgKeyPair, SpEcdsaPublic as InstanceMsgPublicKey, SpEcdsaSignature as InstanceSignedMsgSignature, }; - - impl super::KeySignExt for InstanceMsgKeyPair { - fn sign_prehash(&self, prehash: &[u8; 32]) -> InstanceSignedMsgSignature { - InstanceSignedMsgSignature(self.0.sign_prehashed(prehash)) - } - } } #[cfg(all( @@ -48,12 +42,6 @@ pub mod key_types { SpSr25519 as Curve, SpSr25519Pair as InstanceMsgKeyPair, SpSr25519Public as InstanceMsgPublicKey, SpSr25519Signature as InstanceSignedMsgSignature, }; - - impl super::KeySignExt for InstanceMsgKeyPair { - fn sign_prehash(&self, prehash: &[u8; 32]) -> InstanceSignedMsgSignature { - InstanceSignedMsgSignature(self.0.sign_prehashed(prehash)) - } - } } #[cfg(all( @@ -66,12 +54,6 @@ pub mod key_types { SpEd25519 as Curve, SpEd25519Pair as InstanceMsgKeyPair, SpEd25519Public as InstanceMsgPublicKey, SpEd25519Signature as InstanceSignedMsgSignature, }; - - impl super::KeySignExt for InstanceMsgKeyPair { - fn sign_prehash(&self, prehash: &[u8; 32]) -> InstanceSignedMsgSignature { - InstanceSignedMsgSignature(self.0.sign_prehashed(prehash)) - } - } } #[cfg(all( @@ -102,7 +84,3 @@ pub mod key_types { compile_error!( "Only one of 'sp-core-ecdsa', 'sp-core-sr25519', or 'sp-core-ed25519' features can be enabled at a time" ); - -pub(crate) trait KeySignExt { - fn sign_prehash(&self, prehash: &[u8; 32]) -> InstanceSignedMsgSignature; -} From 80c4bdc11f251e2eed08a5e613162ceae223de9e Mon Sep 17 00:00:00 2001 From: Serial <69764315+Serial-ATA@users.noreply.github.com> Date: Wed, 19 Feb 2025 12:16:38 -0500 Subject: [PATCH 18/52] refactor(networking): move testing utils --- crates/networking/src/lib.rs | 2 - crates/networking/src/test_helpers.rs | 144 ------------------- crates/networking/src/tests/discovery.rs | 22 +-- crates/networking/src/tests/handshake.rs | 12 ++ crates/networking/src/tests/mod.rs | 174 ++++++++++++++++++++++- 5 files changed, 189 insertions(+), 165 deletions(-) delete mode 100644 crates/networking/src/test_helpers.rs create mode 100644 crates/networking/src/tests/handshake.rs diff --git a/crates/networking/src/lib.rs b/crates/networking/src/lib.rs index 7aa838009..52f445344 100644 --- a/crates/networking/src/lib.rs +++ b/crates/networking/src/lib.rs @@ -9,8 +9,6 @@ pub mod service; pub mod service_handle; pub mod types; -#[cfg(test)] -pub(crate) mod test_helpers; #[cfg(test)] mod tests; diff --git a/crates/networking/src/test_helpers.rs b/crates/networking/src/test_helpers.rs deleted file mode 100644 index 008156415..000000000 --- a/crates/networking/src/test_helpers.rs +++ /dev/null @@ -1,144 +0,0 @@ -use std::{collections::HashSet, mem, time::Duration}; - -use gadget_crypto::{sp_core::SpEcdsa, KeyType}; -use libp2p::{identity, Multiaddr, PeerId}; -use tokio::time::timeout; -use tracing::info; - -use crate::{ - key_types::InstanceMsgPublicKey, service_handle::NetworkServiceHandle, NetworkConfig, - NetworkService, -}; - -/// Test node configuration for network tests -pub struct TestNode { - pub service: Option, - pub peer_id: PeerId, - pub listen_addr: Option, -} - -impl TestNode { - pub async fn new( - network_name: &str, - instance_id: &str, - allowed_keys: HashSet, - bootstrap_peers: Vec, - ) -> Self { - let local_key = identity::Keypair::generate_ed25519(); - let peer_id = local_key.public().to_peer_id(); - - let listen_addr: Multiaddr = "/ip4/127.0.0.1/tcp/0".parse().unwrap(); - info!("Creating test node {peer_id} with TCP address: {listen_addr}"); - - let instance_key_pair = SpEcdsa::generate_with_seed(None).unwrap(); - - let config = NetworkConfig { - network_name: network_name.to_string(), - instance_id: instance_id.to_string(), - instance_key_pair, - local_key, - listen_addr: listen_addr.clone(), - target_peer_count: 10, - bootstrap_peers, - enable_mdns: true, - enable_kademlia: true, - }; - - let (allowed_keys_tx, allowed_keys_rx) = crossbeam_channel::unbounded(); - allowed_keys_tx.send(allowed_keys.clone()).unwrap(); - let service = NetworkService::new(config, allowed_keys) // TODO: allowed_keys_rx - .expect("Failed to create network service"); - - Self { - service: Some(service), - peer_id, - listen_addr: None, // To be set later - } - } - - /// Start the node - pub async fn start(&mut self) -> Result { - // Take ownership of the service - let service = self.service.take().ok_or("Service already started")?; - let handle = service.start(); - - // Wait for the actual listening address - let timeout_duration = Duration::from_secs(5); - match timeout(timeout_duration, async { - while self.listen_addr.is_none() { - if let Some(addr) = handle.get_listen_addr() { - info!("Node {} listening on {}", self.peer_id, addr); - self.listen_addr = Some(addr); - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - { - Ok(_) => Ok(handle), - Err(_) => Err("Timeout waiting for node to start listening"), - } - } - - /// Get the actual listening address - pub fn get_listen_addr(&self) -> Option { - self.listen_addr.clone() - } -} - -/// Wait for a condition with timeout -pub async fn wait_for_condition(timeout: Duration, mut condition: F) -> Result<(), &'static str> -where - F: FnMut() -> bool, -{ - let start = std::time::Instant::now(); - while !condition() { - if start.elapsed() > timeout { - return Err("Timeout waiting for condition"); - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - Ok(()) -} - -/// Wait for peers to discover each other -pub async fn wait_for_peer_discovery( - nodes: &[&TestNode], - timeout: Duration, -) -> Result<(), &'static str> { - wait_for_condition(timeout, || { - for (i, node1) in nodes.iter().enumerate() { - for (j, node2) in nodes.iter().enumerate() { - if i != j - && !node1 - .service - .as_ref() - .map(|s| s.peer_manager.get_peers()) - .unwrap_or_default() - .contains_key(&node2.peer_id) - { - return false; - } - } - } - true - }) - .await -} - -/// Wait for peer info to be updated -pub async fn wait_for_peer_info( - node: &TestNode, - peer_id: &PeerId, - timeout: Duration, -) -> Result<(), &'static str> { - wait_for_condition(timeout, || { - node.service - .as_ref() - .map(|s| s.peer_manager.get_peer_info(peer_id)) - .unwrap_or(None) - .map_or(false, |info| info.identify_info.is_some()) - }) - .await -} diff --git a/crates/networking/src/tests/discovery.rs b/crates/networking/src/tests/discovery.rs index 49fbdc4cd..ecbeb83f4 100644 --- a/crates/networking/src/tests/discovery.rs +++ b/crates/networking/src/tests/discovery.rs @@ -1,19 +1,11 @@ use crate::service::NetworkMessage; -use crate::test_helpers::TestNode; +use super::TestNode; use std::{collections::HashSet, time::Duration}; use tokio::time::timeout; use tracing::{debug, info}; use tracing_subscriber::{fmt, EnvFilter}; - -fn init_tracing() { - let _ = fmt() - .with_env_filter(EnvFilter::from_default_env()) - .with_target(true) - .with_thread_ids(true) - .with_file(true) - .with_line_number(true) - .try_init(); -} +use super::init_tracing; +use super::NetworkServiceHandleExt; #[tokio::test] async fn test_peer_discovery_mdns() { @@ -146,13 +138,7 @@ async fn test_peer_info_updates() { .await .expect("Basic peer discovery timed out"); - let node2_listen_addr = node2 - .listen_addr - .expect("Node2 did not return a listening address"); - - handle1 - .send_network_message(NetworkMessage::Dial(node2_listen_addr.clone())) - .expect("Node1 failed to dial Node2"); + handle1.dial(&handle2); info!("Waiting for identify info..."); // Now wait for identify info to be populated diff --git a/crates/networking/src/tests/handshake.rs b/crates/networking/src/tests/handshake.rs new file mode 100644 index 000000000..62910d1a3 --- /dev/null +++ b/crates/networking/src/tests/handshake.rs @@ -0,0 +1,12 @@ +use tokio::time::timeout; +use std::collections::HashSet; +use std::time::Duration; +use super::init_tracing; +use super::TestNode; + +#[tokio::test] +async fn test_peer_handshake() { + init_tracing(); + + todo!() +} \ No newline at end of file diff --git a/crates/networking/src/tests/mod.rs b/crates/networking/src/tests/mod.rs index 58184b5c9..6cf4476b1 100644 --- a/crates/networking/src/tests/mod.rs +++ b/crates/networking/src/tests/mod.rs @@ -1 +1,173 @@ -mod discovery; \ No newline at end of file +use crate::service::NetworkMessage; +use std::{collections::HashSet, mem, time::Duration}; + +use gadget_crypto::{sp_core::SpEcdsa, KeyType}; +use libp2p::{identity, Multiaddr, PeerId}; +use tokio::time::timeout; +use tracing::info; + +use crate::{ + key_types::InstanceMsgPublicKey, service_handle::NetworkServiceHandle, NetworkConfig, + NetworkService, +}; + +mod discovery; +mod handshake; + +fn init_tracing() { + let _ = tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_target(true) + .with_thread_ids(true) + .with_file(true) + .with_line_number(true) + .try_init(); +} + +/// Test node configuration for network tests +pub struct TestNode { + pub service: Option, + pub peer_id: PeerId, + pub listen_addr: Option, +} + +impl TestNode { + pub async fn new( + network_name: &str, + instance_id: &str, + allowed_keys: HashSet, + bootstrap_peers: Vec, + ) -> Self { + let local_key = identity::Keypair::generate_ed25519(); + let peer_id = local_key.public().to_peer_id(); + + let listen_addr: Multiaddr = "/ip4/127.0.0.1/tcp/0".parse().unwrap(); + info!("Creating test node {peer_id} with TCP address: {listen_addr}"); + + let instance_key_pair = SpEcdsa::generate_with_seed(None).unwrap(); + + let config = NetworkConfig { + network_name: network_name.to_string(), + instance_id: instance_id.to_string(), + instance_key_pair, + local_key, + listen_addr: listen_addr.clone(), + target_peer_count: 10, + bootstrap_peers, + enable_mdns: true, + enable_kademlia: true, + }; + + let (allowed_keys_tx, allowed_keys_rx) = crossbeam_channel::unbounded(); + allowed_keys_tx.send(allowed_keys.clone()).unwrap(); + let service = NetworkService::new(config, allowed_keys) // TODO: allowed_keys_rx + .expect("Failed to create network service"); + + Self { + service: Some(service), + peer_id, + listen_addr: None, // To be set later + } + } + + /// Start the node + pub async fn start(&mut self) -> Result { + // Take ownership of the service + let service = self.service.take().ok_or("Service already started")?; + let handle = service.start(); + + // Wait for the actual listening address + let timeout_duration = Duration::from_secs(5); + match timeout(timeout_duration, async { + while self.listen_addr.is_none() { + if let Some(addr) = handle.get_listen_addr() { + info!("Node {} listening on {}", self.peer_id, addr); + self.listen_addr = Some(addr); + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + { + Ok(_) => Ok(handle), + Err(_) => Err("Timeout waiting for node to start listening"), + } + } + + /// Get the actual listening address + pub fn get_listen_addr(&self) -> Option { + self.listen_addr.clone() + } +} + +/// Wait for a condition with timeout +pub async fn wait_for_condition(timeout: Duration, mut condition: F) -> Result<(), &'static str> +where + F: FnMut() -> bool, +{ + let start = std::time::Instant::now(); + while !condition() { + if start.elapsed() > timeout { + return Err("Timeout waiting for condition"); + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + Ok(()) +} + +/// Wait for peers to discover each other +pub async fn wait_for_peer_discovery( + nodes: &[&TestNode], + timeout: Duration, +) -> Result<(), &'static str> { + wait_for_condition(timeout, || { + for (i, node1) in nodes.iter().enumerate() { + for (j, node2) in nodes.iter().enumerate() { + if i != j + && !node1 + .service + .as_ref() + .map(|s| s.peer_manager.get_peers()) + .unwrap_or_default() + .contains_key(&node2.peer_id) + { + return false; + } + } + } + true + }) + .await +} + +/// Wait for peer info to be updated +pub async fn wait_for_peer_info( + node: &TestNode, + peer_id: &PeerId, + timeout: Duration, +) -> Result<(), &'static str> { + wait_for_condition(timeout, || { + node.service + .as_ref() + .map(|s| s.peer_manager.get_peer_info(peer_id)) + .unwrap_or(None) + .map_or(false, |info| info.identify_info.is_some()) + }) + .await +} + +trait NetworkServiceHandleExt { + fn dial(&self, other: &NetworkServiceHandle); +} + +impl NetworkServiceHandleExt for NetworkServiceHandle { + fn dial(&self, other: &NetworkServiceHandle) { + let listen_addr = other.get_listen_addr() + .expect("Node2 did not return a listening address"); + + self + .send_network_message(NetworkMessage::Dial(listen_addr.clone())) + .expect("Node1 failed to dial Node2"); + } +} \ No newline at end of file From 1150fca7674d13eec261f6c94b5135903ebc502e Mon Sep 17 00:00:00 2001 From: Serial <69764315+Serial-ATA@users.noreply.github.com> Date: Wed, 19 Feb 2025 12:48:28 -0500 Subject: [PATCH 19/52] feat: add handshake test --- crates/networking/src/tests/discovery.rs | 76 ++++-------------------- crates/networking/src/tests/handshake.rs | 38 ++++++++++-- crates/networking/src/tests/mod.rs | 68 +++++++++++++-------- 3 files changed, 89 insertions(+), 93 deletions(-) diff --git a/crates/networking/src/tests/discovery.rs b/crates/networking/src/tests/discovery.rs index ecbeb83f4..7687010f5 100644 --- a/crates/networking/src/tests/discovery.rs +++ b/crates/networking/src/tests/discovery.rs @@ -1,11 +1,12 @@ -use crate::service::NetworkMessage; +use super::init_tracing; +use super::NetworkServiceHandleExt; use super::TestNode; +use super::{wait_for_peer_discovery, wait_for_peer_info}; +use crate::service::NetworkMessage; use std::{collections::HashSet, time::Duration}; use tokio::time::timeout; use tracing::{debug, info}; use tracing_subscriber::{fmt, EnvFilter}; -use super::init_tracing; -use super::NetworkServiceHandleExt; #[tokio::test] async fn test_peer_discovery_mdns() { @@ -23,24 +24,11 @@ async fn test_peer_discovery_mdns() { let handle1 = node1.start().await.expect("Failed to start node1"); let handle2 = node2.start().await.expect("Failed to start node2"); - // Wait for peer discovery (with timeout) - let discovery_timeout = Duration::from_secs(10); - match timeout(discovery_timeout, async { - loop { - let peers1 = handle1.peers(); - let peers2 = handle2.peers(); - - if peers1.contains(&node2.peer_id) && peers2.contains(&node1.peer_id) { - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - { - Ok(_) => println!("Peers discovered each other successfully"), - Err(_) => panic!("Peer discovery timed out"), - } + // First wait for basic peer discovery (they see each other) + let discovery_timeout = Duration::from_secs(20); + wait_for_peer_discovery(&[&handle1, &handle2], discovery_timeout) + .await + .expect("Basic peer discovery timed out"); } #[tokio::test] @@ -118,53 +106,15 @@ async fn test_peer_info_updates() { let mut handle1 = node1.start().await.expect("Failed to start node1"); let handle2 = node2.start().await.expect("Failed to start node2"); - info!("Waiting for peer discovery..."); // First wait for basic peer discovery (they see each other) let discovery_timeout = Duration::from_secs(20); - timeout(discovery_timeout, async { - loop { - let peers1 = handle1.peers(); - let peers2 = handle2.peers(); - - debug!("Node1 peers: {:?}, Node2 peers: {:?}", peers1, peers2); - - if peers1.contains(&node2.peer_id) && peers2.contains(&node1.peer_id) { - info!("Basic peer discovery successful"); - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Basic peer discovery timed out"); + wait_for_peer_discovery(&[&handle1, &handle2], discovery_timeout) + .await + .expect("Basic peer discovery timed out"); handle1.dial(&handle2); - info!("Waiting for identify info..."); // Now wait for identify info to be populated let identify_timeout = Duration::from_secs(20); - match timeout(identify_timeout, async { - loop { - let peer_info1 = handle1.peer_info(&node2.peer_id); - let peer_info2 = handle2.peer_info(&node1.peer_id); - - if let Some(peer_info) = peer_info1 { - if peer_info.identify_info.is_some() { - // Also verify reverse direction - if let Some(peer_info) = peer_info2 { - if peer_info.identify_info.is_some() { - info!("Identify info populated in both directions"); - break; - } - } - } - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - { - Ok(_) => info!("Peer info updated successfully in both directions"), - Err(_) => panic!("Peer info update timed out"), - } + wait_for_peer_info(&handle1, &handle2, identify_timeout).await; } diff --git a/crates/networking/src/tests/handshake.rs b/crates/networking/src/tests/handshake.rs index 62910d1a3..8077247ea 100644 --- a/crates/networking/src/tests/handshake.rs +++ b/crates/networking/src/tests/handshake.rs @@ -1,12 +1,40 @@ -use tokio::time::timeout; +use super::TestNode; +use super::{init_tracing, wait_for_peer_discovery}; +use super::{wait_for_peer_info, NetworkServiceHandleExt}; use std::collections::HashSet; use std::time::Duration; -use super::init_tracing; -use super::TestNode; +use tokio::time::timeout; #[tokio::test] async fn test_peer_handshake() { init_tracing(); - todo!() -} \ No newline at end of file + let network_name = "test-network"; + let instance_id = "test-instance"; + let allowed_keys = HashSet::new(); + + // Create two nodes + let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; + let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; + + // Start both nodes and wait for them to be listening + let handle1 = node1.start().await.expect("Failed to start node1"); + let handle2 = node2.start().await.expect("Failed to start node2"); + + // First wait for basic peer discovery (they see each other) + let discovery_timeout = Duration::from_secs(20); + wait_for_peer_discovery(&[&handle1, &handle2], discovery_timeout) + .await + .expect("Basic peer discovery timed out"); + + handle1.dial(&handle2); + + let identify_timeout = Duration::from_secs(20); + wait_for_peer_info(&handle1, &handle2, identify_timeout).await; + + let node2_peer = handle2.local_peer_id; + assert!( + handle1.peer_manager.is_peer_verified(&node2_peer), + "Node2 was not verified, handshake failed" + ); +} diff --git a/crates/networking/src/tests/mod.rs b/crates/networking/src/tests/mod.rs index 6cf4476b1..9b864cf32 100644 --- a/crates/networking/src/tests/mod.rs +++ b/crates/networking/src/tests/mod.rs @@ -88,7 +88,7 @@ impl TestNode { tokio::time::sleep(Duration::from_millis(100)).await; } }) - .await + .await { Ok(_) => Ok(handle), Err(_) => Err("Timeout waiting for node to start listening"), @@ -118,19 +118,19 @@ where /// Wait for peers to discover each other pub async fn wait_for_peer_discovery( - nodes: &[&TestNode], + handles: &[&NetworkServiceHandle], timeout: Duration, ) -> Result<(), &'static str> { + info!("Waiting for peer discovery..."); + wait_for_condition(timeout, || { - for (i, node1) in nodes.iter().enumerate() { - for (j, node2) in nodes.iter().enumerate() { + for (i, handle1) in handles.iter().enumerate() { + for (j, handle2) in handles.iter().enumerate() { if i != j - && !node1 - .service - .as_ref() - .map(|s| s.peer_manager.get_peers()) - .unwrap_or_default() - .contains_key(&node2.peer_id) + && !handle1 + .peers() + .iter() + .any(|id| *id == handle2.local_peer_id) { return false; } @@ -138,23 +138,41 @@ pub async fn wait_for_peer_discovery( } true }) - .await + .await } /// Wait for peer info to be updated pub async fn wait_for_peer_info( - node: &TestNode, - peer_id: &PeerId, + handle1: &NetworkServiceHandle, + handle2: &NetworkServiceHandle, timeout: Duration, -) -> Result<(), &'static str> { - wait_for_condition(timeout, || { - node.service - .as_ref() - .map(|s| s.peer_manager.get_peer_info(peer_id)) - .unwrap_or(None) - .map_or(false, |info| info.identify_info.is_some()) +) { + info!("Waiting for identify info..."); + + match tokio::time::timeout(timeout, async { + loop { + let peer_info1 = handle1.peer_info(&handle2.local_peer_id); + let peer_info2 = handle2.peer_info(&handle1.local_peer_id); + + if let Some(peer_info) = peer_info1 { + if peer_info.identify_info.is_some() { + // Also verify reverse direction + if let Some(peer_info) = peer_info2 { + if peer_info.identify_info.is_some() { + info!("Identify info populated in both directions"); + break; + } + } + } + } + tokio::time::sleep(Duration::from_millis(100)).await; + } }) - .await + .await + { + Ok(_) => info!("Peer info updated successfully in both directions"), + Err(_) => panic!("Peer info update timed out"), + } } trait NetworkServiceHandleExt { @@ -163,11 +181,11 @@ trait NetworkServiceHandleExt { impl NetworkServiceHandleExt for NetworkServiceHandle { fn dial(&self, other: &NetworkServiceHandle) { - let listen_addr = other.get_listen_addr() + let listen_addr = other + .get_listen_addr() .expect("Node2 did not return a listening address"); - self - .send_network_message(NetworkMessage::Dial(listen_addr.clone())) + self.send_network_message(NetworkMessage::Dial(listen_addr.clone())) .expect("Node1 failed to dial Node2"); } -} \ No newline at end of file +} From 8e0d458250289d5f076fda2f367a84fe41391291 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 10:50:32 -0700 Subject: [PATCH 20/52] chore: add handshake tests / dial on connection --- crates/networking/src/service.rs | 45 +++++++++++++++--------------- crates/networking/src/tests/mod.rs | 2 +- 2 files changed, 24 insertions(+), 23 deletions(-) diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index beeee7161..ad9b00bc8 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -19,8 +19,11 @@ use crate::{ use crossbeam_channel::{self, Receiver, Sender}; use futures::StreamExt; use libp2p::{ - identify, identity::Keypair, kad, mdns, ping, swarm::SwarmEvent, Multiaddr, PeerId, Swarm, - SwarmBuilder, + identify, + identity::Keypair, + kad, mdns, ping, + swarm::{dial_opts::DialOpts, SwarmEvent}, + Multiaddr, PeerId, Swarm, SwarmBuilder, }; use tracing::trace; use tracing::{debug, info, warn}; @@ -338,14 +341,7 @@ async fn handle_behaviour_event( match event { GadgetBehaviourEvent::ConnectionLimits(_) => {} GadgetBehaviourEvent::Discovery(discovery_event) => { - handle_discovery_event( - &swarm.behaviour().discovery.peer_info, - peer_manager, - discovery_event, - event_sender, - &swarm.behaviour().blueprint_protocol.blueprint_protocol_name, - ) - .await?; + handle_discovery_event(swarm, peer_manager, discovery_event, event_sender).await?; } GadgetBehaviourEvent::BlueprintProtocol(blueprint_event) => { handle_blueprint_protocol_event(swarm, peer_manager, blueprint_event, event_sender) @@ -361,17 +357,16 @@ async fn handle_behaviour_event( /// Handle a discovery event async fn handle_discovery_event( - peer_info_map: &HashMap, + swarm: &mut Swarm, peer_manager: &Arc, event: DiscoveryEvent, event_sender: &Sender, - blueprint_protocol_name: &str, ) -> Result<(), Error> { match event { DiscoveryEvent::PeerConnected(peer_id) => { info!("Peer connected, {peer_id}"); // Update peer info when connected - if let Some(info) = peer_info_map.get(&peer_id) { + if let Some(info) = swarm.behaviour().discovery.peer_info.get(&peer_id) { peer_manager.update_peer(peer_id, info.clone()); } event_sender.send(NetworkEvent::PeerConnected(peer_id))?; @@ -396,6 +391,8 @@ async fn handle_discovery_event( debug!(%peer_id, ?protocols, "Supported protocols"); + let blueprint_protocol_name = + &swarm.behaviour().blueprint_protocol.blueprint_protocol_name; if !protocols.contains(blueprint_protocol_name) { warn!(%peer_id, %blueprint_protocol_name, "Peer does not support required protocol"); peer_manager @@ -430,8 +427,15 @@ async fn handle_discovery_event( // Process newly discovered peers for peer_info in &ok.peers { if !peer_manager.get_peers().contains_key(&peer_info.peer_id) { + info!(%peer_info.peer_id, "Newly discovered peer from Kademlia"); let info = PeerInfo::default(); peer_manager.update_peer(peer_info.peer_id, info); + let addrs: Vec<_> = peer_info.addrs.iter().cloned().collect(); + for addr in addrs { + if let Err(e) = swarm.dial(DialOpts::from(addr)) { + warn!("Failed to dial address: {}", e); + } + } } } } @@ -439,9 +443,13 @@ async fn handle_discovery_event( // Add newly discovered peers from mDNS for (peer_id, addr) in list { if !peer_manager.get_peers().contains_key(peer_id) { + info!(%peer_id, %addr, "Newly discovered peer from Mdns"); let mut info = PeerInfo::default(); info.addresses.insert(addr.clone()); peer_manager.update_peer(*peer_id, info); + if let Err(e) = swarm.dial(DialOpts::from(addr.clone())) { + warn!("Failed to dial address: {}", e); + } } } } @@ -515,17 +523,10 @@ async fn handle_ping_event( /// Handle a network message async fn handle_network_message( - swarm: &mut Swarm, - msg: NetworkMessage, + _swarm: &mut Swarm, + _msg: NetworkMessage, _peer_manager: &Arc, event_sender: &Sender, ) -> Result<(), Error> { - match msg { - NetworkMessage::Dial(addr) => { - swarm.dial(addr)?; - } - _ => {} - } - Ok(()) } diff --git a/crates/networking/src/tests/mod.rs b/crates/networking/src/tests/mod.rs index 9b864cf32..339d6d901 100644 --- a/crates/networking/src/tests/mod.rs +++ b/crates/networking/src/tests/mod.rs @@ -1,5 +1,5 @@ use crate::service::NetworkMessage; -use std::{collections::HashSet, mem, time::Duration}; +use std::{collections::HashSet, time::Duration}; use gadget_crypto::{sp_core::SpEcdsa, KeyType}; use libp2p::{identity, Multiaddr, PeerId}; From b2f449f93040e125cc00e6ef7bd933e24630c408 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 10:50:43 -0700 Subject: [PATCH 21/52] chore: add handshake tests --- crates/networking/src/tests/handshake.rs | 233 +++++++++++++++++++++-- 1 file changed, 214 insertions(+), 19 deletions(-) diff --git a/crates/networking/src/tests/handshake.rs b/crates/networking/src/tests/handshake.rs index 8077247ea..fa0239822 100644 --- a/crates/networking/src/tests/handshake.rs +++ b/crates/networking/src/tests/handshake.rs @@ -1,13 +1,32 @@ -use super::TestNode; -use super::{init_tracing, wait_for_peer_discovery}; -use super::{wait_for_peer_info, NetworkServiceHandleExt}; -use std::collections::HashSet; -use std::time::Duration; +use crate::{ + blueprint_protocol::{InstanceMessageRequest, InstanceMessageResponse}, + key_types::{Curve, InstanceMsgKeyPair, InstanceMsgPublicKey}, + service::NetworkMessage, + service_handle::NetworkServiceHandle, + tests::TestNode, +}; +use gadget_crypto::KeyType; +use std::{collections::HashSet, time::Duration}; use tokio::time::timeout; +use tracing::{debug, info}; +use tracing_subscriber::{fmt, EnvFilter}; + +const TEST_TIMEOUT: Duration = Duration::from_secs(5); + +fn init_tracing() { + let _ = fmt() + .with_env_filter(EnvFilter::from_default_env()) + .with_target(true) + .with_thread_ids(true) + .with_file(true) + .with_line_number(true) + .try_init(); +} #[tokio::test] -async fn test_peer_handshake() { +async fn test_automatic_handshake() { init_tracing(); + info!("Starting automatic handshake test"); let network_name = "test-network"; let instance_id = "test-instance"; @@ -17,24 +36,200 @@ async fn test_peer_handshake() { let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; - // Start both nodes and wait for them to be listening + info!("Starting nodes"); + // Start both nodes - this should trigger automatic handshake let handle1 = node1.start().await.expect("Failed to start node1"); let handle2 = node2.start().await.expect("Failed to start node2"); - // First wait for basic peer discovery (they see each other) - let discovery_timeout = Duration::from_secs(20); - wait_for_peer_discovery(&[&handle1, &handle2], discovery_timeout) - .await - .expect("Basic peer discovery timed out"); - - handle1.dial(&handle2); + // Wait for automatic handshake completion + info!("Waiting for automatic handshake completion"); + timeout(TEST_TIMEOUT, async { + loop { + if handle1.peer_manager.is_peer_verified(&node2.peer_id) + && handle2.peer_manager.is_peer_verified(&node1.peer_id) + { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Automatic handshake verification timed out"); - let identify_timeout = Duration::from_secs(20); - wait_for_peer_info(&handle1, &handle2, identify_timeout).await; + // Verify peer info and identify info are present + let peer_info1 = handle1 + .peer_info(&node2.peer_id) + .expect("Missing peer info for node2"); + let peer_info2 = handle2 + .peer_info(&node1.peer_id) + .expect("Missing peer info for node1"); - let node2_peer = handle2.local_peer_id; assert!( - handle1.peer_manager.is_peer_verified(&node2_peer), - "Node2 was not verified, handshake failed" + peer_info1.identify_info.is_some(), + "Missing identify info for node2" ); + assert!( + peer_info2.identify_info.is_some(), + "Missing identify info for node1" + ); + + info!("Automatic handshake test completed successfully"); +} + +#[tokio::test] +async fn test_handshake_with_invalid_peer() { + init_tracing(); + info!("Starting invalid peer handshake test"); + + let network_name = "test-network"; + let instance_id = "test-instance"; + + // Create node1 with empty whitelist + let mut node1 = TestNode::new(network_name, instance_id, HashSet::new(), vec![]).await; + + // Create node2 with non-whitelisted key + let mut node2 = TestNode::new(network_name, instance_id, HashSet::new(), vec![]).await; + + info!("Starting nodes"); + let handle1 = node1.start().await.expect("Failed to start node1"); + let handle2 = node2.start().await.expect("Failed to start node2"); + + // Wait for ban to be applied automatically + info!("Waiting for automatic ban"); + timeout(TEST_TIMEOUT, async { + loop { + if handle2.peer_manager.is_banned(&node1.peer_id) { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Ban was not applied"); + + // Verify peers remain unverified + assert!(!handle1.peer_manager.is_peer_verified(&node2.peer_id)); + assert!(!handle2.peer_manager.is_peer_verified(&node1.peer_id)); + + info!("Invalid peer handshake test completed successfully"); } + +#[tokio::test] +async fn test_handshake_reconnection() { + init_tracing(); + info!("Starting handshake reconnection test"); + + let network_name = "test-network"; + let instance_id = "test-instance"; + let allowed_keys = HashSet::new(); + + // Create two nodes + let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; + let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; + + info!("Starting initial connection"); + // Start both nodes and wait for initial handshake + let handle1 = node1.start().await.expect("Failed to start node1"); + let handle2 = node2.start().await.expect("Failed to start node2"); + + // Wait for initial handshake + timeout(TEST_TIMEOUT, async { + loop { + if handle1.peer_manager.is_peer_verified(&node2.peer_id) + && handle2.peer_manager.is_peer_verified(&node1.peer_id) + { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Initial handshake timed out"); + + info!("Disconnecting node2"); + // Drop node2's handle to simulate disconnect + drop(handle2); + tokio::time::sleep(Duration::from_millis(100)).await; + + // Verify node1 sees node2 as disconnected + assert!(!handle1.peer_manager.is_peer_verified(&node2.peer_id)); + + info!("Reconnecting node2"); + // Restart node2 + let handle2 = node2.start().await.expect("Failed to restart node2"); + + // Wait for automatic reconnection and handshake + info!("Waiting for automatic reconnection handshake"); + timeout(TEST_TIMEOUT, async { + loop { + if handle1.peer_manager.is_peer_verified(&node2.peer_id) + && handle2.peer_manager.is_peer_verified(&node1.peer_id) + { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Reconnection handshake timed out"); + + info!("Handshake reconnection test completed successfully"); +} + +// #[tokio::test] +// async fn test_concurrent_connections() { +// init_tracing(); +// info!("Starting concurrent connections test"); + +// let network_name = "test-network"; +// let instance_id = "test-instance"; +// let allowed_keys = HashSet::new(); + +// // Create three nodes to test multiple concurrent handshakes +// let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; +// let mut node2 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; +// let mut node3 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; + +// info!("Starting all nodes simultaneously"); +// // Start all nodes simultaneously +// let (handle1, handle2, handle3) = tokio::join!(node1.start(), node2.start(), node3.start()); +// let handle1 = handle1.expect("Failed to start node1"); +// let handle2 = handle2.expect("Failed to start node2"); +// let handle3 = handle3.expect("Failed to start node3"); + +// // Wait for all handshakes to complete +// info!("Waiting for all handshakes to complete"); +// timeout(TEST_TIMEOUT, async { +// loop { +// let all_verified = handle1.peer_manager.is_peer_verified(&node2.peer_id) +// && handle1.peer_manager.is_peer_verified(&node3.peer_id) +// && handle2.peer_manager.is_peer_verified(&node1.peer_id) +// && handle2.peer_manager.is_peer_verified(&node3.peer_id) +// && handle3.peer_manager.is_peer_verified(&node1.peer_id) +// && handle3.peer_manager.is_peer_verified(&node2.peer_id); + +// if all_verified { +// break; +// } +// tokio::time::sleep(Duration::from_millis(100)).await; +// } +// }) +// .await +// .expect("Concurrent handshakes timed out"); + +// // Verify all peer info is present +// for (handle, peers) in [ +// (&handle1, vec![&node2.peer_id, &node3.peer_id]), +// (&handle2, vec![&node1.peer_id, &node3.peer_id]), +// (&handle3, vec![&node1.peer_id, &node2.peer_id]), +// ] { +// for peer_id in peers { +// assert!( +// handle.peer_info(peer_id).unwrap().identify_info.is_some(), +// "Missing identify info for peer {peer_id:?}" +// ); +// } +// } + +// info!("Concurrent connections test completed successfully"); +// } From f44f1f9000900e189df0d8366672ecff5490f0c2 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 11:03:55 -0700 Subject: [PATCH 22/52] chore: fix --- crates/networking/src/service.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index ad9b00bc8..6d2f1201c 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -432,6 +432,7 @@ async fn handle_discovery_event( peer_manager.update_peer(peer_info.peer_id, info); let addrs: Vec<_> = peer_info.addrs.iter().cloned().collect(); for addr in addrs { + debug!(%addr, "Dialing peer from Kademlia"); if let Err(e) = swarm.dial(DialOpts::from(addr)) { warn!("Failed to dial address: {}", e); } From 4e2bb4416ce0554f39847969ff725007892bb9b9 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 11:05:32 -0700 Subject: [PATCH 23/52] chore: fix --- crates/networking/src/service.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index 6d2f1201c..c71bdb7c0 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -432,7 +432,7 @@ async fn handle_discovery_event( peer_manager.update_peer(peer_info.peer_id, info); let addrs: Vec<_> = peer_info.addrs.iter().cloned().collect(); for addr in addrs { - debug!(%addr, "Dialing peer from Kademlia"); + debug!(%peer_info.peer_id, %addr, "Dialing peer from Kademlia"); if let Err(e) = swarm.dial(DialOpts::from(addr)) { warn!("Failed to dial address: {}", e); } @@ -448,6 +448,7 @@ async fn handle_discovery_event( let mut info = PeerInfo::default(); info.addresses.insert(addr.clone()); peer_manager.update_peer(*peer_id, info); + debug!(%peer_id, %addr, "Dialing peer from Mdns"); if let Err(e) = swarm.dial(DialOpts::from(addr.clone())) { warn!("Failed to dial address: {}", e); } From ea054f62173e310bf4a6b1136473166f213d1647 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 11:10:46 -0700 Subject: [PATCH 24/52] chore: fix --- crates/networking/src/tests/discovery.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/crates/networking/src/tests/discovery.rs b/crates/networking/src/tests/discovery.rs index 7687010f5..37292631e 100644 --- a/crates/networking/src/tests/discovery.rs +++ b/crates/networking/src/tests/discovery.rs @@ -112,8 +112,6 @@ async fn test_peer_info_updates() { .await .expect("Basic peer discovery timed out"); - handle1.dial(&handle2); - // Now wait for identify info to be populated let identify_timeout = Duration::from_secs(20); wait_for_peer_info(&handle1, &handle2, identify_timeout).await; From 770c5c596edffd85aab9f88b039774c404953f30 Mon Sep 17 00:00:00 2001 From: Serial <69764315+Serial-ATA@users.noreply.github.com> Date: Wed, 19 Feb 2025 13:17:15 -0500 Subject: [PATCH 25/52] fix: add back simple handshake verification test --- crates/networking/src/tests/discovery.rs | 2 +- crates/networking/src/tests/handshake.rs | 44 +++++++++++++++++++----- 2 files changed, 36 insertions(+), 10 deletions(-) diff --git a/crates/networking/src/tests/discovery.rs b/crates/networking/src/tests/discovery.rs index 37292631e..7d8dd4709 100644 --- a/crates/networking/src/tests/discovery.rs +++ b/crates/networking/src/tests/discovery.rs @@ -103,7 +103,7 @@ async fn test_peer_info_updates() { info!("Starting nodes..."); // Start both nodes - let mut handle1 = node1.start().await.expect("Failed to start node1"); + let handle1 = node1.start().await.expect("Failed to start node1"); let handle2 = node2.start().await.expect("Failed to start node2"); // First wait for basic peer discovery (they see each other) diff --git a/crates/networking/src/tests/handshake.rs b/crates/networking/src/tests/handshake.rs index fa0239822..60b9ffb31 100644 --- a/crates/networking/src/tests/handshake.rs +++ b/crates/networking/src/tests/handshake.rs @@ -1,9 +1,11 @@ +use super::{ + init_tracing, wait_for_peer_discovery, wait_for_peer_info, NetworkServiceHandleExt, TestNode, +}; use crate::{ blueprint_protocol::{InstanceMessageRequest, InstanceMessageResponse}, key_types::{Curve, InstanceMsgKeyPair, InstanceMsgPublicKey}, service::NetworkMessage, service_handle::NetworkServiceHandle, - tests::TestNode, }; use gadget_crypto::KeyType; use std::{collections::HashSet, time::Duration}; @@ -13,14 +15,38 @@ use tracing_subscriber::{fmt, EnvFilter}; const TEST_TIMEOUT: Duration = Duration::from_secs(5); -fn init_tracing() { - let _ = fmt() - .with_env_filter(EnvFilter::from_default_env()) - .with_target(true) - .with_thread_ids(true) - .with_file(true) - .with_line_number(true) - .try_init(); +#[tokio::test] +async fn test_peer_handshake() { + init_tracing(); + + let network_name = "test-network"; + let instance_id = "test-instance"; + let allowed_keys = HashSet::new(); + + // Create two nodes + let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; + let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; + + // Start both nodes and wait for them to be listening + let handle1 = node1.start().await.expect("Failed to start node1"); + let handle2 = node2.start().await.expect("Failed to start node2"); + + // First wait for basic peer discovery (they see each other) + let discovery_timeout = Duration::from_secs(20); + wait_for_peer_discovery(&[&handle1, &handle2], discovery_timeout) + .await + .expect("Basic peer discovery timed out"); + + handle1.dial(&handle2); + + let identify_timeout = Duration::from_secs(20); + wait_for_peer_info(&handle1, &handle2, identify_timeout).await; + + let node2_peer = handle2.local_peer_id; + assert!( + handle1.peer_manager.is_peer_verified(&node2_peer), + "Node2 was not verified, handshake failed" + ); } #[tokio::test] From b36427ebc23cf1f4934e95619d836afd3b324902 Mon Sep 17 00:00:00 2001 From: Serial <69764315+Serial-ATA@users.noreply.github.com> Date: Wed, 19 Feb 2025 13:21:12 -0500 Subject: [PATCH 26/52] Revert "fix: add back simple handshake verification test" This reverts commit 770c5c596edffd85aab9f88b039774c404953f30. --- crates/networking/src/tests/discovery.rs | 2 +- crates/networking/src/tests/handshake.rs | 44 +++++------------------- 2 files changed, 10 insertions(+), 36 deletions(-) diff --git a/crates/networking/src/tests/discovery.rs b/crates/networking/src/tests/discovery.rs index 7d8dd4709..37292631e 100644 --- a/crates/networking/src/tests/discovery.rs +++ b/crates/networking/src/tests/discovery.rs @@ -103,7 +103,7 @@ async fn test_peer_info_updates() { info!("Starting nodes..."); // Start both nodes - let handle1 = node1.start().await.expect("Failed to start node1"); + let mut handle1 = node1.start().await.expect("Failed to start node1"); let handle2 = node2.start().await.expect("Failed to start node2"); // First wait for basic peer discovery (they see each other) diff --git a/crates/networking/src/tests/handshake.rs b/crates/networking/src/tests/handshake.rs index 60b9ffb31..fa0239822 100644 --- a/crates/networking/src/tests/handshake.rs +++ b/crates/networking/src/tests/handshake.rs @@ -1,11 +1,9 @@ -use super::{ - init_tracing, wait_for_peer_discovery, wait_for_peer_info, NetworkServiceHandleExt, TestNode, -}; use crate::{ blueprint_protocol::{InstanceMessageRequest, InstanceMessageResponse}, key_types::{Curve, InstanceMsgKeyPair, InstanceMsgPublicKey}, service::NetworkMessage, service_handle::NetworkServiceHandle, + tests::TestNode, }; use gadget_crypto::KeyType; use std::{collections::HashSet, time::Duration}; @@ -15,38 +13,14 @@ use tracing_subscriber::{fmt, EnvFilter}; const TEST_TIMEOUT: Duration = Duration::from_secs(5); -#[tokio::test] -async fn test_peer_handshake() { - init_tracing(); - - let network_name = "test-network"; - let instance_id = "test-instance"; - let allowed_keys = HashSet::new(); - - // Create two nodes - let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; - let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; - - // Start both nodes and wait for them to be listening - let handle1 = node1.start().await.expect("Failed to start node1"); - let handle2 = node2.start().await.expect("Failed to start node2"); - - // First wait for basic peer discovery (they see each other) - let discovery_timeout = Duration::from_secs(20); - wait_for_peer_discovery(&[&handle1, &handle2], discovery_timeout) - .await - .expect("Basic peer discovery timed out"); - - handle1.dial(&handle2); - - let identify_timeout = Duration::from_secs(20); - wait_for_peer_info(&handle1, &handle2, identify_timeout).await; - - let node2_peer = handle2.local_peer_id; - assert!( - handle1.peer_manager.is_peer_verified(&node2_peer), - "Node2 was not verified, handshake failed" - ); +fn init_tracing() { + let _ = fmt() + .with_env_filter(EnvFilter::from_default_env()) + .with_target(true) + .with_thread_ids(true) + .with_file(true) + .with_line_number(true) + .try_init(); } #[tokio::test] From a9e4d77657a6fa36873ce7dc98b66355eb0db7c7 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 11:21:15 -0700 Subject: [PATCH 27/52] chore: add back keys to test node --- crates/networking/src/test_helpers.rs | 1 + crates/networking/src/tests/discovery.rs | 1 - crates/networking/src/tests/mod.rs | 43 ++++++++++-------------- 3 files changed, 18 insertions(+), 27 deletions(-) create mode 100644 crates/networking/src/test_helpers.rs diff --git a/crates/networking/src/test_helpers.rs b/crates/networking/src/test_helpers.rs new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/crates/networking/src/test_helpers.rs @@ -0,0 +1 @@ + diff --git a/crates/networking/src/tests/discovery.rs b/crates/networking/src/tests/discovery.rs index 37292631e..1feaae016 100644 --- a/crates/networking/src/tests/discovery.rs +++ b/crates/networking/src/tests/discovery.rs @@ -1,5 +1,4 @@ use super::init_tracing; -use super::NetworkServiceHandleExt; use super::TestNode; use super::{wait_for_peer_discovery, wait_for_peer_info}; use crate::service::NetworkMessage; diff --git a/crates/networking/src/tests/mod.rs b/crates/networking/src/tests/mod.rs index 339d6d901..1dc100a5a 100644 --- a/crates/networking/src/tests/mod.rs +++ b/crates/networking/src/tests/mod.rs @@ -1,16 +1,17 @@ use crate::service::NetworkMessage; -use std::{collections::HashSet, time::Duration}; - +use crate::{ + key_types::InstanceMsgPublicKey, service_handle::NetworkServiceHandle, InstanceMsgKeyPair, + NetworkConfig, NetworkService, +}; use gadget_crypto::{sp_core::SpEcdsa, KeyType}; -use libp2p::{identity, Multiaddr, PeerId}; +use libp2p::{ + identity::{self, Keypair}, + Multiaddr, PeerId, +}; +use std::{collections::HashSet, time::Duration}; use tokio::time::timeout; use tracing::info; -use crate::{ - key_types::InstanceMsgPublicKey, service_handle::NetworkServiceHandle, NetworkConfig, - NetworkService, -}; - mod discovery; mod handshake; @@ -24,11 +25,14 @@ fn init_tracing() { .try_init(); } +/// Test node configuration for network tests /// Test node configuration for network tests pub struct TestNode { pub service: Option, pub peer_id: PeerId, pub listen_addr: Option, + pub instance_key_pair: InstanceMsgKeyPair, + pub local_key: Keypair, } impl TestNode { @@ -49,8 +53,8 @@ impl TestNode { let config = NetworkConfig { network_name: network_name.to_string(), instance_id: instance_id.to_string(), - instance_key_pair, - local_key, + instance_key_pair: instance_key_pair.clone(), + local_key: local_key.clone(), listen_addr: listen_addr.clone(), target_peer_count: 10, bootstrap_peers, @@ -66,7 +70,9 @@ impl TestNode { Self { service: Some(service), peer_id, - listen_addr: None, // To be set later + listen_addr: None, + instance_key_pair, + local_key, } } @@ -174,18 +180,3 @@ pub async fn wait_for_peer_info( Err(_) => panic!("Peer info update timed out"), } } - -trait NetworkServiceHandleExt { - fn dial(&self, other: &NetworkServiceHandle); -} - -impl NetworkServiceHandleExt for NetworkServiceHandle { - fn dial(&self, other: &NetworkServiceHandle) { - let listen_addr = other - .get_listen_addr() - .expect("Node2 did not return a listening address"); - - self.send_network_message(NetworkMessage::Dial(listen_addr.clone())) - .expect("Node1 failed to dial Node2"); - } -} From 8fc3638dbbd54bc4b007e20cff4a09ff70764c73 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 11:56:45 -0700 Subject: [PATCH 28/52] chore: invalid handshake test working --- Cargo.toml | 6 ++ crates/crypto/sp-core/src/lib.rs | 6 ++ .../src/blueprint_protocol/behaviour.rs | 13 ++- .../src/blueprint_protocol/handler.rs | 6 ++ crates/networking/src/tests/discovery.rs | 39 ++++++-- crates/networking/src/tests/handshake.rs | 56 ++++++++--- crates/networking/src/tests/mod.rs | 96 +++++++++++++++---- 7 files changed, 185 insertions(+), 37 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index adc070995..0be443114 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,9 @@ +# This line needs to come before anything else in Cargo.toml +cargo-features = ["codegen-backend"] + +[profile.dev] +codegen-backend = "cranelift" + [workspace] resolver = "2" members = ["cli", "blueprints/*", "crates/*"] diff --git a/crates/crypto/sp-core/src/lib.rs b/crates/crypto/sp-core/src/lib.rs index 6fb82e4c4..4438abcce 100644 --- a/crates/crypto/sp-core/src/lib.rs +++ b/crates/crypto/sp-core/src/lib.rs @@ -118,6 +118,12 @@ macro_rules! impl_sp_core_pair_public { write!(f, "{:?}", self.to_bytes()) } } + + impl gadget_std::fmt::Display for [] { + fn fmt(&self, f: &mut gadget_std::fmt::Formatter<'_>) -> gadget_std::fmt::Result { + write!(f, "{}", hex::encode(self.to_bytes())) + } + } } }; } diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index cf7b6d0b0..19ac90634 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -142,7 +142,14 @@ impl BlueprintProtocolBehaviour { ) -> InstanceSignedMsgSignature { let msg = peer.to_bytes(); match ::sign_with_secret(key_pair, &msg) { - Ok(signature) => signature, + Ok(signature) => { + let public_key = key_pair.public(); + let hex_msg = hex::encode(msg); + let hex_public_key = hex::encode(public_key.0.to_raw()); + let hex_signature = hex::encode(signature.0.to_raw()); + debug!(%peer, ?hex_msg, %hex_public_key, %hex_signature, "signing handshake"); + signature + } Err(e) => { warn!("Failed to sign handshake message: {e}"); InstanceSignedMsgSignature::default() @@ -198,6 +205,10 @@ impl BlueprintProtocolBehaviour { signature: &InstanceSignedMsgSignature, ) -> Result<(), InstanceMessageResponse> { let msg = peer.to_bytes(); + let hex_msg = hex::encode(msg.clone()); + let hex_public_key = hex::encode(public_key.0.to_raw()); + let hex_signature = hex::encode(signature.0.to_raw()); + debug!(%peer, ?hex_msg, %hex_public_key, %hex_signature, "verifying handshake"); let valid = ::verify(public_key, &msg, signature); if !valid { warn!("Invalid initial handshake signature from peer: {peer}"); diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index 4baeaf0d0..4c09c9dba 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -42,6 +42,12 @@ impl BlueprintProtocolBehaviour { debug!(%peer, "Responding to inbound handshake request while outbound is pending"); } + if !self.peer_manager.is_key_whitelisted(&public_key) { + warn!(%peer, %public_key, "Received handshake response from unwhitelisted peer"); + self.peer_manager.handle_nonwhitelisted_peer(&peer); + return; + } + // Verify the handshake match self.verify_handshake(&peer, &public_key, &signature) { Ok(()) => { diff --git a/crates/networking/src/tests/discovery.rs b/crates/networking/src/tests/discovery.rs index 1feaae016..90c3f5efd 100644 --- a/crates/networking/src/tests/discovery.rs +++ b/crates/networking/src/tests/discovery.rs @@ -100,18 +100,41 @@ async fn test_peer_info_updates() { let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; - info!("Starting nodes..."); - // Start both nodes - let mut handle1 = node1.start().await.expect("Failed to start node1"); + info!("Starting node1..."); + let handle1 = node1.start().await.expect("Failed to start node1"); + info!("Node1 started successfully"); + + info!("Starting node2..."); let handle2 = node2.start().await.expect("Failed to start node2"); + info!("Node2 started successfully"); + + info!("Both nodes started, waiting for peer discovery..."); // First wait for basic peer discovery (they see each other) - let discovery_timeout = Duration::from_secs(20); - wait_for_peer_discovery(&[&handle1, &handle2], discovery_timeout) - .await - .expect("Basic peer discovery timed out"); + let discovery_timeout = Duration::from_secs(30); // Increased timeout + match wait_for_peer_discovery(&[&handle1, &handle2], discovery_timeout).await { + Ok(_) => info!("Peer discovery successful"), + Err(e) => { + // Log peer states before failing + info!("Node1 peers: {:?}", handle1.peers()); + info!("Node2 peers: {:?}", handle2.peers()); + panic!("Peer discovery failed: {}", e); + } + } + + info!("Peers discovered each other, waiting for identify info..."); // Now wait for identify info to be populated - let identify_timeout = Duration::from_secs(20); + let identify_timeout = Duration::from_secs(30); // Increased timeout wait_for_peer_info(&handle1, &handle2, identify_timeout).await; + + info!("Test completed successfully - both nodes have identify info"); + + // Log final state + if let Some(info) = handle1.peer_info(&handle2.local_peer_id) { + info!("Node1's info about Node2: {:?}", info); + } + if let Some(info) = handle2.peer_info(&handle1.local_peer_id) { + info!("Node2's info about Node1: {:?}", info); + } } diff --git a/crates/networking/src/tests/handshake.rs b/crates/networking/src/tests/handshake.rs index fa0239822..499648170 100644 --- a/crates/networking/src/tests/handshake.rs +++ b/crates/networking/src/tests/handshake.rs @@ -30,11 +30,27 @@ async fn test_automatic_handshake() { let network_name = "test-network"; let instance_id = "test-instance"; - let allowed_keys = HashSet::new(); - // Create two nodes - let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; - let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; + // Generate node2's key pair first + let instance_key_pair2 = Curve::generate_with_seed(None).unwrap(); + let mut allowed_keys1 = HashSet::new(); + allowed_keys1.insert(instance_key_pair2.public()); + + // Create node1 with node2's key whitelisted + let mut node1 = TestNode::new(network_name, instance_id, allowed_keys1, vec![]).await; + + // Create node2 with node1's key whitelisted and pre-generated key + let mut allowed_keys2 = HashSet::new(); + allowed_keys2.insert(node1.instance_key_pair.public()); + let mut node2 = TestNode::new_with_keys( + network_name, + instance_id, + allowed_keys2, + vec![], + Some(instance_key_pair2), + None, + ) + .await; info!("Starting nodes"); // Start both nodes - this should trigger automatic handshake @@ -87,8 +103,10 @@ async fn test_handshake_with_invalid_peer() { // Create node1 with empty whitelist let mut node1 = TestNode::new(network_name, instance_id, HashSet::new(), vec![]).await; - // Create node2 with non-whitelisted key - let mut node2 = TestNode::new(network_name, instance_id, HashSet::new(), vec![]).await; + // Create node2 with node1's key whitelisted (but node2's key is not whitelisted by node1) + let mut allowed_keys2 = HashSet::new(); + allowed_keys2.insert(node1.instance_key_pair.public()); + let mut node2 = TestNode::new(network_name, instance_id, allowed_keys2, vec![]).await; info!("Starting nodes"); let handle1 = node1.start().await.expect("Failed to start node1"); @@ -98,7 +116,7 @@ async fn test_handshake_with_invalid_peer() { info!("Waiting for automatic ban"); timeout(TEST_TIMEOUT, async { loop { - if handle2.peer_manager.is_banned(&node1.peer_id) { + if handle1.peer_manager.is_banned(&node2.peer_id) { break; } tokio::time::sleep(Duration::from_millis(100)).await; @@ -121,11 +139,27 @@ async fn test_handshake_reconnection() { let network_name = "test-network"; let instance_id = "test-instance"; - let allowed_keys = HashSet::new(); - // Create two nodes - let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; - let mut node2 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; + // Generate node2's key pair first + let instance_key_pair2 = Curve::generate_with_seed(None).unwrap(); + let mut allowed_keys1 = HashSet::new(); + allowed_keys1.insert(instance_key_pair2.public()); + + // Create node1 with node2's key whitelisted + let mut node1 = TestNode::new(network_name, instance_id, allowed_keys1, vec![]).await; + + // Create node2 with node1's key whitelisted and pre-generated key + let mut allowed_keys2 = HashSet::new(); + allowed_keys2.insert(node1.instance_key_pair.public()); + let mut node2 = TestNode::new_with_keys( + network_name, + instance_id, + allowed_keys2, + vec![], + Some(instance_key_pair2), + None, + ) + .await; info!("Starting initial connection"); // Start both nodes and wait for initial handshake diff --git a/crates/networking/src/tests/mod.rs b/crates/networking/src/tests/mod.rs index 1dc100a5a..5aa035330 100644 --- a/crates/networking/src/tests/mod.rs +++ b/crates/networking/src/tests/mod.rs @@ -1,9 +1,10 @@ use crate::service::NetworkMessage; use crate::{ - key_types::InstanceMsgPublicKey, service_handle::NetworkServiceHandle, InstanceMsgKeyPair, + key_types::{Curve, InstanceMsgKeyPair, InstanceMsgPublicKey}, + service_handle::NetworkServiceHandle, NetworkConfig, NetworkService, }; -use gadget_crypto::{sp_core::SpEcdsa, KeyType}; +use gadget_crypto::KeyType; use libp2p::{ identity::{self, Keypair}, Multiaddr, PeerId, @@ -25,7 +26,6 @@ fn init_tracing() { .try_init(); } -/// Test node configuration for network tests /// Test node configuration for network tests pub struct TestNode { pub service: Option, @@ -36,19 +36,42 @@ pub struct TestNode { } impl TestNode { + /// Create a new test node with auto-generated keys pub async fn new( network_name: &str, instance_id: &str, allowed_keys: HashSet, bootstrap_peers: Vec, ) -> Self { - let local_key = identity::Keypair::generate_ed25519(); + Self::new_with_keys( + network_name, + instance_id, + allowed_keys, + bootstrap_peers, + None, + None, + ) + .await + } + + /// Create a new test node with specified keys + pub async fn new_with_keys( + network_name: &str, + instance_id: &str, + allowed_keys: HashSet, + bootstrap_peers: Vec, + instance_key_pair: Option, + local_key: Option, + ) -> Self { + let local_key = local_key.unwrap_or_else(|| identity::Keypair::generate_ed25519()); let peer_id = local_key.public().to_peer_id(); - let listen_addr: Multiaddr = "/ip4/127.0.0.1/tcp/0".parse().unwrap(); + // Bind to all interfaces instead of just localhost + let listen_addr: Multiaddr = "/ip4/0.0.0.0/tcp/0".parse().unwrap(); info!("Creating test node {peer_id} with TCP address: {listen_addr}"); - let instance_key_pair = SpEcdsa::generate_with_seed(None).unwrap(); + let instance_key_pair = + instance_key_pair.unwrap_or_else(|| Curve::generate_with_seed(None).unwrap()); let config = NetworkConfig { network_name: network_name.to_string(), @@ -62,10 +85,8 @@ impl TestNode { enable_kademlia: true, }; - let (allowed_keys_tx, allowed_keys_rx) = crossbeam_channel::unbounded(); - allowed_keys_tx.send(allowed_keys.clone()).unwrap(); - let service = NetworkService::new(config, allowed_keys) // TODO: allowed_keys_rx - .expect("Failed to create network service"); + let service = + NetworkService::new(config, allowed_keys).expect("Failed to create network service"); Self { service: Some(service), @@ -76,28 +97,69 @@ impl TestNode { } } - /// Start the node + /// Start the node and wait for it to be fully initialized pub async fn start(&mut self) -> Result { // Take ownership of the service let service = self.service.take().ok_or("Service already started")?; let handle = service.start(); - // Wait for the actual listening address - let timeout_duration = Duration::from_secs(5); + // Wait for the node to be fully initialized + let timeout_duration = Duration::from_secs(10); // Increased timeout match timeout(timeout_duration, async { + // First wait for the listening address while self.listen_addr.is_none() { if let Some(addr) = handle.get_listen_addr() { info!("Node {} listening on {}", self.peer_id, addr); - self.listen_addr = Some(addr); - break; + self.listen_addr = Some(addr.clone()); + + // Extract port from multiaddr + let addr_str = addr.to_string(); + let port = addr_str.split("/").nth(4).unwrap_or("0").to_string(); + + // Try localhost first + let localhost_addr = format!("127.0.0.1:{}", port); + match tokio::net::TcpStream::connect(&localhost_addr).await { + Ok(_) => { + info!("Successfully verified localhost port for {}", self.peer_id); + break; + } + Err(e) => { + info!("Localhost port not ready for {}: {}", self.peer_id, e); + // Try external IP + let external_addr = format!("10.0.1.142:{}", port); + match tokio::net::TcpStream::connect(&external_addr).await { + Ok(_) => { + info!( + "Successfully verified external port for {}", + self.peer_id + ); + break; + } + Err(e) => { + info!("External port not ready for {}: {}", self.peer_id, e); + tokio::time::sleep(Duration::from_millis(100)).await; + continue; + } + } + } + } } tokio::time::sleep(Duration::from_millis(100)).await; } + + // Give the node a moment to initialize protocols + tokio::time::sleep(Duration::from_millis(500)).await; + + Ok::<(), &'static str>(()) }) .await { - Ok(_) => Ok(handle), - Err(_) => Err("Timeout waiting for node to start listening"), + Ok(Ok(_)) => { + info!("Node {} fully initialized", self.peer_id); + Ok(handle) + } + Ok(Err(e)) => Err(e), + Err(_) => Err("Timeout waiting for node to initialize"), } } From 184bad2d5372b5620bc48cda11209ef3f43ba8b8 Mon Sep 17 00:00:00 2001 From: Shady Khalifa Date: Wed, 19 Feb 2025 21:04:24 +0200 Subject: [PATCH 29/52] fix: remove cranelift from cargo.toml --- Cargo.toml | 6 ------ 1 file changed, 6 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 0be443114..adc070995 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,9 +1,3 @@ -# This line needs to come before anything else in Cargo.toml -cargo-features = ["codegen-backend"] - -[profile.dev] -codegen-backend = "cranelift" - [workspace] resolver = "2" members = ["cli", "blueprints/*", "crates/*"] From 985efbe6636760dcea9a2f5dea4db3734ade4d87 Mon Sep 17 00:00:00 2001 From: Shady Khalifa Date: Wed, 19 Feb 2025 22:03:07 +0200 Subject: [PATCH 30/52] fix: handshakes --- Cargo.lock | 3 + crates/networking/Cargo.toml | 1 + .../src/blueprint_protocol/behaviour.rs | 45 +++++++++----- .../src/blueprint_protocol/handler.rs | 19 ++++-- .../networking/src/blueprint_protocol/mod.rs | 60 +++++++++++++++++++ 5 files changed, 108 insertions(+), 20 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 517005e0c..19dabfe4d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9483,6 +9483,7 @@ dependencies = [ "quick-protobuf-codec", "rand 0.8.5", "regex", + "serde", "sha2 0.10.8", "tracing", "web-time", @@ -9521,6 +9522,7 @@ dependencies = [ "multihash", "quick-protobuf", "rand 0.8.5", + "serde", "sha2 0.10.8", "thiserror 1.0.69", "tracing", @@ -9546,6 +9548,7 @@ dependencies = [ "quick-protobuf", "quick-protobuf-codec", "rand 0.8.5", + "serde", "sha2 0.10.8", "smallvec", "thiserror 2.0.11", diff --git a/crates/networking/Cargo.toml b/crates/networking/Cargo.toml index 28b1b0e9e..eb1426ca4 100644 --- a/crates/networking/Cargo.toml +++ b/crates/networking/Cargo.toml @@ -61,6 +61,7 @@ features = [ "dns", "autonat", "upnp", + "serde", ] [dev-dependencies] diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index 19ac90634..a69e3f326 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -1,4 +1,5 @@ use super::{InstanceMessageRequest, InstanceMessageResponse}; +use crate::blueprint_protocol::HandshakeMessage; use crate::discovery::PeerManager; use crate::{ types::ProtocolMessage, Curve, InstanceMsgKeyPair, InstanceMsgPublicKey, @@ -23,7 +24,7 @@ use std::{ task::Poll, time::{Duration, Instant}, }; -use tracing::{debug, info, trace, warn}; +use tracing::{debug, error, info, trace, warn}; #[derive(NetworkBehaviour)] pub struct DerivedBlueprintProtocolBehaviour { @@ -139,15 +140,17 @@ impl BlueprintProtocolBehaviour { &self, key_pair: &mut InstanceMsgKeyPair, peer: &PeerId, + handshake_msg: &HandshakeMessage, ) -> InstanceSignedMsgSignature { - let msg = peer.to_bytes(); + let msg = handshake_msg.to_bytes(peer); match ::sign_with_secret(key_pair, &msg) { Ok(signature) => { let public_key = key_pair.public(); let hex_msg = hex::encode(msg); let hex_public_key = hex::encode(public_key.0.to_raw()); let hex_signature = hex::encode(signature.0.to_raw()); - debug!(%peer, ?hex_msg, %hex_public_key, %hex_signature, "signing handshake"); + + debug!(%peer, %hex_msg, %hex_public_key, %hex_signature, "signing handshake"); signature } Err(e) => { @@ -200,38 +203,47 @@ impl BlueprintProtocolBehaviour { /// Verify and handle a handshake with a peer pub fn verify_handshake( &self, - peer: &PeerId, + msg: &HandshakeMessage, public_key: &InstanceMsgPublicKey, signature: &InstanceSignedMsgSignature, ) -> Result<(), InstanceMessageResponse> { - let msg = peer.to_bytes(); - let hex_msg = hex::encode(msg.clone()); + if msg.is_expired(HandshakeMessage::MAX_AGE) { + error!(%msg.sender, "Handshake message expired"); + return Err(InstanceMessageResponse::Error { + code: 400, + message: "Handshake message expired".to_string(), + }); + } + let msg_bytes = msg.to_bytes(&self.local_peer_id); + let hex_msg = hex::encode(msg_bytes.clone()); let hex_public_key = hex::encode(public_key.0.to_raw()); let hex_signature = hex::encode(signature.0.to_raw()); - debug!(%peer, ?hex_msg, %hex_public_key, %hex_signature, "verifying handshake"); - let valid = ::verify(public_key, &msg, signature); + debug!(%hex_msg, %hex_public_key, %hex_signature, "verifying handshake"); + + debug!("Verifying handshake with public key: {:?}", hex_public_key); + + let valid = ::verify(public_key, &msg_bytes, signature); if !valid { - warn!("Invalid initial handshake signature from peer: {peer}"); + warn!(%msg.sender, "Invalid handshake signature for peer"); return Err(InstanceMessageResponse::Error { code: 400, message: "Invalid handshake signature".to_string(), }); } - trace!("Received valid handshake from peer: {peer}"); - + trace!(%msg.sender, "Handshake signature verified successfully"); Ok(()) } pub fn handle_handshake( &self, - peer: &PeerId, + msg: &HandshakeMessage, public_key: &InstanceMsgPublicKey, signature: &InstanceSignedMsgSignature, ) -> Result<(), InstanceMessageResponse> { - self.verify_handshake(peer, public_key, signature)?; + self.verify_handshake(msg, public_key, signature)?; self.peer_manager - .add_peer_id_to_public_key(peer, public_key); + .add_peer_id_to_public_key(&msg.sender, public_key); Ok(()) } @@ -334,11 +346,14 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { e.peer_id ); let mut key_pair = self.instance_key_pair.clone(); + let handshake_msg = HandshakeMessage::new(self.local_peer_id); + let signature = self.sign_handshake(&mut key_pair, &e.peer_id, &handshake_msg); self.send_request( &e.peer_id, InstanceMessageRequest::Handshake { public_key: key_pair.public(), - signature: self.sign_handshake(&mut key_pair, &e.peer_id), + signature, + msg: handshake_msg, }, ); self.outbound_handshakes.insert(e.peer_id, Instant::now()); diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index 4c09c9dba..1dca8d56b 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -6,6 +6,7 @@ use libp2p::{ }; use tracing::{debug, warn}; +use crate::blueprint_protocol::HandshakeMessage; use crate::{key_types::InstanceMsgPublicKey, types::ProtocolMessage}; use super::{BlueprintProtocolBehaviour, InstanceMessageRequest, InstanceMessageResponse}; @@ -14,6 +15,7 @@ const INBOUND_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30); const OUTBOUND_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30); impl BlueprintProtocolBehaviour { + #[allow(clippy::too_many_lines)] pub fn handle_request_response_event( &mut self, event: request_response::Event, @@ -27,6 +29,7 @@ impl BlueprintProtocolBehaviour { InstanceMessageRequest::Handshake { public_key, signature, + msg, }, channel, .. @@ -49,7 +52,7 @@ impl BlueprintProtocolBehaviour { } // Verify the handshake - match self.verify_handshake(&peer, &public_key, &signature) { + match self.verify_handshake(&msg, &public_key, &signature) { Ok(()) => { // Store the handshake request self.inbound_handshakes.insert(peer, Instant::now()); @@ -58,9 +61,12 @@ impl BlueprintProtocolBehaviour { // Send handshake response let mut key_pair = self.instance_key_pair.clone(); + let handshake_msg = HandshakeMessage::new(self.local_peer_id); + let signature = self.sign_handshake(&mut key_pair, &peer, &handshake_msg); let response = InstanceMessageResponse::Handshake { - public_key: key_pair.public().clone(), - signature: self.sign_handshake(&mut key_pair, &peer), + public_key: key_pair.public(), + signature, + msg: handshake_msg, }; if let Err(e) = self.send_response(channel, response) { @@ -88,6 +94,7 @@ impl BlueprintProtocolBehaviour { InstanceMessageResponse::Handshake { public_key, signature, + msg, }, .. }, @@ -108,7 +115,7 @@ impl BlueprintProtocolBehaviour { } // Verify the handshake - match self.verify_handshake(&peer, &public_key, &signature) { + match self.verify_handshake(&msg, &public_key, &signature) { Ok(()) => { // Mark handshake as completed self.complete_handshake(&peer, &public_key); @@ -169,7 +176,9 @@ impl BlueprintProtocolBehaviour { }; debug!(%peer, %protocol, %protocol_message, "Received protocol request"); - self.protocol_message_sender.send(protocol_message); + if let Err(e) = self.protocol_message_sender.send(protocol_message) { + warn!(%peer, "Failed to send protocol message: {:?}", e); + } } request_response::Event::Message { peer, diff --git a/crates/networking/src/blueprint_protocol/mod.rs b/crates/networking/src/blueprint_protocol/mod.rs index 8973e8fb5..dc1776f96 100644 --- a/crates/networking/src/blueprint_protocol/mod.rs +++ b/crates/networking/src/blueprint_protocol/mod.rs @@ -2,6 +2,7 @@ mod behaviour; mod handler; pub use behaviour::{BlueprintProtocolBehaviour, BlueprintProtocolEvent}; +use libp2p::PeerId; use crate::key_types::{InstanceMsgPublicKey, InstanceSignedMsgSignature}; use serde::{Deserialize, Serialize}; @@ -15,6 +16,8 @@ pub enum InstanceMessageRequest { public_key: InstanceMsgPublicKey, /// Signature for verification signature: InstanceSignedMsgSignature, + /// Handshake message + msg: HandshakeMessage, }, /// Protocol-specific message with custom payload Protocol { @@ -36,6 +39,8 @@ pub enum InstanceMessageResponse { public_key: InstanceMsgPublicKey, /// Signature for verification signature: InstanceSignedMsgSignature, + /// Handshake message + msg: HandshakeMessage, }, /// Success response with optional data Success { @@ -52,3 +57,58 @@ pub enum InstanceMessageResponse { message: String, }, } + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HandshakeMessage { + /// Sender [`PeerId`] + pub sender: PeerId, + /// A Unix timestamp in milliseconds + pub timestamp: u128, +} + +impl HandshakeMessage { + /// Maximum age for a handshake message in milliseconds + pub const MAX_AGE: u128 = 30_000; + + /// Creates a new handshake message + /// + /// # Panics + /// - If the system time is before the Unix epoch + #[must_use] + pub fn new(sender: PeerId) -> Self { + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("time went backwards") + .as_millis(); + Self { sender, timestamp } + } + + /// Checks if the handshake message is expired + /// + /// # Arguments + /// - `max_age`: Maximum age in milliseconds + /// + /// # Returns + /// - `true` if the message is expired, `false` otherwise + /// + /// # Panics + /// - If the system time is before the Unix epoch + #[must_use] + pub fn is_expired(&self, max_age: u128) -> bool { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("time went backwards") + .as_millis(); + now.saturating_sub(self.timestamp) > max_age + } + + /// Converts the handshake message to a byte array + #[must_use] + pub fn to_bytes(&self, other_peer_id: &PeerId) -> Vec { + let mut bytes = Vec::new(); + bytes.extend(&self.sender.to_bytes()); + bytes.extend(other_peer_id.to_bytes()); + bytes.extend(&self.timestamp.to_be_bytes()); + bytes + } +} From 138246a89d1b2baae82a190e0e6385152a2b5ddc Mon Sep 17 00:00:00 2001 From: Shady Khalifa Date: Wed, 19 Feb 2025 22:23:45 +0200 Subject: [PATCH 31/52] fix: use Sr25519 keys instead --- crates/clients/tangle/src/client.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/clients/tangle/src/client.rs b/crates/clients/tangle/src/client.rs index 135e00dda..cac1a9f6f 100644 --- a/crates/clients/tangle/src/client.rs +++ b/crates/clients/tangle/src/client.rs @@ -154,17 +154,17 @@ impl TangleClient { let parties = self.get_operators().await?; let my_id = self .keystore - .first_local::() + .first_local::() .map_err(Error::Keystore)?; gadget_logging::trace!( - "Looking for {my_id:?} in parties: {:?}", + "Looking for {my_id} in parties: {:?}", parties.keys().collect::>() ); let index_of_my_id = parties .iter() - .position(|(_id, key)| key == &my_id.0) + .position(|(id, _key)| id.0 == my_id.0.to_raw()) .ok_or(Error::PartyNotFound)?; Ok((index_of_my_id, parties)) From 0bb024b5defbe1b461d881388e13d0cd1962901c Mon Sep 17 00:00:00 2001 From: Alex <69764315+Serial-ATA@users.noreply.github.com> Date: Wed, 19 Feb 2025 15:33:39 -0500 Subject: [PATCH 32/52] feat(networking): round-based tests (#669) * fix(networking): stop using `InstanceSignedMsgSignature::default()` * fix(networking): remove old trait * feat(networking): re-export gadget_crypto::KeyType * fix(crypto): feature parity for k256 keys * fix(networking): start round-based test * feat(crypto): implement Display for sp-core signatures --- Cargo.lock | 1 + crates/crypto/k256/src/lib.rs | 73 ++++++++----------- crates/crypto/sp-core/src/lib.rs | 6 ++ .../Cargo.toml | 24 ++++++ .../src/lib.rs | 3 + .../src/tests.rs | 61 ++++++++++++++++ crates/networking/Cargo.toml | 1 - .../src/blueprint_protocol/behaviour.rs | 22 +++--- .../src/blueprint_protocol/handler.rs | 8 +- crates/networking/src/lib.rs | 7 +- 10 files changed, 144 insertions(+), 62 deletions(-) create mode 100644 crates/networking-round-based-extension/src/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 19dabfe4d..67d219022 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7178,6 +7178,7 @@ dependencies = [ "dashmap", "futures", "gadget-networking", + "libp2p", "round-based", "serde", "serde_json", diff --git a/crates/crypto/k256/src/lib.rs b/crates/crypto/k256/src/lib.rs index c848bcf01..55c47a8e7 100644 --- a/crates/crypto/k256/src/lib.rs +++ b/crates/crypto/k256/src/lib.rs @@ -5,6 +5,7 @@ pub mod error; #[cfg(test)] mod tests; +use std::hash::{Hash, Hasher}; use crate::error::{K256Error, Result}; use alloy_signer_local::LocalSigner; use gadget_crypto_core::KeyEncoding; @@ -13,50 +14,10 @@ use gadget_std::string::{String, ToString}; use gadget_std::UniformRand; use k256::ecdsa::signature::SignerMut; use k256::ecdsa::{SigningKey, VerifyingKey}; -use serde::{Deserialize, Serialize}; /// ECDSA key type pub struct K256Ecdsa; -#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize)] -pub struct K256VerifyingKey(pub VerifyingKey); - -impl From for VerifyingKey { - fn from(key: K256VerifyingKey) -> Self { - key.0 - } -} - -impl KeyEncoding for K256VerifyingKey { - fn to_bytes(&self) -> Vec { - self.0.to_sec1_bytes().to_vec() - } - - fn from_bytes(bytes: &[u8]) -> core::result::Result { - let vk = VerifyingKey::from_sec1_bytes(bytes) - .map_err(|e| serde::de::Error::custom(e.to_string()))?; - Ok(K256VerifyingKey(vk)) - } -} - -impl PartialOrd for K256VerifyingKey { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for K256VerifyingKey { - fn cmp(&self, other: &Self) -> gadget_std::cmp::Ordering { - self.0.to_sec1_bytes().cmp(&other.0.to_sec1_bytes()) - } -} - -impl gadget_std::hash::Hash for K256VerifyingKey { - fn hash(&self, state: &mut H) { - self.0.to_sec1_bytes().hash(state); - } -} - macro_rules! impl_serde_bytes { ($wrapper:ident, $inner:path) => { #[derive(Clone, PartialEq, Eq, Debug)] @@ -101,14 +62,42 @@ macro_rules! impl_serde_bytes { D: serde::Deserializer<'de>, { let bytes = Vec::::deserialize(deserializer)?; - let inner = <$inner>::from_slice(&bytes) + let ret = <$wrapper as KeyEncoding>::from_bytes(&bytes) .map_err(|e| serde::de::Error::custom(e.to_string()))?; - Ok($wrapper(inner)) + Ok(ret) + } + } + + impl gadget_std::fmt::Display for $wrapper { + fn fmt(&self, f: &mut gadget_std::fmt::Formatter<'_>) -> gadget_std::fmt::Result { + write!(f, "{}", hex::encode(self.to_bytes())) } } }; } +impl_serde_bytes!(K256VerifyingKey, k256::ecdsa::VerifyingKey); + +impl K256VerifyingKey { + fn to_bytes_impl(&self) -> Vec { + self.0.to_sec1_bytes().to_vec() + } + + fn from_bytes_impl(bytes: &[u8]) -> Result { + let vk = VerifyingKey::from_sec1_bytes(bytes) + .map_err(|e| K256Error::InvalidSigner(e.to_string()))?; + Ok(K256VerifyingKey(vk)) + } +} + +impl Hash for K256VerifyingKey { + fn hash(&self, state: &mut H) { + state.write(self.to_bytes().as_slice()); + } +} + +impl Copy for K256VerifyingKey {} + impl_serde_bytes!(K256SigningKey, k256::ecdsa::SigningKey); impl K256SigningKey { diff --git a/crates/crypto/sp-core/src/lib.rs b/crates/crypto/sp-core/src/lib.rs index 4438abcce..d41a67932 100644 --- a/crates/crypto/sp-core/src/lib.rs +++ b/crates/crypto/sp-core/src/lib.rs @@ -152,6 +152,12 @@ macro_rules! impl_sp_core_signature { write!(f, "{:?}", self.0.0) } } + + impl gadget_std::fmt::Display for [] { + fn fmt(&self, f: &mut gadget_std::fmt::Formatter<'_>) -> gadget_std::fmt::Result { + write!(f, "{}", hex::encode(self.to_bytes())) + } + } } }; } diff --git a/crates/networking-round-based-extension/Cargo.toml b/crates/networking-round-based-extension/Cargo.toml index 1ca17c5d3..1ee2231ad 100644 --- a/crates/networking-round-based-extension/Cargo.toml +++ b/crates/networking-round-based-extension/Cargo.toml @@ -20,5 +20,29 @@ crossbeam = { workspace = true } crossbeam-channel = { workspace = true } thiserror = { workspace = true } +[dev-dependencies] +round-based = { workspace = true, features = ["derive"] } +libp2p = { workspace = true, features = [ + "tokio", + "gossipsub", + "mdns", + "noise", + "macros", + "yamux", + "tcp", + "quic", + "request-response", + "cbor", + "identify", + "kad", + "dcutr", + "relay", + "ping", + "dns", + "autonat", + "upnp", +] } + + [lints] workspace = true diff --git a/crates/networking-round-based-extension/src/lib.rs b/crates/networking-round-based-extension/src/lib.rs index 892a0ba08..2eaddc5ad 100644 --- a/crates/networking-round-based-extension/src/lib.rs +++ b/crates/networking-round-based-extension/src/lib.rs @@ -1,3 +1,6 @@ +#[cfg(test)] +mod tests; + use crossbeam_channel::{self, Receiver, Sender}; use dashmap::DashMap; use futures::Future; diff --git a/crates/networking-round-based-extension/src/tests.rs b/crates/networking-round-based-extension/src/tests.rs new file mode 100644 index 000000000..c8cb23149 --- /dev/null +++ b/crates/networking-round-based-extension/src/tests.rs @@ -0,0 +1,61 @@ +use std::collections::HashSet; +use libp2p::Multiaddr; +use round_based::ProtocolMessage; +use serde::{Deserialize, Serialize}; +use gadget_networking::{KeyType, NetworkConfig, NetworkService}; + +#[derive(Debug, Serialize, Deserialize, Clone, ProtocolMessage)] +enum Msg { + Round1(Round1Msg), + Round2(Round2Msg), + Round3(Round3Msg), +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +struct Round1Msg { + pub power: u16, + pub hitpoints: u16, + pub armor: u16, + pub name: String, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +struct Round2Msg { + pub x: u16, + pub y: u16, + pub z: u16, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +struct Round3Msg { + rotation: u16, + velocity: (u16, u16, u16), +} + +const TOPIC: &str = "/gadget/test/1.0.0"; + +fn node() -> NetworkService { + let local_key = libp2p::identity::Keypair::generate_ed25519(); + let listen_addr: Multiaddr = "/ip4/127.0.0.1/tcp/0".parse().unwrap(); + + let instance_key_pair = gadget_networking::Curve::generate_with_seed(None).unwrap(); + + let config = NetworkConfig { + network_name: TOPIC.to_string(), + instance_id: String::from("0"), + instance_key_pair, + local_key, + listen_addr, + target_peer_count: 0, + bootstrap_peers: vec![], + enable_mdns: true, + enable_kademlia: true, + }; + + NetworkService::new(config, HashSet::default()).unwrap() +} + +#[tokio::test] +async fn round_based() { + let node = node(); +} \ No newline at end of file diff --git a/crates/networking/Cargo.toml b/crates/networking/Cargo.toml index eb1426ca4..46917174c 100644 --- a/crates/networking/Cargo.toml +++ b/crates/networking/Cargo.toml @@ -39,7 +39,6 @@ gadget-crypto = { workspace = true, features = ["k256", "hashing"] } gadget-crypto-core = { workspace = true, features = ["k256"] } k256 = { workspace = true } - [target.'cfg(not(target_family = "wasm"))'.dependencies.libp2p] workspace = true features = [ diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index a69e3f326..f95ce62d9 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -141,21 +141,19 @@ impl BlueprintProtocolBehaviour { key_pair: &mut InstanceMsgKeyPair, peer: &PeerId, handshake_msg: &HandshakeMessage, - ) -> InstanceSignedMsgSignature { + ) -> Option { let msg = handshake_msg.to_bytes(peer); match ::sign_with_secret(key_pair, &msg) { Ok(signature) => { let public_key = key_pair.public(); let hex_msg = hex::encode(msg); - let hex_public_key = hex::encode(public_key.0.to_raw()); - let hex_signature = hex::encode(signature.0.to_raw()); - debug!(%peer, %hex_msg, %hex_public_key, %hex_signature, "signing handshake"); - signature + debug!(%peer, ?hex_msg, %public_key, %signature, "signing handshake"); + Some(signature) } Err(e) => { warn!("Failed to sign handshake message: {e}"); - InstanceSignedMsgSignature::default() + None } } } @@ -214,13 +212,11 @@ impl BlueprintProtocolBehaviour { message: "Handshake message expired".to_string(), }); } + let msg_bytes = msg.to_bytes(&self.local_peer_id); let hex_msg = hex::encode(msg_bytes.clone()); - let hex_public_key = hex::encode(public_key.0.to_raw()); - let hex_signature = hex::encode(signature.0.to_raw()); - debug!(%hex_msg, %hex_public_key, %hex_signature, "verifying handshake"); - debug!("Verifying handshake with public key: {:?}", hex_public_key); + debug!(%hex_msg, %public_key, %signature, "verifying handshake"); let valid = ::verify(public_key, &msg_bytes, signature); if !valid { @@ -346,8 +342,12 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { e.peer_id ); let mut key_pair = self.instance_key_pair.clone(); + let handshake_msg = HandshakeMessage::new(self.local_peer_id); - let signature = self.sign_handshake(&mut key_pair, &e.peer_id, &handshake_msg); + let Some(signature) = self.sign_handshake(&mut key_pair, &e.peer_id, &handshake_msg) else { + return; + }; + self.send_request( &e.peer_id, InstanceMessageRequest::Handshake { diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index 1dca8d56b..83a70d1a6 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -61,10 +61,14 @@ impl BlueprintProtocolBehaviour { // Send handshake response let mut key_pair = self.instance_key_pair.clone(); + let handshake_msg = HandshakeMessage::new(self.local_peer_id); - let signature = self.sign_handshake(&mut key_pair, &peer, &handshake_msg); + let Some(signature) = self.sign_handshake(&mut key_pair, &peer, &handshake_msg) else { + return; + }; + let response = InstanceMessageResponse::Handshake { - public_key: key_pair.public(), + public_key: key_pair.public().clone(), signature, msg: handshake_msg, }; diff --git a/crates/networking/src/lib.rs b/crates/networking/src/lib.rs index 52f445344..003985975 100644 --- a/crates/networking/src/lib.rs +++ b/crates/networking/src/lib.rs @@ -17,6 +17,7 @@ pub use gadget_networking_round_based_extension as round_based_compat; pub use key_types::*; pub use service::{NetworkConfig, NetworkEvent, NetworkService}; +pub use gadget_crypto::KeyType; #[cfg(all( feature = "sp-core-ecdsa", @@ -65,12 +66,6 @@ pub mod key_types { K256Ecdsa as Curve, K256Signature as InstanceSignedMsgSignature, K256SigningKey as InstanceMsgKeyPair, K256VerifyingKey as InstanceMsgPublicKey, }; - - impl super::KeySignExt for InstanceMsgKeyPair { - fn sign_prehash(&self, prehash: &[u8; 32]) -> InstanceSignedMsgSignature { - self.sign_prehash(prehash) - } - } } // Compile-time assertion to ensure only one feature is enabled From 8a60e1c4904c2611faae8d3d29628c4c378e5ae9 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 13:34:45 -0700 Subject: [PATCH 33/52] chore: fix and start protocol tests --- Cargo.toml | 5 + crates/crypto/sp-core/src/lib.rs | 2 +- crates/networking/README.md | 42 +- .../src/tests/blueprint_protocol.rs | 607 ++++++++++++++++++ crates/networking/src/tests/handshake.rs | 114 ++-- crates/networking/src/tests/mod.rs | 8 + 6 files changed, 711 insertions(+), 67 deletions(-) create mode 100644 crates/networking/src/tests/blueprint_protocol.rs diff --git a/Cargo.toml b/Cargo.toml index adc070995..b1b9d247c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,8 @@ +cargo-features = ["codegen-backend"] + +[profile.dev] +codegen-backend = "cranelift" + [workspace] resolver = "2" members = ["cli", "blueprints/*", "crates/*"] diff --git a/crates/crypto/sp-core/src/lib.rs b/crates/crypto/sp-core/src/lib.rs index d41a67932..ce4c16c70 100644 --- a/crates/crypto/sp-core/src/lib.rs +++ b/crates/crypto/sp-core/src/lib.rs @@ -155,7 +155,7 @@ macro_rules! impl_sp_core_signature { impl gadget_std::fmt::Display for [] { fn fmt(&self, f: &mut gadget_std::fmt::Formatter<'_>) -> gadget_std::fmt::Result { - write!(f, "{}", hex::encode(self.to_bytes())) + write!(f, "{}", hex::encode(self.0.0)) } } } diff --git a/crates/networking/README.md b/crates/networking/README.md index 99b4ea8ae..88136d528 100644 --- a/crates/networking/README.md +++ b/crates/networking/README.md @@ -13,21 +13,31 @@ sequenceDiagram Note over A,B: Initial TCP/QUIC Connection Established - Note over A: Generate signature of B's peer ID + Note over A: Create handshake message with:
1. A's peer ID
2. Current timestamp + Note over A: Sign(A_id | B_id | timestamp) A->>+B: HandshakeRequest { public_key: A_pub, - signature: sign(B_peer_id) + signature: sign(msg), + msg: HandshakeMessage { + sender: A_id, + timestamp: now + } } - Note over B: 1. Verify A's signature
2. Store A's public key + Note over B: 1. Verify timestamp is fresh
2. Verify A_pub derives to A_id
3. Verify signature
4. Store A's public key - Note over B: Generate signature of A's peer ID + Note over B: Create handshake message with:
1. B's peer ID
2. Current timestamp + Note over B: Sign(B_id | A_id | timestamp) B-->>-A: HandshakeResponse { public_key: B_pub, - signature: sign(A_peer_id) + signature: sign(msg), + msg: HandshakeMessage { + sender: B_id, + timestamp: now + } } - Note over A: 1. Verify B's signature
2. Store B's public key + Note over A: 1. Verify timestamp is fresh
2. Verify B_pub derives to B_id
3. Verify signature
4. Store B's public key Note over A,B: ✓ Handshake Complete Note over A,B: ✓ Protocol Messages Allowed @@ -43,12 +53,15 @@ stateDiagram-v2 Connected --> OutboundPending: Send Handshake Connected --> InboundPending: Receive Handshake - OutboundPending --> Verified: Valid Response + OutboundPending --> Verifying: Valid Response OutboundPending --> Failed: Invalid/Timeout - InboundPending --> Verified: Valid Request & Response + InboundPending --> Verifying: Valid Request InboundPending --> Failed: Invalid/Timeout + Verifying --> Verified: All Checks Pass + Verifying --> Failed: Checks Fail + Verified --> [*]: Connection Closed Failed --> [*]: Connection Closed @@ -56,6 +69,14 @@ stateDiagram-v2 Initial TCP/QUIC connection established end note + note right of Verifying + Checks: + 1. Timestamp fresh + 2. PubKey matches PeerId + 3. Signature valid + 4. Key whitelisted + end note + note right of Verified Both peers authenticated Protocol messages allowed @@ -130,7 +151,8 @@ stateDiagram-v2 - Initiated on first connection - Mutual authentication using public key cryptography -- Signatures verify peer identity +- Signatures verify peer identity and ownership +- Timestamps prevent replay attacks - Timeouts after 30 seconds - Handles concurrent handshakes gracefully @@ -148,6 +170,8 @@ stateDiagram-v2 ### Security Features - Peer verification before message acceptance +- Public key to peer ID verification +- Timestamp-based replay protection - Signature verification for handshakes - Banned peer tracking - Connection limits diff --git a/crates/networking/src/tests/blueprint_protocol.rs b/crates/networking/src/tests/blueprint_protocol.rs new file mode 100644 index 000000000..01880f3fc --- /dev/null +++ b/crates/networking/src/tests/blueprint_protocol.rs @@ -0,0 +1,607 @@ +use gadget_crypto::KeyType; +use libp2p::PeerId; +use serde::{Deserialize, Serialize}; +use std::{collections::HashSet, time::Duration}; +use tokio::time::timeout; +use tracing::info; + +use crate::{ + blueprint_protocol::{InstanceMessageRequest, InstanceMessageResponse}, + key_types::{Curve, InstanceMsgKeyPair, InstanceMsgPublicKey}, + service::NetworkMessage, + service_handle::NetworkServiceHandle, + tests::TestNode, + types::{MessageRouting, ParticipantId, ParticipantInfo, ProtocolMessage}, +}; + +const TEST_TIMEOUT: Duration = Duration::from_secs(10); +const PROTOCOL_NAME: &str = "summation/1.0.0"; + +// Protocol message types +#[derive(Debug, Clone, Serialize, Deserialize)] +enum SummationMessage { + Number(u64), + Verification { sum: u64 }, +} + +// Helper to create a whitelisted test node +async fn create_node_with_keys( + network: &str, + instance: &str, + allowed_keys: HashSet, + key_pair: Option, +) -> TestNode { + TestNode::new_with_keys(network, instance, allowed_keys, vec![], key_pair, None).await +} + +// Helper to create a protocol message +fn create_protocol_message( + protocol: &str, + msg: SummationMessage, + sender: &NetworkServiceHandle, + recipient: Option, +) -> ProtocolMessage { + ProtocolMessage { + protocol: protocol.to_string(), + routing: MessageRouting { + message_id: 0, + round_id: 0, + sender: ParticipantInfo { + id: ParticipantId(0), + public_key: None, + }, + recipient: recipient.map(|peer_id| ParticipantInfo { + id: ParticipantId(0), + public_key: None, + }), + }, + payload: bincode::serialize(&msg).expect("Failed to serialize message"), + } +} + +// Helper to wait for handshake completion between multiple nodes +async fn wait_for_all_handshakes(handles: &[&NetworkServiceHandle]) { + timeout(TEST_TIMEOUT, async { + loop { + let mut all_verified = true; + for (i, h1) in handles.iter().enumerate() { + for (j, h2) in handles.iter().enumerate() { + if i != j && !h1.peer_manager.is_peer_verified(&h2.local_peer_id) { + all_verified = false; + break; + } + } + } + if all_verified { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Handshake verification timed out"); +} + +// Helper to wait for handshake completion between two nodes +async fn wait_for_handshake_completion( + handle1: &NetworkServiceHandle, + handle2: &NetworkServiceHandle, +) { + timeout(TEST_TIMEOUT, async { + loop { + if handle1 + .peer_manager + .is_peer_verified(&handle2.local_peer_id) + && handle2 + .peer_manager + .is_peer_verified(&handle1.local_peer_id) + { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Handshake verification timed out"); +} + +// Helper to extract number from message +fn extract_number_from_message(msg: &ProtocolMessage) -> u64 { + match bincode::deserialize::(&msg.payload).expect("Failed to deserialize") { + SummationMessage::Number(n) => n, + _ => panic!("Expected number message"), + } +} + +// Helper to extract sum from verification message +fn extract_sum_from_verification(msg: &ProtocolMessage) -> u64 { + match bincode::deserialize::(&msg.payload).expect("Failed to deserialize") { + SummationMessage::Verification { sum } => sum, + _ => panic!("Expected verification message"), + } +} + +// Helper to create a set of nodes with whitelisted keys +async fn create_whitelisted_nodes(count: usize) -> Vec { + let mut nodes = Vec::with_capacity(count); + let mut key_pairs = Vec::with_capacity(count); + let mut allowed_keys = vec![HashSet::new(); count]; + + // Generate all key pairs first + for _ in 0..count { + key_pairs.push(Curve::generate_with_seed(None).unwrap()); + } + + // Create allowed keys sets + for i in 0..count { + for j in 0..count { + if i != j { + allowed_keys[i].insert(key_pairs[j].public()); + } + } + } + + // Create nodes with whitelisted keys + for i in 0..count { + nodes.push( + create_node_with_keys( + "test-net", + "sum-test", + allowed_keys[i].clone(), + Some(key_pairs[i].clone()), + ) + .await, + ); + } + + nodes +} + +#[tokio::test] +async fn test_summation_protocol_basic() { + super::init_tracing(); + info!("Starting summation protocol test"); + + // Create nodes with whitelisted keys + let instance_key_pair2 = Curve::generate_with_seed(None).unwrap(); + let mut allowed_keys1 = HashSet::new(); + allowed_keys1.insert(instance_key_pair2.public()); + + let mut node1 = TestNode::new("test-net", "sum-test", allowed_keys1, vec![]).await; + + let mut allowed_keys2 = HashSet::new(); + allowed_keys2.insert(node1.instance_key_pair.public()); + let mut node2 = TestNode::new_with_keys( + "test-net", + "sum-test", + allowed_keys2, + vec![], + Some(instance_key_pair2), + None, + ) + .await; + + info!("Starting nodes"); + let mut handle1 = node1.start().await.expect("Failed to start node1"); + let mut handle2 = node2.start().await.expect("Failed to start node2"); + + info!("Waiting for handshake completion"); + wait_for_handshake_completion(&handle1, &handle2).await; + + // Generate test numbers + let num1 = 42; + let num2 = 58; + let expected_sum = num1 + num2; + + info!("Sending numbers via gossip"); + // Send numbers via gossip + handle1 + .send_protocol_message(create_protocol_message( + PROTOCOL_NAME, + SummationMessage::Number(num1), + &handle1, + None, + )) + .expect("Failed to send number from node1"); + + handle2 + .send_protocol_message(create_protocol_message( + PROTOCOL_NAME, + SummationMessage::Number(num2), + &handle2, + None, + )) + .expect("Failed to send number from node2"); + + info!("Waiting for messages to be processed"); + // Wait for messages and compute sums + let mut sum1 = 0; + let mut sum2 = 0; + let mut node1_received = false; + let mut node2_received = false; + + timeout(TEST_TIMEOUT, async { + loop { + // Process incoming messages + if let Some(msg) = handle1.next_protocol_message() { + if !node1_received { + sum1 += extract_number_from_message(&msg); + node1_received = true; + } + } + if let Some(msg) = handle2.next_protocol_message() { + if !node2_received { + sum2 += extract_number_from_message(&msg); + node2_received = true; + } + } + + // Check if both nodes have received messages + if node1_received && node2_received { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Timeout waiting for summation completion"); + + info!("Verifying sums via P2P messages"); + // Verify sums via P2P messages + handle1 + .send_protocol_message(create_protocol_message( + PROTOCOL_NAME, + SummationMessage::Verification { sum: sum1 }, + &handle1, + Some(node2.peer_id), + )) + .expect("Failed to send verification from node1"); + + handle2 + .send_protocol_message(create_protocol_message( + PROTOCOL_NAME, + SummationMessage::Verification { sum: sum2 }, + &handle2, + Some(node1.peer_id), + )) + .expect("Failed to send verification from node2"); + + info!("Waiting for verification messages"); + // Wait for verification messages + timeout(TEST_TIMEOUT, async { + let mut node1_verified = false; + let mut node2_verified = false; + + loop { + // Process verification messages + if let Some(msg) = handle1.next_protocol_message() { + if !node1_verified { + assert_eq!(extract_sum_from_verification(&msg), expected_sum); + node1_verified = true; + } + } + if let Some(msg) = handle2.next_protocol_message() { + if !node2_verified { + assert_eq!(extract_sum_from_verification(&msg), expected_sum); + node2_verified = true; + } + } + + if node1_verified && node2_verified { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Timeout waiting for verification completion"); + + info!("Summation protocol test completed successfully"); +} + +#[tokio::test] +async fn test_summation_protocol_multi_node() { + super::init_tracing(); + info!("Starting multi-node summation protocol test"); + + // Create 3 nodes with whitelisted keys + let mut nodes = create_whitelisted_nodes(3).await; + + // Start all nodes + let mut handles = Vec::new(); + for node in nodes.iter_mut() { + handles.push(node.start().await.expect("Failed to start node")); + } + + // Convert handles to mutable references + let handles: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); + + // Wait for all handshakes to complete + info!("Waiting for handshake completion"); + wait_for_all_handshakes(&handles).await; + + // Generate test numbers + let numbers = vec![42, 58, 100]; + let expected_sum: u64 = numbers.iter().sum(); + + info!("Sending numbers via gossip"); + // Each node broadcasts its number + for (i, handle) in handles.iter().enumerate() { + handle + .send_protocol_message(create_protocol_message( + PROTOCOL_NAME, + SummationMessage::Number(numbers[i]), + handle, + None, + )) + .expect("Failed to send number"); + } + + info!("Waiting for messages to be processed"); + // Wait for all nodes to receive all numbers + let mut sums = vec![0; handles.len()]; + let mut received = vec![0; handles.len()]; + + timeout(TEST_TIMEOUT, async { + loop { + for (i, handle) in handles.iter().enumerate() { + if let Some(msg) = handle.next_protocol_message() { + if received[i] < handles.len() - 1 { + sums[i] += extract_number_from_message(&msg); + received[i] += 1; + } + } + } + + if received.iter().all(|&r| r == handles.len() - 1) { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Timeout waiting for summation completion"); + + info!("Verifying sums via P2P messages"); + // Each node verifies with every other node + for (i, sender) in handles.iter().enumerate() { + for (j, recipient) in handles.iter().enumerate() { + if i != j { + sender + .send_protocol_message(create_protocol_message( + PROTOCOL_NAME, + SummationMessage::Verification { sum: sums[i] }, + sender, + Some(recipient.local_peer_id), + )) + .expect("Failed to send verification"); + } + } + } + + info!("Waiting for verification messages"); + // Wait for all verifications + timeout(TEST_TIMEOUT, async { + let mut verified = vec![0; handles.len()]; + loop { + for (i, handle) in handles.iter().enumerate() { + if let Some(msg) = handle.next_protocol_message() { + if verified[i] < handles.len() - 1 { + assert_eq!(extract_sum_from_verification(&msg), expected_sum); + verified[i] += 1; + } + } + } + + if verified.iter().all(|&v| v == handles.len() - 1) { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Timeout waiting for verification completion"); + + info!("Multi-node summation protocol test completed successfully"); +} + +#[tokio::test] +async fn test_summation_protocol_late_join() { + super::init_tracing(); + info!("Starting late join summation protocol test"); + + // Create 3 nodes but only start 2 initially + let mut nodes = create_whitelisted_nodes(3).await; + + // Start first two nodes + let mut handles = Vec::new(); + for node in nodes[..2].iter_mut() { + handles.push(node.start().await.expect("Failed to start node")); + } + + // Convert handles to mutable references + let mut handles_refs: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); + + // Wait for initial handshakes + info!("Waiting for initial handshake completion"); + wait_for_all_handshakes(&handles_refs).await; + + // Initial nodes send their numbers + let numbers = vec![42, 58, 100]; + let expected_sum: u64 = numbers.iter().sum(); + + info!("Initial nodes sending numbers"); + for (i, handle) in handles.iter().enumerate() { + handle + .send_protocol_message(create_protocol_message( + PROTOCOL_NAME, + SummationMessage::Number(numbers[i]), + handle, + None, + )) + .expect("Failed to send number"); + } + + // Wait for initial nodes to process messages + timeout(TEST_TIMEOUT, async { + let mut received = vec![false; 2]; + loop { + for (i, handle) in handles.iter().enumerate() { + if let Some(msg) = handle.next_protocol_message() { + if !received[i] { + assert_eq!(extract_number_from_message(&msg), numbers[1 - i]); + received[i] = true; + } + } + } + if received.iter().all(|&r| r) { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Timeout waiting for initial summation"); + + // Start the late joining node + info!("Starting late joining node"); + handles.push(nodes[2].start().await.expect("Failed to start late node")); + let all_handles: Vec<&NetworkServiceHandle> = handles.iter().collect(); + + // Wait for the new node to complete handshakes + wait_for_all_handshakes(&all_handles).await; + + // Late node sends its number and receives history + info!("Late node sending number and receiving history"); + handles[2] + .send_protocol_message(create_protocol_message( + PROTOCOL_NAME, + SummationMessage::Number(numbers[2]), + &handles[2], + None, + )) + .expect("Failed to send number from late node"); + + // Verify final state + timeout(TEST_TIMEOUT, async { + let mut verified = vec![false; handles.len()]; + loop { + for (i, handle) in handles.iter().enumerate() { + if let Some(msg) = handle.next_protocol_message() { + let num = extract_number_from_message(&msg); + if !verified[i] && (num == numbers[2] || numbers.contains(&num)) { + verified[i] = true; + } + } + } + if verified.iter().all(|&v| v) { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Timeout waiting for late node synchronization"); + + info!("Late join test completed successfully"); +} + +#[tokio::test] +async fn test_summation_protocol_node_disconnect() { + super::init_tracing(); + info!("Starting node disconnect test"); + + // Create 3 nodes + let mut nodes = create_whitelisted_nodes(3).await; + + // Start all nodes + let mut handles = Vec::new(); + for node in nodes.iter_mut() { + handles.push(node.start().await.expect("Failed to start node")); + } + + // Convert handles to mutable references + let mut handles_refs: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); + + // Wait for all handshakes + wait_for_all_handshakes(&handles_refs).await; + + // Send initial numbers + let numbers = vec![42, 58, 100]; + for (i, handle) in handles.iter().enumerate() { + handle + .send_protocol_message(create_protocol_message( + PROTOCOL_NAME, + SummationMessage::Number(numbers[i]), + handle, + None, + )) + .expect("Failed to send number"); + } + + // Wait for initial processing + timeout(TEST_TIMEOUT, async { + let mut received = vec![0; handles.len()]; + loop { + for (i, handle) in handles.iter().enumerate() { + if let Some(msg) = handle.next_protocol_message() { + if received[i] < 2 { + extract_number_from_message(&msg); + received[i] += 1; + } + } + } + if received.iter().all(|&r| r == 2) { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Timeout waiting for initial messages"); + + // Disconnect one node + info!("Disconnecting node"); + drop(handles.pop()); + let remaining_handles: Vec<&NetworkServiceHandle> = handles.iter().collect(); + + // Verify remaining nodes can still communicate + info!("Verifying remaining nodes can communicate"); + for (i, sender) in remaining_handles.iter().enumerate() { + for (j, recipient) in remaining_handles.iter().enumerate() { + if i != j { + sender + .send_protocol_message(create_protocol_message( + PROTOCOL_NAME, + SummationMessage::Verification { sum: numbers[i] }, + sender, + Some(recipient.local_peer_id), + )) + .expect("Failed to send verification"); + } + } + } + + // Wait for verification messages + timeout(TEST_TIMEOUT, async { + let mut verified = vec![false; remaining_handles.len()]; + loop { + for (i, handle) in remaining_handles.iter().enumerate() { + if let Some(msg) = handle.next_protocol_message() { + if !verified[i] { + extract_sum_from_verification(&msg); + verified[i] = true; + } + } + } + if verified.iter().all(|&v| v) { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Timeout waiting for verification after disconnect"); + + info!("Node disconnect test completed successfully"); +} diff --git a/crates/networking/src/tests/handshake.rs b/crates/networking/src/tests/handshake.rs index 499648170..9280499d8 100644 --- a/crates/networking/src/tests/handshake.rs +++ b/crates/networking/src/tests/handshake.rs @@ -210,60 +210,60 @@ async fn test_handshake_reconnection() { info!("Handshake reconnection test completed successfully"); } -// #[tokio::test] -// async fn test_concurrent_connections() { -// init_tracing(); -// info!("Starting concurrent connections test"); - -// let network_name = "test-network"; -// let instance_id = "test-instance"; -// let allowed_keys = HashSet::new(); - -// // Create three nodes to test multiple concurrent handshakes -// let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; -// let mut node2 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; -// let mut node3 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; - -// info!("Starting all nodes simultaneously"); -// // Start all nodes simultaneously -// let (handle1, handle2, handle3) = tokio::join!(node1.start(), node2.start(), node3.start()); -// let handle1 = handle1.expect("Failed to start node1"); -// let handle2 = handle2.expect("Failed to start node2"); -// let handle3 = handle3.expect("Failed to start node3"); - -// // Wait for all handshakes to complete -// info!("Waiting for all handshakes to complete"); -// timeout(TEST_TIMEOUT, async { -// loop { -// let all_verified = handle1.peer_manager.is_peer_verified(&node2.peer_id) -// && handle1.peer_manager.is_peer_verified(&node3.peer_id) -// && handle2.peer_manager.is_peer_verified(&node1.peer_id) -// && handle2.peer_manager.is_peer_verified(&node3.peer_id) -// && handle3.peer_manager.is_peer_verified(&node1.peer_id) -// && handle3.peer_manager.is_peer_verified(&node2.peer_id); - -// if all_verified { -// break; -// } -// tokio::time::sleep(Duration::from_millis(100)).await; -// } -// }) -// .await -// .expect("Concurrent handshakes timed out"); - -// // Verify all peer info is present -// for (handle, peers) in [ -// (&handle1, vec![&node2.peer_id, &node3.peer_id]), -// (&handle2, vec![&node1.peer_id, &node3.peer_id]), -// (&handle3, vec![&node1.peer_id, &node2.peer_id]), -// ] { -// for peer_id in peers { -// assert!( -// handle.peer_info(peer_id).unwrap().identify_info.is_some(), -// "Missing identify info for peer {peer_id:?}" -// ); -// } -// } - -// info!("Concurrent connections test completed successfully"); -// } +#[tokio::test] +async fn test_concurrent_connections() { + init_tracing(); + info!("Starting concurrent connections test"); + + let network_name = "test-network"; + let instance_id = "test-instance"; + let allowed_keys = HashSet::new(); + + // Create three nodes to test multiple concurrent handshakes + let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; + let mut node2 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; + let mut node3 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; + + info!("Starting all nodes simultaneously"); + // Start all nodes simultaneously + let (handle1, handle2, handle3) = tokio::join!(node1.start(), node2.start(), node3.start()); + let handle1 = handle1.expect("Failed to start node1"); + let handle2 = handle2.expect("Failed to start node2"); + let handle3 = handle3.expect("Failed to start node3"); + + // Wait for all handshakes to complete + info!("Waiting for all handshakes to complete"); + timeout(TEST_TIMEOUT, async { + loop { + let all_verified = handle1.peer_manager.is_peer_verified(&node2.peer_id) + && handle1.peer_manager.is_peer_verified(&node3.peer_id) + && handle2.peer_manager.is_peer_verified(&node1.peer_id) + && handle2.peer_manager.is_peer_verified(&node3.peer_id) + && handle3.peer_manager.is_peer_verified(&node1.peer_id) + && handle3.peer_manager.is_peer_verified(&node2.peer_id); + + if all_verified { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Concurrent handshakes timed out"); + + // Verify all peer info is present + for (handle, peers) in [ + (&handle1, vec![&node2.peer_id, &node3.peer_id]), + (&handle2, vec![&node1.peer_id, &node3.peer_id]), + (&handle3, vec![&node1.peer_id, &node2.peer_id]), + ] { + for peer_id in peers { + assert!( + handle.peer_info(peer_id).unwrap().identify_info.is_some(), + "Missing identify info for peer {peer_id:?}" + ); + } + } + + info!("Concurrent connections test completed successfully"); +} diff --git a/crates/networking/src/tests/mod.rs b/crates/networking/src/tests/mod.rs index 5aa035330..a4acc15cc 100644 --- a/crates/networking/src/tests/mod.rs +++ b/crates/networking/src/tests/mod.rs @@ -13,6 +13,7 @@ use std::{collections::HashSet, time::Duration}; use tokio::time::timeout; use tracing::info; +mod blueprint_protocol; mod discovery; mod handshake; @@ -167,6 +168,13 @@ impl TestNode { pub fn get_listen_addr(&self) -> Option { self.listen_addr.clone() } + + /// Update the allowed keys for this node + pub fn update_allowed_keys(&self, allowed_keys: HashSet) { + if let Some(service) = &self.service { + service.peer_manager.update_whitelisted_keys(allowed_keys); + } + } } /// Wait for a condition with timeout From bd83895a20a946a487aecbf4c129df4fee510c13 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 13:34:57 -0700 Subject: [PATCH 34/52] chore: fmt --- crates/crypto/k256/src/lib.rs | 2 +- crates/networking-round-based-extension/src/tests.rs | 8 ++++---- crates/networking/src/blueprint_protocol/behaviour.rs | 4 +++- crates/networking/src/blueprint_protocol/handler.rs | 4 +++- crates/networking/src/lib.rs | 2 +- 5 files changed, 12 insertions(+), 8 deletions(-) diff --git a/crates/crypto/k256/src/lib.rs b/crates/crypto/k256/src/lib.rs index 55c47a8e7..5afe60da6 100644 --- a/crates/crypto/k256/src/lib.rs +++ b/crates/crypto/k256/src/lib.rs @@ -5,7 +5,6 @@ pub mod error; #[cfg(test)] mod tests; -use std::hash::{Hash, Hasher}; use crate::error::{K256Error, Result}; use alloy_signer_local::LocalSigner; use gadget_crypto_core::KeyEncoding; @@ -14,6 +13,7 @@ use gadget_std::string::{String, ToString}; use gadget_std::UniformRand; use k256::ecdsa::signature::SignerMut; use k256::ecdsa::{SigningKey, VerifyingKey}; +use std::hash::{Hash, Hasher}; /// ECDSA key type pub struct K256Ecdsa; diff --git a/crates/networking-round-based-extension/src/tests.rs b/crates/networking-round-based-extension/src/tests.rs index c8cb23149..68c5d8210 100644 --- a/crates/networking-round-based-extension/src/tests.rs +++ b/crates/networking-round-based-extension/src/tests.rs @@ -1,8 +1,8 @@ -use std::collections::HashSet; +use gadget_networking::{KeyType, NetworkConfig, NetworkService}; use libp2p::Multiaddr; use round_based::ProtocolMessage; use serde::{Deserialize, Serialize}; -use gadget_networking::{KeyType, NetworkConfig, NetworkService}; +use std::collections::HashSet; #[derive(Debug, Serialize, Deserialize, Clone, ProtocolMessage)] enum Msg { @@ -51,11 +51,11 @@ fn node() -> NetworkService { enable_mdns: true, enable_kademlia: true, }; - + NetworkService::new(config, HashSet::default()).unwrap() } #[tokio::test] async fn round_based() { let node = node(); -} \ No newline at end of file +} diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index f95ce62d9..4352fe32f 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -344,7 +344,9 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { let mut key_pair = self.instance_key_pair.clone(); let handshake_msg = HandshakeMessage::new(self.local_peer_id); - let Some(signature) = self.sign_handshake(&mut key_pair, &e.peer_id, &handshake_msg) else { + let Some(signature) = + self.sign_handshake(&mut key_pair, &e.peer_id, &handshake_msg) + else { return; }; diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index 83a70d1a6..63eb7ec1b 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -63,7 +63,9 @@ impl BlueprintProtocolBehaviour { let mut key_pair = self.instance_key_pair.clone(); let handshake_msg = HandshakeMessage::new(self.local_peer_id); - let Some(signature) = self.sign_handshake(&mut key_pair, &peer, &handshake_msg) else { + let Some(signature) = + self.sign_handshake(&mut key_pair, &peer, &handshake_msg) + else { return; }; diff --git a/crates/networking/src/lib.rs b/crates/networking/src/lib.rs index 003985975..119dc6aee 100644 --- a/crates/networking/src/lib.rs +++ b/crates/networking/src/lib.rs @@ -15,9 +15,9 @@ mod tests; #[cfg(feature = "round-based-compat")] pub use gadget_networking_round_based_extension as round_based_compat; +pub use gadget_crypto::KeyType; pub use key_types::*; pub use service::{NetworkConfig, NetworkEvent, NetworkService}; -pub use gadget_crypto::KeyType; #[cfg(all( feature = "sp-core-ecdsa", From 94164bb3ada6897580bf2a5a30deaa007399c32c Mon Sep 17 00:00:00 2001 From: Shady Khalifa Date: Wed, 19 Feb 2025 22:37:25 +0200 Subject: [PATCH 35/52] test: add tests to round-based compat --- Cargo.lock | 16 ++ .../Cargo.toml | 42 ++-- .../tests/rand_protocol.rs | 209 ++++++++++++++++++ 3 files changed, 249 insertions(+), 18 deletions(-) create mode 100644 crates/networking-round-based-extension/tests/rand_protocol.rs diff --git a/Cargo.lock b/Cargo.lock index 67d219022..fe7d3230c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7178,13 +7178,18 @@ dependencies = [ "dashmap", "futures", "gadget-networking", + "hex", "libp2p", + "rand 0.8.5", + "rand_dev", "round-based", "serde", "serde_json", + "sha2 0.10.8", "thiserror 2.0.11", "tokio", "tracing", + "tracing-subscriber 0.3.19", ] [[package]] @@ -14734,6 +14739,17 @@ dependencies = [ "zerocopy 0.8.18", ] +[[package]] +name = "rand_dev" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cbee97c27dada05f03db49ffe6516872f6c926e0fd525f9ce0cb3c051adf145c" +dependencies = [ + "hex", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand_xorshift" version = "0.3.0" diff --git a/crates/networking-round-based-extension/Cargo.toml b/crates/networking-round-based-extension/Cargo.toml index 1ee2231ad..929ce5751 100644 --- a/crates/networking-round-based-extension/Cargo.toml +++ b/crates/networking-round-based-extension/Cargo.toml @@ -23,26 +23,32 @@ thiserror = { workspace = true } [dev-dependencies] round-based = { workspace = true, features = ["derive"] } libp2p = { workspace = true, features = [ - "tokio", - "gossipsub", - "mdns", - "noise", - "macros", - "yamux", - "tcp", - "quic", - "request-response", - "cbor", - "identify", - "kad", - "dcutr", - "relay", - "ping", - "dns", - "autonat", - "upnp", + "tokio", + "gossipsub", + "mdns", + "noise", + "macros", + "yamux", + "tcp", + "quic", + "request-response", + "cbor", + "identify", + "kad", + "dcutr", + "relay", + "ping", + "dns", + "autonat", + "upnp", ] } +tokio = { workspace = true, features = ["full"] } +tracing-subscriber = { workspace = true } +sha2 = { workspace = true } +rand_dev = "0.1" +rand = { workspace = true } +hex = { workspace = true } [lints] workspace = true diff --git a/crates/networking-round-based-extension/tests/rand_protocol.rs b/crates/networking-round-based-extension/tests/rand_protocol.rs new file mode 100644 index 000000000..e29190f0f --- /dev/null +++ b/crates/networking-round-based-extension/tests/rand_protocol.rs @@ -0,0 +1,209 @@ +//! Simple protocol in which parties cooperate to generate randomness + +#![no_std] +#![forbid(unused_crate_dependencies, missing_docs)] + +#[cfg(test)] +extern crate std; + +extern crate alloc; + +mod _unused_deps { + // We don't use it directly, but we need to enable `serde` feature + use generic_array as _; +} + +use alloc::{vec, vec::Vec}; + +use serde::{Deserialize, Serialize}; +use sha2::{digest::Output, Digest, Sha256}; + +use round_based::rounds_router::{ + simple_store::{RoundInput, RoundInputError}, + CompleteRoundError, RoundsRouter, +}; +use round_based::{Delivery, Mpc, MpcParty, MsgId, Outgoing, PartyIndex, ProtocolMessage, SinkExt}; + +/// Protocol message +#[derive(Clone, Debug, PartialEq, ProtocolMessage, Serialize, Deserialize)] +pub enum Msg { + /// Round 1 + CommitMsg(CommitMsg), + /// Round 2 + DecommitMsg(DecommitMsg), +} + +/// Message from round 1 +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct CommitMsg { + /// Party commitment + pub commitment: Output, +} + +/// Message from round 2 +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct DecommitMsg { + /// Randomness generated by party + pub randomness: [u8; 32], +} + +/// Carries out the randomness generation protocol +pub async fn protocol_of_random_generation( + party: M, + i: PartyIndex, + n: u16, + mut rng: R, +) -> Result<[u8; 32], Error> +where + M: Mpc, + R: rand_core::RngCore, +{ + let MpcParty { delivery, .. } = party.into_party(); + let (incoming, mut outgoing) = delivery.split(); + + // Define rounds + let mut rounds = RoundsRouter::::builder(); + let round1 = rounds.add_round(RoundInput::::broadcast(i, n)); + let round2 = rounds.add_round(RoundInput::::broadcast(i, n)); + let mut rounds = rounds.listen(incoming); + + // --- The Protocol --- + + // 1. Generate local randomness + let mut local_randomness = [0u8; 32]; + rng.fill_bytes(&mut local_randomness); + + // 2. Commit local randomness (broadcast m=sha256(randomness)) + let commitment = Sha256::digest(local_randomness); + outgoing + .send(Outgoing::broadcast(Msg::CommitMsg(CommitMsg { + commitment, + }))) + .await + .map_err(Error::Round1Send)?; + + // 3. Receive committed randomness from other parties + let commitments = rounds + .complete(round1) + .await + .map_err(Error::Round1Receive)?; + + // 4. Open local randomness + outgoing + .send(Outgoing::broadcast(Msg::DecommitMsg(DecommitMsg { + randomness: local_randomness, + }))) + .await + .map_err(Error::Round2Send)?; + + // 5. Receive opened local randomness from other parties, verify them, and output protocol randomness + let randomness = rounds + .complete(round2) + .await + .map_err(Error::Round2Receive)?; + + let mut guilty_parties = vec![]; + let mut output = local_randomness; + for ((party_i, com_msg_id, commit), (_, decom_msg_id, decommit)) in commitments + .into_iter_indexed() + .zip(randomness.into_iter_indexed()) + { + let commitment_expected = Sha256::digest(decommit.randomness); + if commit.commitment != commitment_expected { + guilty_parties.push(Blame { + guilty_party: party_i, + commitment_msg: com_msg_id, + decommitment_msg: decom_msg_id, + }); + continue; + } + + output + .iter_mut() + .zip(decommit.randomness) + .for_each(|(x, r)| *x ^= r); + } + + if !guilty_parties.is_empty() { + Err(Error::PartiesOpenedRandomnessDoesntMatchCommitment { guilty_parties }) + } else { + Ok(output) + } +} + +/// Protocol error +#[derive(Debug, thiserror::Error)] +pub enum Error { + /// Couldn't send a message in the first round + #[error("send a message at round 1")] + Round1Send(#[source] SendErr), + /// Couldn't receive a message in the first round + #[error("receive messages at round 1")] + Round1Receive(#[source] CompleteRoundError), + /// Couldn't send a message in the second round + #[error("send a message at round 2")] + Round2Send(#[source] SendErr), + /// Couldn't receive a message in the second round + #[error("receive messages at round 2")] + Round2Receive(#[source] CompleteRoundError), + + /// Some of the parties cheated + #[error("malicious parties: {guilty_parties:?}")] + PartiesOpenedRandomnessDoesntMatchCommitment { + /// List of cheated parties + guilty_parties: Vec, + }, +} + +/// Blames a party in cheating during the protocol +#[derive(Debug)] +pub struct Blame { + /// Index of the cheated party + pub guilty_party: PartyIndex, + /// ID of the message that party sent in the first round + pub commitment_msg: MsgId, + /// ID of the message that party sent in the second round + pub decommitment_msg: MsgId, +} + +#[cfg(test)] +mod tests { + use rand::Rng; + use sha2::{Digest, Sha256}; + + use super::protocol_of_random_generation; + + #[test] + fn simulation() { + let mut rng = rand_dev::DevRng::new(); + + let n: u16 = 5; + + let randomness = round_based::sim::run_with_setup( + core::iter::repeat_with(|| rng.fork()).take(n.into()), + |i, party, rng| protocol_of_random_generation(party, i, n, rng), + ) + .unwrap() + .expect_ok() + .expect_eq(); + + std::println!("Output randomness: {}", hex::encode(randomness)); + } + + #[tokio::test] + async fn simulation_async() { + let mut rng = rand_dev::DevRng::new(); + + let n: u16 = 5; + + let randomness = round_based::sim::async_env::run_with_setup( + core::iter::repeat_with(|| rng.fork()).take(n.into()), + |i, party, rng| protocol_of_random_generation(party, i, n, rng), + ) + .await + .expect_ok() + .expect_eq(); + + std::println!("Output randomness: {}", hex::encode(randomness)); + } +} From 2fb99a3335065ce4ff0970217c1d5729394a18c9 Mon Sep 17 00:00:00 2001 From: Shady Khalifa Date: Wed, 19 Feb 2025 22:38:52 +0200 Subject: [PATCH 36/52] fix: remove cranelift from cargo.toml --- Cargo.toml | 5 ----- 1 file changed, 5 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index b1b9d247c..adc070995 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,8 +1,3 @@ -cargo-features = ["codegen-backend"] - -[profile.dev] -codegen-backend = "cranelift" - [workspace] resolver = "2" members = ["cli", "blueprints/*", "crates/*"] From 67457dfb2df3943b99d8b23a292cb9d8d1d18f3c Mon Sep 17 00:00:00 2001 From: Shady Khalifa Date: Wed, 19 Feb 2025 23:09:08 +0200 Subject: [PATCH 37/52] test: add tests for round based networking --- Cargo.lock | 6 + .../Cargo.toml | 6 +- .../tests/common/mod.rs | 242 ++++++++++++++++++ .../tests/rand_protocol.rs | 92 ++++++- 4 files changed, 338 insertions(+), 8 deletions(-) create mode 100644 crates/networking-round-based-extension/tests/common/mod.rs diff --git a/Cargo.lock b/Cargo.lock index fe7d3230c..48d10266c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7177,7 +7177,10 @@ dependencies = [ "crossbeam-channel", "dashmap", "futures", + "gadget-crypto", + "gadget-crypto-core", "gadget-networking", + "generic-array", "hex", "libp2p", "rand 0.8.5", @@ -7447,6 +7450,7 @@ version = "0.14.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ + "serde", "typenum", "version_check", "zeroize", @@ -15211,6 +15215,8 @@ dependencies = [ "phantom-type", "round-based-derive", "thiserror 2.0.11", + "tokio", + "tokio-stream", "tracing", ] diff --git a/crates/networking-round-based-extension/Cargo.toml b/crates/networking-round-based-extension/Cargo.toml index 929ce5751..491addc40 100644 --- a/crates/networking-round-based-extension/Cargo.toml +++ b/crates/networking-round-based-extension/Cargo.toml @@ -21,7 +21,7 @@ crossbeam-channel = { workspace = true } thiserror = { workspace = true } [dev-dependencies] -round-based = { workspace = true, features = ["derive"] } +round-based = { workspace = true, features = ["derive", "sim", "sim-async"] } libp2p = { workspace = true, features = [ "tokio", "gossipsub", @@ -42,6 +42,8 @@ libp2p = { workspace = true, features = [ "autonat", "upnp", ] } +gadget-crypto = { workspace = true, features = ["sp-core"] } +gadget-crypto-core = { workspace = true, features = ["tangle"] } tokio = { workspace = true, features = ["full"] } tracing-subscriber = { workspace = true } @@ -49,6 +51,8 @@ sha2 = { workspace = true } rand_dev = "0.1" rand = { workspace = true } hex = { workspace = true } +# We don't use it directly, but we need to enable `serde` feature +generic-array = { version = "0.14", features = ["serde"] } [lints] workspace = true diff --git a/crates/networking-round-based-extension/tests/common/mod.rs b/crates/networking-round-based-extension/tests/common/mod.rs new file mode 100644 index 000000000..0b4e2be25 --- /dev/null +++ b/crates/networking-round-based-extension/tests/common/mod.rs @@ -0,0 +1,242 @@ +use gadget_crypto::KeyType; +use gadget_networking::service::NetworkMessage; +use gadget_networking::{ + key_types::{Curve, InstanceMsgKeyPair, InstanceMsgPublicKey}, + service_handle::NetworkServiceHandle, + NetworkConfig, NetworkService, +}; +use libp2p::{ + identity::{self, Keypair}, + Multiaddr, PeerId, +}; +use std::string::ToString; +use std::{collections::HashSet, time::Duration}; +use tokio::time::timeout; +use tracing::info; + +pub fn init_tracing() { + let _ = tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_target(true) + .with_thread_ids(true) + .with_file(true) + .with_line_number(true) + .try_init(); +} + +/// Test node configuration for network tests +pub struct TestNode { + pub service: Option, + pub peer_id: PeerId, + pub listen_addr: Option, + pub instance_key_pair: InstanceMsgKeyPair, + pub local_key: Keypair, +} + +impl TestNode { + /// Create a new test node with auto-generated keys + pub async fn new( + network_name: &str, + instance_id: &str, + allowed_keys: HashSet, + bootstrap_peers: Vec, + ) -> Self { + Self::new_with_keys( + network_name, + instance_id, + allowed_keys, + bootstrap_peers, + None, + None, + ) + .await + } + + /// Create a new test node with specified keys + pub async fn new_with_keys( + network_name: &str, + instance_id: &str, + allowed_keys: HashSet, + bootstrap_peers: Vec, + instance_key_pair: Option, + local_key: Option, + ) -> Self { + let local_key = local_key.unwrap_or_else(|| identity::Keypair::generate_ed25519()); + let peer_id = local_key.public().to_peer_id(); + + // Bind to all interfaces instead of just localhost + let listen_addr: Multiaddr = "/ip4/0.0.0.0/tcp/0".parse().unwrap(); + info!("Creating test node {peer_id} with TCP address: {listen_addr}"); + + let instance_key_pair = + instance_key_pair.unwrap_or_else(|| Curve::generate_with_seed(None).unwrap()); + + let config = NetworkConfig { + network_name: network_name.to_string(), + instance_id: instance_id.to_string(), + instance_key_pair: instance_key_pair.clone(), + local_key: local_key.clone(), + listen_addr: listen_addr.clone(), + target_peer_count: 10, + bootstrap_peers, + enable_mdns: true, + enable_kademlia: true, + }; + + let service = + NetworkService::new(config, allowed_keys).expect("Failed to create network service"); + + Self { + service: Some(service), + peer_id, + listen_addr: None, + instance_key_pair, + local_key, + } + } + + /// Start the node and wait for it to be fully initialized + pub async fn start(&mut self) -> Result { + // Take ownership of the service + let service = self.service.take().ok_or("Service already started")?; + let handle = service.start(); + + // Wait for the node to be fully initialized + let timeout_duration = Duration::from_secs(10); // Increased timeout + match timeout(timeout_duration, async { + // First wait for the listening address + while self.listen_addr.is_none() { + if let Some(addr) = handle.get_listen_addr() { + info!("Node {} listening on {}", self.peer_id, addr); + self.listen_addr = Some(addr.clone()); + + // Extract port from multiaddr + let addr_str = addr.to_string(); + let port = addr_str.split("/").nth(4).unwrap_or("0").to_string(); + + // Try localhost first + let localhost_addr = format!("127.0.0.1:{}", port); + match tokio::net::TcpStream::connect(&localhost_addr).await { + Ok(_) => { + info!("Successfully verified localhost port for {}", self.peer_id); + break; + } + Err(e) => { + info!("Localhost port not ready for {}: {}", self.peer_id, e); + // Try external IP + let external_addr = format!("10.0.1.142:{}", port); + match tokio::net::TcpStream::connect(&external_addr).await { + Ok(_) => { + info!( + "Successfully verified external port for {}", + self.peer_id + ); + break; + } + Err(e) => { + info!("External port not ready for {}: {}", self.peer_id, e); + tokio::time::sleep(Duration::from_millis(100)).await; + continue; + } + } + } + } + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + + // Give the node a moment to initialize protocols + tokio::time::sleep(Duration::from_millis(500)).await; + + Ok::<(), &'static str>(()) + }) + .await + { + Ok(Ok(_)) => { + info!("Node {} fully initialized", self.peer_id); + Ok(handle) + } + Ok(Err(e)) => Err(e), + Err(_) => Err("Timeout waiting for node to initialize"), + } + } + + /// Get the actual listening address + pub fn get_listen_addr(&self) -> Option { + self.listen_addr.clone() + } +} + +/// Wait for a condition with timeout +pub async fn wait_for_condition(timeout: Duration, mut condition: F) -> Result<(), &'static str> +where + F: FnMut() -> bool, +{ + let start = std::time::Instant::now(); + while !condition() { + if start.elapsed() > timeout { + return Err("Timeout waiting for condition"); + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + Ok(()) +} + +/// Wait for peers to discover each other +pub async fn wait_for_peer_discovery( + handles: &[&NetworkServiceHandle], + timeout: Duration, +) -> Result<(), &'static str> { + info!("Waiting for peer discovery..."); + + wait_for_condition(timeout, || { + for (i, handle1) in handles.iter().enumerate() { + for (j, handle2) in handles.iter().enumerate() { + if i != j + && !handle1 + .peers() + .iter() + .any(|id| *id == handle2.local_peer_id) + { + return false; + } + } + } + true + }) + .await +} + +/// Wait for peer info to be updated +pub async fn wait_for_peer_info( + handle1: &NetworkServiceHandle, + handle2: &NetworkServiceHandle, + timeout: Duration, +) { + info!("Waiting for identify info..."); + + match tokio::time::timeout(timeout, async { + loop { + let peer_info1 = handle1.peer_info(&handle2.local_peer_id); + let peer_info2 = handle2.peer_info(&handle1.local_peer_id); + + if let Some(peer_info) = peer_info1 { + if peer_info.identify_info.is_some() { + // Also verify reverse direction + if let Some(peer_info) = peer_info2 { + if peer_info.identify_info.is_some() { + info!("Identify info populated in both directions"); + break; + } + } + } + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + { + Ok(_) => info!("Peer info updated successfully in both directions"), + Err(_) => panic!("Peer info update timed out"), + } +} diff --git a/crates/networking-round-based-extension/tests/rand_protocol.rs b/crates/networking-round-based-extension/tests/rand_protocol.rs index e29190f0f..1d1bf4c01 100644 --- a/crates/networking-round-based-extension/tests/rand_protocol.rs +++ b/crates/networking-round-based-extension/tests/rand_protocol.rs @@ -1,13 +1,9 @@ //! Simple protocol in which parties cooperate to generate randomness -#![no_std] -#![forbid(unused_crate_dependencies, missing_docs)] - -#[cfg(test)] -extern crate std; - extern crate alloc; +mod common; + mod _unused_deps { // We don't use it directly, but we need to enable `serde` feature use generic_array as _; @@ -56,7 +52,7 @@ pub async fn protocol_of_random_generation( ) -> Result<[u8; 32], Error> where M: Mpc, - R: rand_core::RngCore, + R: rand::RngCore, { let MpcParty { delivery, .. } = party.into_party(); let (incoming, mut outgoing) = delivery.split(); @@ -168,8 +164,15 @@ pub struct Blame { #[cfg(test)] mod tests { + use std::collections::{HashMap, HashSet}; + + use super::common::*; + use gadget_networking::{Curve, KeyType}; + use gadget_networking_round_based_extension::RoundBasedNetworkAdapter; use rand::Rng; + use round_based::MpcParty; use sha2::{Digest, Sha256}; + use tracing::{debug, info}; use super::protocol_of_random_generation; @@ -206,4 +209,79 @@ mod tests { std::println!("Output randomness: {}", hex::encode(randomness)); } + + #[tokio::test] + async fn p2p_networking() { + init_tracing(); + let network_name = "rand-test-network"; + let instance_id = "rand-test-instance"; + + // Generate node2's key pair first + let instance_key_pair2 = Curve::generate_with_seed(None).unwrap(); + let mut allowed_keys1 = HashSet::new(); + allowed_keys1.insert(instance_key_pair2.public()); + + // Create node1 with node2's key whitelisted + let mut node1 = TestNode::new(network_name, instance_id, allowed_keys1, vec![]).await; + + // Create node2 with node1's key whitelisted and pre-generated key + let mut allowed_keys2 = HashSet::new(); + allowed_keys2.insert(node1.instance_key_pair.public()); + let mut node2 = TestNode::new_with_keys( + network_name, + instance_id, + allowed_keys2, + vec![], + Some(instance_key_pair2), + None, + ) + .await; + + info!("Starting nodes"); + // Start both nodes - this should trigger automatic handshake + let handle1 = node1.start().await.expect("Failed to start node1"); + let handle2 = node2.start().await.expect("Failed to start node2"); + + let parties = HashMap::from_iter([ + (0, node1.instance_key_pair.public()), + (1, node2.instance_key_pair.public()), + ]); + + let node1_network = RoundBasedNetworkAdapter::new(handle1, 0, parties.clone(), instance_id); + let node2_network = RoundBasedNetworkAdapter::new(handle2, 1, parties, instance_id); + + let mut tasks = vec![]; + tasks.push(tokio::spawn(async move { + let mut rng = rand_dev::DevRng::new(); + let mpc_party = MpcParty::connected(node1_network); + let randomness = protocol_of_random_generation(mpc_party, 0, 2, &mut rng) + .await + .expect("Failed to generate randomness"); + debug!("Node1 generated randomness: {:?}", randomness); + randomness + })); + + tasks.push(tokio::spawn(async move { + let mut rng = rand_dev::DevRng::new(); + let mpc_party = MpcParty::connected(node2_network); + let randomness = protocol_of_random_generation(mpc_party, 1, 2, &mut rng) + .await + .expect("Failed to generate randomness"); + debug!("Node2 generated randomness: {:?}", randomness); + randomness + })); + + let results = futures::future::join_all(tasks).await; + + for result in results { + match result { + Ok(randomness) => { + debug!("Randomness result: {:?}", randomness); + } + Err(e) => { + panic!("Error in randomness generation: {:?}", e); + } + } + } + } } From ada706400a4b5904d27ce2287850accb887d4c96 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 14:42:54 -0700 Subject: [PATCH 38/52] feat: start implementing inbound gossip handling --- .../src/blueprint_protocol/behaviour.rs | 47 +++++- .../src/blueprint_protocol/handler.rs | 25 ---- .../src/tests/blueprint_protocol.rs | 115 +++++++++++---- crates/networking/src/tests/handshake.rs | 136 ------------------ 4 files changed, 131 insertions(+), 192 deletions(-) diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index 4352fe32f..804bf76ce 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -5,6 +5,7 @@ use crate::{ types::ProtocolMessage, Curve, InstanceMsgKeyPair, InstanceMsgPublicKey, InstanceSignedMsgSignature, }; +use bincode; use crossbeam_channel::Sender; use dashmap::DashMap; use gadget_crypto::KeyType; @@ -102,11 +103,18 @@ impl BlueprintProtocolBehaviour { let gossipsub_config = gossipsub::ConfigBuilder::default() .heartbeat_interval(Duration::from_secs(1)) .validation_mode(gossipsub::ValidationMode::Strict) + .mesh_n_low(2) + .mesh_n(4) + .mesh_n_high(8) + .gossip_lazy(3) + .history_length(10) + .history_gossip(3) + .flood_publish(true) .build() .expect("Valid gossipsub config"); let gossipsub = gossipsub::Behaviour::new( - MessageAuthenticity::Signed(local_key.clone()), + gossipsub::MessageAuthenticity::Signed(local_key.clone()), gossipsub_config, ) .expect("Valid gossipsub behaviour"); @@ -257,6 +265,42 @@ impl BlueprintProtocolBehaviour { } } } + + pub fn handle_gossipsub_event(&mut self, event: gossipsub::Event) { + match event { + gossipsub::Event::Message { + propagation_source, + message_id: _, + message, + } => { + // Only accept gossip from verified peers + if !self.peer_manager.is_peer_verified(&propagation_source) { + warn!(%propagation_source, "Received gossip from unverified peer"); + return; + } + + debug!(%propagation_source, "Received gossip message"); + + // Deserialize the protocol message + if let Ok(protocol_message) = bincode::deserialize::(&message.data) + { + debug!(%propagation_source, %protocol_message, "Forwarding gossip message to protocol handler"); + if let Err(e) = self.protocol_message_sender.send(protocol_message) { + warn!(%propagation_source, "Failed to forward gossip message: {}", e); + } + } else { + warn!(%propagation_source, "Failed to deserialize gossip message"); + } + } + gossipsub::Event::Subscribed { peer_id, topic } => { + debug!(%peer_id, %topic, "Peer subscribed to topic"); + } + gossipsub::Event::Unsubscribed { peer_id, topic } => { + debug!(%peer_id, %topic, "Peer unsubscribed from topic"); + } + _ => {} + } + } } impl NetworkBehaviour for BlueprintProtocolBehaviour { @@ -449,7 +493,6 @@ impl NetworkBehaviour for BlueprintProtocolBehaviour { _ => {} } } - Poll::Pending } } diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index 63eb7ec1b..fbe7d86bf 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -268,29 +268,4 @@ impl BlueprintProtocolBehaviour { // Add to verified peers self.peer_manager.verify_peer(peer); } - - pub fn handle_gossipsub_event(&mut self, event: gossipsub::Event) { - match event { - gossipsub::Event::Message { - propagation_source, - message_id, - message, - } => { - // Only accept gossip from verified peers - if !self.peer_manager.is_peer_verified(&propagation_source) { - warn!(%propagation_source, "Received gossip from unverified peer"); - return; - } - - debug!(%propagation_source, "Received gossip message"); - } - gossipsub::Event::Subscribed { peer_id, topic } => { - debug!(%peer_id, %topic, "Peer subscribed to topic"); - } - gossipsub::Event::Unsubscribed { peer_id, topic } => { - debug!(%peer_id, %topic, "Peer unsubscribed from topic"); - } - _ => {} - } - } } diff --git a/crates/networking/src/tests/blueprint_protocol.rs b/crates/networking/src/tests/blueprint_protocol.rs index 01880f3fc..50769da12 100644 --- a/crates/networking/src/tests/blueprint_protocol.rs +++ b/crates/networking/src/tests/blueprint_protocol.rs @@ -60,19 +60,30 @@ fn create_protocol_message( } // Helper to wait for handshake completion between multiple nodes -async fn wait_for_all_handshakes(handles: &[&NetworkServiceHandle]) { +async fn wait_for_all_handshakes(handles: &[&mut NetworkServiceHandle]) { + info!("Starting handshake wait for {} nodes", handles.len()); timeout(TEST_TIMEOUT, async { loop { let mut all_verified = true; - for (i, h1) in handles.iter().enumerate() { - for (j, h2) in handles.iter().enumerate() { - if i != j && !h1.peer_manager.is_peer_verified(&h2.local_peer_id) { - all_verified = false; - break; + for (i, handle1) in handles.iter().enumerate() { + for (j, handle2) in handles.iter().enumerate() { + if i != j { + let verified = handle1 + .peer_manager + .is_peer_verified(&handle2.local_peer_id); + if !verified { + info!("Node {} -> Node {}: handshake not verified yet", i, j); + all_verified = false; + break; + } } } + if !all_verified { + break; + } } if all_verified { + info!("All handshakes completed successfully"); break; } tokio::time::sleep(Duration::from_millis(100)).await; @@ -305,28 +316,45 @@ async fn test_summation_protocol_multi_node() { info!("Starting multi-node summation protocol test"); // Create 3 nodes with whitelisted keys + info!("Creating whitelisted nodes"); let mut nodes = create_whitelisted_nodes(3).await; + info!("Created {} nodes successfully", nodes.len()); // Start all nodes + info!("Starting all nodes"); let mut handles = Vec::new(); - for node in nodes.iter_mut() { + for (i, node) in nodes.iter_mut().enumerate() { + info!("Starting node {}", i); handles.push(node.start().await.expect("Failed to start node")); + info!("Node {} started successfully", i); } // Convert handles to mutable references - let handles: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); + info!("Converting handles to mutable references"); + let mut handles: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); + let handles_len = handles.len(); + info!("Converted {} handles", handles_len); // Wait for all handshakes to complete - info!("Waiting for handshake completion"); + info!( + "Waiting for handshake completion between {} nodes", + handles_len + ); wait_for_all_handshakes(&handles).await; + info!("All handshakes completed successfully"); // Generate test numbers let numbers = vec![42, 58, 100]; let expected_sum: u64 = numbers.iter().sum(); + info!( + "Generated test numbers: {:?}, expected sum: {}", + numbers, expected_sum + ); info!("Sending numbers via gossip"); // Each node broadcasts its number for (i, handle) in handles.iter().enumerate() { + info!("Node {} broadcasting number {}", i, numbers[i]); handle .send_protocol_message(create_protocol_message( PROTOCOL_NAME, @@ -335,25 +363,40 @@ async fn test_summation_protocol_multi_node() { None, )) .expect("Failed to send number"); + info!("Node {} successfully broadcast its number", i); + // Add a small delay between broadcasts to avoid message collisions + tokio::time::sleep(Duration::from_millis(100)).await; } info!("Waiting for messages to be processed"); // Wait for all nodes to receive all numbers - let mut sums = vec![0; handles.len()]; - let mut received = vec![0; handles.len()]; + let mut sums = vec![0; handles_len]; + let mut received = vec![0; handles_len]; timeout(TEST_TIMEOUT, async { loop { - for (i, handle) in handles.iter().enumerate() { + for (i, handle) in handles.iter_mut().enumerate() { if let Some(msg) = handle.next_protocol_message() { - if received[i] < handles.len() - 1 { - sums[i] += extract_number_from_message(&msg); + if received[i] < handles_len - 1 { + let num = extract_number_from_message(&msg); + sums[i] += num; received[i] += 1; + info!( + "Node {} received number {}, total sum: {}, received count: {}", + i, num, sums[i], received[i] + ); } } } - if received.iter().all(|&r| r == handles.len() - 1) { + let all_received = received.iter().all(|&r| r == handles_len - 1); + info!( + "Current received counts: {:?}, target count: {}", + received, + handles_len - 1 + ); + if all_received { + info!("All nodes have received all numbers"); break; } tokio::time::sleep(Duration::from_millis(100)).await; @@ -363,10 +406,15 @@ async fn test_summation_protocol_multi_node() { .expect("Timeout waiting for summation completion"); info!("Verifying sums via P2P messages"); + info!("Final sums: {:?}", sums); // Each node verifies with every other node for (i, sender) in handles.iter().enumerate() { for (j, recipient) in handles.iter().enumerate() { if i != j { + info!( + "Node {} sending verification sum {} to node {}", + i, sums[i], j + ); sender .send_protocol_message(create_protocol_message( PROTOCOL_NAME, @@ -382,18 +430,27 @@ async fn test_summation_protocol_multi_node() { info!("Waiting for verification messages"); // Wait for all verifications timeout(TEST_TIMEOUT, async { - let mut verified = vec![0; handles.len()]; + let mut verified = vec![0; handles_len]; loop { - for (i, handle) in handles.iter().enumerate() { + for (i, handle) in handles.iter_mut().enumerate() { if let Some(msg) = handle.next_protocol_message() { - if verified[i] < handles.len() - 1 { - assert_eq!(extract_sum_from_verification(&msg), expected_sum); + if verified[i] < handles_len - 1 { + let sum = extract_sum_from_verification(&msg); + info!( + "Node {} received verification sum {}, expected {}", + i, sum, expected_sum + ); + assert_eq!(sum, expected_sum); verified[i] += 1; + info!("Node {} verification count: {}", i, verified[i]); } } } - if verified.iter().all(|&v| v == handles.len() - 1) { + let all_verified = verified.iter().all(|&v| v == handles_len - 1); + info!("Current verification counts: {:?}", verified); + if all_verified { + info!("All nodes have verified all sums"); break; } tokio::time::sleep(Duration::from_millis(100)).await; @@ -420,7 +477,7 @@ async fn test_summation_protocol_late_join() { } // Convert handles to mutable references - let mut handles_refs: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); + let handles_refs: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); // Wait for initial handshakes info!("Waiting for initial handshake completion"); @@ -428,7 +485,7 @@ async fn test_summation_protocol_late_join() { // Initial nodes send their numbers let numbers = vec![42, 58, 100]; - let expected_sum: u64 = numbers.iter().sum(); + let _expected_sum: u64 = numbers.iter().sum(); info!("Initial nodes sending numbers"); for (i, handle) in handles.iter().enumerate() { @@ -446,7 +503,7 @@ async fn test_summation_protocol_late_join() { timeout(TEST_TIMEOUT, async { let mut received = vec![false; 2]; loop { - for (i, handle) in handles.iter().enumerate() { + for (i, handle) in handles.iter_mut().enumerate() { if let Some(msg) = handle.next_protocol_message() { if !received[i] { assert_eq!(extract_number_from_message(&msg), numbers[1 - i]); @@ -466,7 +523,7 @@ async fn test_summation_protocol_late_join() { // Start the late joining node info!("Starting late joining node"); handles.push(nodes[2].start().await.expect("Failed to start late node")); - let all_handles: Vec<&NetworkServiceHandle> = handles.iter().collect(); + let all_handles: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); // Wait for the new node to complete handshakes wait_for_all_handshakes(&all_handles).await; @@ -486,7 +543,7 @@ async fn test_summation_protocol_late_join() { timeout(TEST_TIMEOUT, async { let mut verified = vec![false; handles.len()]; loop { - for (i, handle) in handles.iter().enumerate() { + for (i, handle) in handles.iter_mut().enumerate() { if let Some(msg) = handle.next_protocol_message() { let num = extract_number_from_message(&msg); if !verified[i] && (num == numbers[2] || numbers.contains(&num)) { @@ -521,7 +578,7 @@ async fn test_summation_protocol_node_disconnect() { } // Convert handles to mutable references - let mut handles_refs: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); + let handles_refs: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); // Wait for all handshakes wait_for_all_handshakes(&handles_refs).await; @@ -543,7 +600,7 @@ async fn test_summation_protocol_node_disconnect() { timeout(TEST_TIMEOUT, async { let mut received = vec![0; handles.len()]; loop { - for (i, handle) in handles.iter().enumerate() { + for (i, handle) in handles.iter_mut().enumerate() { if let Some(msg) = handle.next_protocol_message() { if received[i] < 2 { extract_number_from_message(&msg); @@ -563,7 +620,7 @@ async fn test_summation_protocol_node_disconnect() { // Disconnect one node info!("Disconnecting node"); drop(handles.pop()); - let remaining_handles: Vec<&NetworkServiceHandle> = handles.iter().collect(); + let mut remaining_handles: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); // Verify remaining nodes can still communicate info!("Verifying remaining nodes can communicate"); @@ -586,7 +643,7 @@ async fn test_summation_protocol_node_disconnect() { timeout(TEST_TIMEOUT, async { let mut verified = vec![false; remaining_handles.len()]; loop { - for (i, handle) in remaining_handles.iter().enumerate() { + for (i, handle) in remaining_handles.iter_mut().enumerate() { if let Some(msg) = handle.next_protocol_message() { if !verified[i] { extract_sum_from_verification(&msg); diff --git a/crates/networking/src/tests/handshake.rs b/crates/networking/src/tests/handshake.rs index 9280499d8..dbe01d714 100644 --- a/crates/networking/src/tests/handshake.rs +++ b/crates/networking/src/tests/handshake.rs @@ -131,139 +131,3 @@ async fn test_handshake_with_invalid_peer() { info!("Invalid peer handshake test completed successfully"); } - -#[tokio::test] -async fn test_handshake_reconnection() { - init_tracing(); - info!("Starting handshake reconnection test"); - - let network_name = "test-network"; - let instance_id = "test-instance"; - - // Generate node2's key pair first - let instance_key_pair2 = Curve::generate_with_seed(None).unwrap(); - let mut allowed_keys1 = HashSet::new(); - allowed_keys1.insert(instance_key_pair2.public()); - - // Create node1 with node2's key whitelisted - let mut node1 = TestNode::new(network_name, instance_id, allowed_keys1, vec![]).await; - - // Create node2 with node1's key whitelisted and pre-generated key - let mut allowed_keys2 = HashSet::new(); - allowed_keys2.insert(node1.instance_key_pair.public()); - let mut node2 = TestNode::new_with_keys( - network_name, - instance_id, - allowed_keys2, - vec![], - Some(instance_key_pair2), - None, - ) - .await; - - info!("Starting initial connection"); - // Start both nodes and wait for initial handshake - let handle1 = node1.start().await.expect("Failed to start node1"); - let handle2 = node2.start().await.expect("Failed to start node2"); - - // Wait for initial handshake - timeout(TEST_TIMEOUT, async { - loop { - if handle1.peer_manager.is_peer_verified(&node2.peer_id) - && handle2.peer_manager.is_peer_verified(&node1.peer_id) - { - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Initial handshake timed out"); - - info!("Disconnecting node2"); - // Drop node2's handle to simulate disconnect - drop(handle2); - tokio::time::sleep(Duration::from_millis(100)).await; - - // Verify node1 sees node2 as disconnected - assert!(!handle1.peer_manager.is_peer_verified(&node2.peer_id)); - - info!("Reconnecting node2"); - // Restart node2 - let handle2 = node2.start().await.expect("Failed to restart node2"); - - // Wait for automatic reconnection and handshake - info!("Waiting for automatic reconnection handshake"); - timeout(TEST_TIMEOUT, async { - loop { - if handle1.peer_manager.is_peer_verified(&node2.peer_id) - && handle2.peer_manager.is_peer_verified(&node1.peer_id) - { - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Reconnection handshake timed out"); - - info!("Handshake reconnection test completed successfully"); -} - -#[tokio::test] -async fn test_concurrent_connections() { - init_tracing(); - info!("Starting concurrent connections test"); - - let network_name = "test-network"; - let instance_id = "test-instance"; - let allowed_keys = HashSet::new(); - - // Create three nodes to test multiple concurrent handshakes - let mut node1 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; - let mut node2 = TestNode::new(network_name, instance_id, allowed_keys.clone(), vec![]).await; - let mut node3 = TestNode::new(network_name, instance_id, allowed_keys, vec![]).await; - - info!("Starting all nodes simultaneously"); - // Start all nodes simultaneously - let (handle1, handle2, handle3) = tokio::join!(node1.start(), node2.start(), node3.start()); - let handle1 = handle1.expect("Failed to start node1"); - let handle2 = handle2.expect("Failed to start node2"); - let handle3 = handle3.expect("Failed to start node3"); - - // Wait for all handshakes to complete - info!("Waiting for all handshakes to complete"); - timeout(TEST_TIMEOUT, async { - loop { - let all_verified = handle1.peer_manager.is_peer_verified(&node2.peer_id) - && handle1.peer_manager.is_peer_verified(&node3.peer_id) - && handle2.peer_manager.is_peer_verified(&node1.peer_id) - && handle2.peer_manager.is_peer_verified(&node3.peer_id) - && handle3.peer_manager.is_peer_verified(&node1.peer_id) - && handle3.peer_manager.is_peer_verified(&node2.peer_id); - - if all_verified { - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Concurrent handshakes timed out"); - - // Verify all peer info is present - for (handle, peers) in [ - (&handle1, vec![&node2.peer_id, &node3.peer_id]), - (&handle2, vec![&node1.peer_id, &node3.peer_id]), - (&handle3, vec![&node1.peer_id, &node2.peer_id]), - ] { - for peer_id in peers { - assert!( - handle.peer_info(peer_id).unwrap().identify_info.is_some(), - "Missing identify info for peer {peer_id:?}" - ); - } - } - - info!("Concurrent connections test completed successfully"); -} From 625e3828ccddae30f08eae953dd9b9995828690b Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 14:50:39 -0700 Subject: [PATCH 39/52] chore: add network message handler for instance p2p requests and gossip messages --- crates/networking/src/service.rs | 54 +++++++++++++++++++++----------- 1 file changed, 36 insertions(+), 18 deletions(-) diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index c71bdb7c0..dfda3f9e4 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -72,30 +72,15 @@ pub enum NetworkEvent { /// Network message types #[derive(Debug)] pub enum NetworkMessage { - Dial(Multiaddr), InstanceRequest { peer: PeerId, request: InstanceMessageRequest, }, - InstanceResponse { - peer: PeerId, - response: InstanceMessageResponse, - }, GossipMessage { source: PeerId, topic: String, message: Vec, }, - HandshakeRequest { - peer: PeerId, - public_key: InstanceMsgPublicKey, - signature: InstanceSignedMsgSignature, - }, - HandshakeResponse { - peer: PeerId, - public_key: InstanceMsgPublicKey, - signature: InstanceSignedMsgSignature, - }, } /// Configuration for the network service @@ -525,10 +510,43 @@ async fn handle_ping_event( /// Handle a network message async fn handle_network_message( - _swarm: &mut Swarm, - _msg: NetworkMessage, - _peer_manager: &Arc, + swarm: &mut Swarm, + msg: NetworkMessage, + peer_manager: &Arc, event_sender: &Sender, ) -> Result<(), Error> { + match msg { + NetworkMessage::InstanceRequest { peer, request } => { + // Only send requests to verified peers + if !peer_manager.is_peer_verified(&peer) { + warn!(%peer, "Attempted to send request to unverified peer"); + return Ok(()); + } + + debug!(%peer, ?request, "Sending instance request"); + swarm + .behaviour_mut() + .blueprint_protocol + .send_request(&peer, request.clone()); + event_sender.send(NetworkEvent::InstanceRequestOutbound { peer, request })?; + } + NetworkMessage::GossipMessage { + source, + topic, + message, + } => { + debug!(%source, %topic, "Publishing gossip message"); + if let Err(e) = swarm + .behaviour_mut() + .blueprint_protocol + .publish(&topic, message.clone()) + { + warn!(%source, %topic, "Failed to publish gossip message: {:?}", e); + return Ok(()); + } + event_sender.send(NetworkEvent::GossipSent { topic, message })?; + } + } + Ok(()) } From 1d6e0c49032f7ab4f6fb5ba8ef7f41b62090a55a Mon Sep 17 00:00:00 2001 From: Serial <69764315+Serial-ATA@users.noreply.github.com> Date: Wed, 19 Feb 2025 17:01:59 -0500 Subject: [PATCH 40/52] chore: change some log levels --- .../tests/rand_protocol.rs | 9 --------- .../networking/src/blueprint_protocol/behaviour.rs | 14 +++++++------- crates/networking/src/service.rs | 8 ++++---- crates/networking/src/service_handle.rs | 6 +++--- 4 files changed, 14 insertions(+), 23 deletions(-) diff --git a/crates/networking-round-based-extension/tests/rand_protocol.rs b/crates/networking-round-based-extension/tests/rand_protocol.rs index 1d1bf4c01..978a5fd07 100644 --- a/crates/networking-round-based-extension/tests/rand_protocol.rs +++ b/crates/networking-round-based-extension/tests/rand_protocol.rs @@ -1,16 +1,7 @@ //! Simple protocol in which parties cooperate to generate randomness -extern crate alloc; - mod common; -mod _unused_deps { - // We don't use it directly, but we need to enable `serde` feature - use generic_array as _; -} - -use alloc::{vec, vec::Vec}; - use serde::{Deserialize, Serialize}; use sha2::{digest::Output, Digest, Sha256}; diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index 804bf76ce..8d92605f4 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -282,14 +282,14 @@ impl BlueprintProtocolBehaviour { debug!(%propagation_source, "Received gossip message"); // Deserialize the protocol message - if let Ok(protocol_message) = bincode::deserialize::(&message.data) - { - debug!(%propagation_source, %protocol_message, "Forwarding gossip message to protocol handler"); - if let Err(e) = self.protocol_message_sender.send(protocol_message) { - warn!(%propagation_source, "Failed to forward gossip message: {}", e); - } - } else { + let Ok(protocol_message) = bincode::deserialize::(&message.data) else { warn!(%propagation_source, "Failed to deserialize gossip message"); + return; + }; + + debug!(%propagation_source, %protocol_message, "Forwarding gossip message to protocol handler"); + if let Err(e) = self.protocol_message_sender.send(protocol_message) { + warn!(%propagation_source, "Failed to forward gossip message: {e}"); } } gossipsub::Event::Subscribed { peer_id, topic } => { diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index dfda3f9e4..333a67c3b 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -374,7 +374,7 @@ async fn handle_discovery_event( .map(std::string::ToString::to_string) .collect(); - debug!(%peer_id, ?protocols, "Supported protocols"); + trace!(%peer_id, ?protocols, "Supported protocols"); let blueprint_protocol_name = &swarm.behaviour().blueprint_protocol.blueprint_protocol_name; @@ -392,15 +392,15 @@ async fn handle_discovery_event( // Update identify info peer_info.identify_info = Some(info.clone()); - debug!(%peer_id, listen_addrs=?info.listen_addrs, "Adding identify addresses"); + trace!(%peer_id, listen_addrs=?info.listen_addrs, "Adding identify addresses"); // Add all addresses from identify info for addr in &info.listen_addrs { peer_info.addresses.insert(addr.clone()); } - debug!(%peer_id, "Updating peer info with identify information"); + trace!(%peer_id, "Updating peer info with identify information"); peer_manager.update_peer(*peer_id, peer_info); - info!(%peer_id, "Successfully processed identify information"); + debug!(%peer_id, "Successfully processed identify information"); } DerivedDiscoveryBehaviourEvent::Identify(_) => { // Ignore other identify events diff --git a/crates/networking/src/service_handle.rs b/crates/networking/src/service_handle.rs index a307b6168..c014db3b3 100644 --- a/crates/networking/src/service_handle.rs +++ b/crates/networking/src/service_handle.rs @@ -8,7 +8,7 @@ use crossbeam_channel::{self, Receiver, Sender}; use libp2p::{Multiaddr, PeerId}; use std::sync::Arc; use tokio::task::JoinHandle; -use tracing::info; +use tracing::{debug, info}; /// Handle for sending outgoing messages to the network #[derive(Clone)] @@ -122,7 +122,7 @@ impl NetworkServiceHandle { peer: peer_id, request: instance_message_request, })?; - info!("Sent outbound p2p `NetworkMessage` to {:?}", peer_id); + debug!("Sent outbound p2p `NetworkMessage` to {:?}", peer_id); } } } else { @@ -132,7 +132,7 @@ impl NetworkServiceHandle { message: raw_payload, }; self.send_network_message(gossip_message)?; - info!("Sent outbound gossip `NetworkMessage`"); + debug!("Sent outbound gossip `NetworkMessage`"); } Ok(()) From 61c8c40481428bc5fb489b93341b86edaa81a58d Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 15:07:41 -0700 Subject: [PATCH 41/52] chore: add gossip test, add subscribe to run, wip --- crates/networking/src/service.rs | 6 + .../src/tests/blueprint_protocol.rs | 117 +--------- crates/networking/src/tests/gossip.rs | 209 ++++++++++++++++++ crates/networking/src/tests/mod.rs | 98 ++++++++ 4 files changed, 321 insertions(+), 109 deletions(-) create mode 100644 crates/networking/src/tests/gossip.rs diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index 333a67c3b..fc8abc949 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -250,6 +250,12 @@ impl NetworkService { } } + self.swarm + .behaviour_mut() + .blueprint_protocol + .subscribe(&self.network_name) + .unwrap(); + loop { tokio::select! { swarm_event = self.swarm.select_next_some() => { diff --git a/crates/networking/src/tests/blueprint_protocol.rs b/crates/networking/src/tests/blueprint_protocol.rs index 50769da12..0c163f9bc 100644 --- a/crates/networking/src/tests/blueprint_protocol.rs +++ b/crates/networking/src/tests/blueprint_protocol.rs @@ -10,7 +10,9 @@ use crate::{ key_types::{Curve, InstanceMsgKeyPair, InstanceMsgPublicKey}, service::NetworkMessage, service_handle::NetworkServiceHandle, - tests::TestNode, + tests::{ + create_whitelisted_nodes, wait_for_all_handshakes, wait_for_handshake_completion, TestNode, + }, types::{MessageRouting, ParticipantId, ParticipantInfo, ProtocolMessage}, }; @@ -24,16 +26,6 @@ enum SummationMessage { Verification { sum: u64 }, } -// Helper to create a whitelisted test node -async fn create_node_with_keys( - network: &str, - instance: &str, - allowed_keys: HashSet, - key_pair: Option, -) -> TestNode { - TestNode::new_with_keys(network, instance, allowed_keys, vec![], key_pair, None).await -} - // Helper to create a protocol message fn create_protocol_message( protocol: &str, @@ -59,63 +51,6 @@ fn create_protocol_message( } } -// Helper to wait for handshake completion between multiple nodes -async fn wait_for_all_handshakes(handles: &[&mut NetworkServiceHandle]) { - info!("Starting handshake wait for {} nodes", handles.len()); - timeout(TEST_TIMEOUT, async { - loop { - let mut all_verified = true; - for (i, handle1) in handles.iter().enumerate() { - for (j, handle2) in handles.iter().enumerate() { - if i != j { - let verified = handle1 - .peer_manager - .is_peer_verified(&handle2.local_peer_id); - if !verified { - info!("Node {} -> Node {}: handshake not verified yet", i, j); - all_verified = false; - break; - } - } - } - if !all_verified { - break; - } - } - if all_verified { - info!("All handshakes completed successfully"); - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Handshake verification timed out"); -} - -// Helper to wait for handshake completion between two nodes -async fn wait_for_handshake_completion( - handle1: &NetworkServiceHandle, - handle2: &NetworkServiceHandle, -) { - timeout(TEST_TIMEOUT, async { - loop { - if handle1 - .peer_manager - .is_peer_verified(&handle2.local_peer_id) - && handle2 - .peer_manager - .is_peer_verified(&handle1.local_peer_id) - { - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Handshake verification timed out"); -} - // Helper to extract number from message fn extract_number_from_message(msg: &ProtocolMessage) -> u64 { match bincode::deserialize::(&msg.payload).expect("Failed to deserialize") { @@ -132,42 +67,6 @@ fn extract_sum_from_verification(msg: &ProtocolMessage) -> u64 { } } -// Helper to create a set of nodes with whitelisted keys -async fn create_whitelisted_nodes(count: usize) -> Vec { - let mut nodes = Vec::with_capacity(count); - let mut key_pairs = Vec::with_capacity(count); - let mut allowed_keys = vec![HashSet::new(); count]; - - // Generate all key pairs first - for _ in 0..count { - key_pairs.push(Curve::generate_with_seed(None).unwrap()); - } - - // Create allowed keys sets - for i in 0..count { - for j in 0..count { - if i != j { - allowed_keys[i].insert(key_pairs[j].public()); - } - } - } - - // Create nodes with whitelisted keys - for i in 0..count { - nodes.push( - create_node_with_keys( - "test-net", - "sum-test", - allowed_keys[i].clone(), - Some(key_pairs[i].clone()), - ) - .await, - ); - } - - nodes -} - #[tokio::test] async fn test_summation_protocol_basic() { super::init_tracing(); @@ -197,7 +96,7 @@ async fn test_summation_protocol_basic() { let mut handle2 = node2.start().await.expect("Failed to start node2"); info!("Waiting for handshake completion"); - wait_for_handshake_completion(&handle1, &handle2).await; + wait_for_handshake_completion(&handle1, &handle2, TEST_TIMEOUT).await; // Generate test numbers let num1 = 42; @@ -340,7 +239,7 @@ async fn test_summation_protocol_multi_node() { "Waiting for handshake completion between {} nodes", handles_len ); - wait_for_all_handshakes(&handles).await; + wait_for_all_handshakes(&handles, TEST_TIMEOUT).await; info!("All handshakes completed successfully"); // Generate test numbers @@ -481,7 +380,7 @@ async fn test_summation_protocol_late_join() { // Wait for initial handshakes info!("Waiting for initial handshake completion"); - wait_for_all_handshakes(&handles_refs).await; + wait_for_all_handshakes(&handles_refs, TEST_TIMEOUT).await; // Initial nodes send their numbers let numbers = vec![42, 58, 100]; @@ -526,7 +425,7 @@ async fn test_summation_protocol_late_join() { let all_handles: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); // Wait for the new node to complete handshakes - wait_for_all_handshakes(&all_handles).await; + wait_for_all_handshakes(&all_handles, TEST_TIMEOUT).await; // Late node sends its number and receives history info!("Late node sending number and receiving history"); @@ -581,7 +480,7 @@ async fn test_summation_protocol_node_disconnect() { let handles_refs: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); // Wait for all handshakes - wait_for_all_handshakes(&handles_refs).await; + wait_for_all_handshakes(&handles_refs, TEST_TIMEOUT).await; // Send initial numbers let numbers = vec![42, 58, 100]; diff --git a/crates/networking/src/tests/gossip.rs b/crates/networking/src/tests/gossip.rs new file mode 100644 index 000000000..efa6e1d4c --- /dev/null +++ b/crates/networking/src/tests/gossip.rs @@ -0,0 +1,209 @@ +use super::{init_tracing, wait_for_handshake_completion, TestNode}; +use crate::{ + key_types::{Curve, InstanceMsgPublicKey}, + service_handle::NetworkServiceHandle, + tests::{create_whitelisted_nodes, wait_for_all_handshakes}, + types::{MessageRouting, ParticipantId, ParticipantInfo, ProtocolMessage}, +}; +use gadget_crypto::KeyType; +use std::{collections::HashSet, time::Duration}; +use tokio::time::timeout; +use tracing::info; + +const TEST_TIMEOUT: Duration = Duration::from_secs(10); +const PROTOCOL_NAME: &str = "test-gossip"; + +#[tokio::test] +async fn test_gossip_between_verified_peers() { + init_tracing(); + info!("Starting gossip test between verified peers"); + + // Create nodes with whitelisted keys + let instance_key_pair2 = Curve::generate_with_seed(None).unwrap(); + let mut allowed_keys1 = HashSet::new(); + allowed_keys1.insert(instance_key_pair2.public()); + + let mut node1 = TestNode::new("test-net", "gossip-test", allowed_keys1, vec![]).await; + + let mut allowed_keys2 = HashSet::new(); + allowed_keys2.insert(node1.instance_key_pair.public()); + let mut node2 = TestNode::new_with_keys( + "test-net", + "gossip-test", + allowed_keys2, + vec![], + Some(instance_key_pair2), + None, + ) + .await; + + info!("Starting nodes"); + let mut handle1 = node1.start().await.expect("Failed to start node1"); + let mut handle2 = node2.start().await.expect("Failed to start node2"); + + info!("Waiting for handshake completion"); + wait_for_handshake_completion(&handle1, &handle2, TEST_TIMEOUT).await; + + // Create test message + let test_payload = b"Hello, gossip network!".to_vec(); + let message = ProtocolMessage { + protocol: PROTOCOL_NAME.to_string(), + routing: MessageRouting { + message_id: 1, + round_id: 0, + sender: ParticipantInfo { + id: ParticipantId(1), + public_key: Some(node1.instance_key_pair.public()), + }, + recipient: None, // No specific recipient for gossip + }, + payload: test_payload.clone(), + }; + + info!("Sending gossip message from node1"); + handle1 + .send_protocol_message(message) + .expect("Failed to send gossip message"); + + info!("Waiting for node2 to receive the message"); + // Wait for node2 to receive the message + let received_message = timeout(TEST_TIMEOUT, async { + loop { + if let Some(msg) = handle2.next_protocol_message() { + if msg.protocol == PROTOCOL_NAME { + return msg; + } + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Timeout waiting for gossip message"); + + // Verify message contents + assert_eq!(received_message.payload, test_payload); + assert_eq!(received_message.protocol, PROTOCOL_NAME); + assert_eq!(received_message.routing.message_id, 1); + assert_eq!(received_message.routing.round_id, 0); + assert_eq!( + received_message.routing.sender.public_key, + Some(node1.instance_key_pair.public()) + ); + assert!(received_message.routing.recipient.is_none()); + + info!("Gossip test completed successfully"); +} + +#[tokio::test] +async fn test_multi_node_gossip() { + init_tracing(); + info!("Starting multi-node gossip test"); + + // Create three nodes with all keys whitelisted + let mut nodes = create_whitelisted_nodes(3).await; + + info!("Starting all nodes"); + let mut handles: Vec<_> = Vec::new(); + for node in nodes.iter_mut() { + handles.push(node.start().await.expect("Failed to start node")); + } + + info!("Waiting for all handshakes to complete"); + let handles_refs: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); + wait_for_all_handshakes(&handles_refs, TEST_TIMEOUT).await; + + // Create test message + let test_payload = b"Multi-node gossip test".to_vec(); + let message = ProtocolMessage { + protocol: PROTOCOL_NAME.to_string(), + routing: MessageRouting { + message_id: 1, + round_id: 0, + sender: ParticipantInfo { + id: ParticipantId(0), + public_key: Some(nodes[0].instance_key_pair.public()), + }, + recipient: None, + }, + payload: test_payload.clone(), + }; + + info!("Sending gossip message from node 0"); + handles[0] + .send_protocol_message(message) + .expect("Failed to send gossip message"); + + info!("Waiting for all nodes to receive the message"); + // Wait for all other nodes to receive the message + timeout(TEST_TIMEOUT, async { + for (i, handle) in handles.iter_mut().enumerate().skip(1) { + let received = loop { + if let Some(msg) = handle.next_protocol_message() { + if msg.protocol == PROTOCOL_NAME { + break msg; + } + } + tokio::time::sleep(Duration::from_millis(100)).await; + }; + + assert_eq!( + received.payload, test_payload, + "Node {} received wrong payload", + i + ); + assert_eq!(received.protocol, PROTOCOL_NAME); + assert_eq!( + received.routing.sender.public_key, + Some(nodes[0].instance_key_pair.public()) + ); + info!("Node {} received the gossip message correctly", i); + } + }) + .await + .expect("Timeout waiting for gossip messages"); + + info!("Multi-node gossip test completed successfully"); +} + +#[tokio::test] +async fn test_unverified_peer_gossip() { + init_tracing(); + info!("Starting unverified peer gossip test"); + + // Create two nodes with no whitelisted keys + let mut node1 = TestNode::new("test-net", "gossip-test", HashSet::new(), vec![]).await; + let mut node2 = TestNode::new("test-net", "gossip-test", HashSet::new(), vec![]).await; + + info!("Starting nodes"); + let mut handle1 = node1.start().await.expect("Failed to start node1"); + let mut handle2 = node2.start().await.expect("Failed to start node2"); + + // Create test message + let test_payload = b"This message should not be received".to_vec(); + let message = ProtocolMessage { + protocol: PROTOCOL_NAME.to_string(), + routing: MessageRouting { + message_id: 1, + round_id: 0, + sender: ParticipantInfo { + id: ParticipantId(1), + public_key: Some(node1.instance_key_pair.public()), + }, + recipient: None, + }, + payload: test_payload, + }; + + info!("Attempting to send gossip message from unverified node"); + handle1 + .send_protocol_message(message) + .expect("Failed to send gossip message"); + + // Wait a bit to ensure message is not received + tokio::time::sleep(Duration::from_secs(2)).await; + + // Verify node2 did not receive the message + assert!(handle2.next_protocol_message().is_none()); + + info!("Unverified peer gossip test completed successfully"); +} diff --git a/crates/networking/src/tests/mod.rs b/crates/networking/src/tests/mod.rs index a4acc15cc..ddb306077 100644 --- a/crates/networking/src/tests/mod.rs +++ b/crates/networking/src/tests/mod.rs @@ -15,6 +15,7 @@ use tracing::info; mod blueprint_protocol; mod discovery; +mod gossip; mod handshake; fn init_tracing() { @@ -250,3 +251,100 @@ pub async fn wait_for_peer_info( Err(_) => panic!("Peer info update timed out"), } } + +// Helper to wait for handshake completion between multiple nodes +async fn wait_for_all_handshakes(handles: &[&mut NetworkServiceHandle], timeout_length: Duration) { + info!("Starting handshake wait for {} nodes", handles.len()); + timeout(timeout_length, async { + loop { + let mut all_verified = true; + for (i, handle1) in handles.iter().enumerate() { + for (j, handle2) in handles.iter().enumerate() { + if i != j { + let verified = handle1 + .peer_manager + .is_peer_verified(&handle2.local_peer_id); + if !verified { + info!("Node {} -> Node {}: handshake not verified yet", i, j); + all_verified = false; + break; + } + } + } + if !all_verified { + break; + } + } + if all_verified { + info!("All handshakes completed successfully"); + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Handshake verification timed out"); +} + +// Helper to wait for handshake completion between two nodes +async fn wait_for_handshake_completion( + handle1: &NetworkServiceHandle, + handle2: &NetworkServiceHandle, + timeout_length: Duration, +) { + timeout(timeout_length, async { + loop { + if handle1 + .peer_manager + .is_peer_verified(&handle2.local_peer_id) + && handle2 + .peer_manager + .is_peer_verified(&handle1.local_peer_id) + { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }) + .await + .expect("Handshake verification timed out"); +} + +// Helper to create a whitelisted test node +async fn create_node_with_keys( + network: &str, + instance: &str, + allowed_keys: HashSet, + key_pair: Option, +) -> TestNode { + TestNode::new_with_keys(network, instance, allowed_keys, vec![], key_pair, None).await +} + +// Helper to create a set of nodes with whitelisted keys +async fn create_whitelisted_nodes(count: usize) -> Vec { + let mut nodes = Vec::with_capacity(count); + let mut key_pairs = Vec::with_capacity(count); + let mut allowed_keys = HashSet::new(); + + // Generate all key pairs first + for _ in 0..count { + let key_pair = Curve::generate_with_seed(None).unwrap(); + key_pairs.push(key_pair.clone()); + allowed_keys.insert(key_pair.public()); + } + + // Create nodes with whitelisted keys + for i in 0..count { + nodes.push( + create_node_with_keys( + "test-net", + "sum-test", + allowed_keys.clone(), + Some(key_pairs[i].clone()), + ) + .await, + ); + } + + nodes +} From a7de9810da2825afee91a326f88a2342a8b570b3 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 15:07:49 -0700 Subject: [PATCH 42/52] chore: fmt --- crates/networking/src/blueprint_protocol/behaviour.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index 8d92605f4..4b033946f 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -282,7 +282,8 @@ impl BlueprintProtocolBehaviour { debug!(%propagation_source, "Received gossip message"); // Deserialize the protocol message - let Ok(protocol_message) = bincode::deserialize::(&message.data) else { + let Ok(protocol_message) = bincode::deserialize::(&message.data) + else { warn!(%propagation_source, "Failed to deserialize gossip message"); return; }; From 1d7dbb7079e03844cb8e9ab0869224b1bc6042ef Mon Sep 17 00:00:00 2001 From: Serial <69764315+Serial-ATA@users.noreply.github.com> Date: Wed, 19 Feb 2025 17:51:09 -0500 Subject: [PATCH 43/52] fix: set protocol name properly --- .../src/lib.rs | 2 +- crates/networking/src/service.rs | 9 +- crates/networking/src/service_handle.rs | 70 +- .../src/tests/blueprint_protocol.rs | 1127 +++++++++-------- crates/networking/src/tests/gossip.rs | 66 +- 5 files changed, 641 insertions(+), 633 deletions(-) diff --git a/crates/networking-round-based-extension/src/lib.rs b/crates/networking-round-based-extension/src/lib.rs index 2eaddc5ad..1a1c1ff5c 100644 --- a/crates/networking-round-based-extension/src/lib.rs +++ b/crates/networking-round-based-extension/src/lib.rs @@ -169,7 +169,7 @@ where }; this.handle - .send_protocol_message(protocol_message) + .send(protocol_message) .map_err(|e| NetworkError::Send(e)) } diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index fc8abc949..ca9cc9cc9 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -182,6 +182,8 @@ impl NetworkService { .unwrap() .build(); + swarm.behaviour_mut().blueprint_protocol.subscribe(&blueprint_protocol_name)?; + // Start listening swarm.listen_on(listen_addr)?; let bootstrap_peers = bootstrap_peers.into_iter().collect(); @@ -213,6 +215,7 @@ impl NetworkService { // Create handle with new interface let handle = NetworkServiceHandle::new( local_peer_id, + self.swarm.behaviour().blueprint_protocol.blueprint_protocol_name.clone(), self.peer_manager.clone(), network_sender, protocol_message_receiver, @@ -250,12 +253,6 @@ impl NetworkService { } } - self.swarm - .behaviour_mut() - .blueprint_protocol - .subscribe(&self.network_name) - .unwrap(); - loop { tokio::select! { swarm_event = self.swarm.select_next_some() => { diff --git a/crates/networking/src/service_handle.rs b/crates/networking/src/service_handle.rs index c014db3b3..35b6660e1 100644 --- a/crates/networking/src/service_handle.rs +++ b/crates/networking/src/service_handle.rs @@ -7,8 +7,10 @@ use crate::{ use crossbeam_channel::{self, Receiver, Sender}; use libp2p::{Multiaddr, PeerId}; use std::sync::Arc; +use libp2p::gossipsub::Sha256Topic; use tokio::task::JoinHandle; use tracing::{debug, info}; +use crate::types::{MessageRouting, ParticipantInfo}; /// Handle for sending outgoing messages to the network #[derive(Clone)] @@ -54,6 +56,7 @@ impl NetworkReceiver { /// Combined handle for the network service pub struct NetworkServiceHandle { pub local_peer_id: PeerId, + pub blueprint_protocol_name: Arc, pub sender: NetworkSender, pub receiver: NetworkReceiver, pub peer_manager: Arc, @@ -63,6 +66,7 @@ impl Clone for NetworkServiceHandle { fn clone(&self) -> Self { Self { local_peer_id: self.local_peer_id, + blueprint_protocol_name: self.blueprint_protocol_name.clone(), sender: self.sender.clone(), receiver: NetworkReceiver::new(self.receiver.protocol_message_receiver.clone()), peer_manager: self.peer_manager.clone(), @@ -74,12 +78,14 @@ impl NetworkServiceHandle { #[must_use] pub fn new( local_peer_id: PeerId, + blueprint_protocol_name: String, peer_manager: Arc, network_message_sender: Sender, protocol_message_receiver: Receiver, ) -> Self { Self { local_peer_id, + blueprint_protocol_name: Arc::from(blueprint_protocol_name), sender: NetworkSender::new(network_message_sender), receiver: NetworkReceiver::new(protocol_message_receiver), peer_manager, @@ -107,32 +113,46 @@ impl NetworkServiceHandle { } #[must_use] - pub fn send_protocol_message(&self, message: ProtocolMessage) -> Result<(), String> { - let raw_payload = bincode::serialize(&message).map_err(|err| err.to_string())?; - if message.routing.recipient.is_some() { - let instance_message_request = InstanceMessageRequest::Protocol { - protocol: message.protocol, - payload: raw_payload, - metadata: None, - }; - let recipient = message.routing.recipient.unwrap(); - if let Some(public_key) = recipient.public_key { - if let Some(peer_id) = self.peer_manager.get_peer_id_from_public_key(&public_key) { - self.send_network_message(NetworkMessage::InstanceRequest { - peer: peer_id, - request: instance_message_request, - })?; - debug!("Sent outbound p2p `NetworkMessage` to {:?}", peer_id); - } + pub fn send(&self, routing: MessageRouting, message: impl Into>) -> Result<(), String> { + let protocol = Sha256Topic::new(&*self.blueprint_protocol_name).to_string(); + let protocol_message = ProtocolMessage { + protocol: protocol.clone(), + routing, + payload: message.into(), + }; + + let raw_payload = bincode::serialize(&protocol_message).map_err(|err| err.to_string())?; + match protocol_message.routing.recipient { + Some(recipient) => { + let instance_message_request = InstanceMessageRequest::Protocol { + protocol, + payload: raw_payload, + metadata: None, + }; + + let Some(public_key) = recipient.public_key else { + return Ok(()); + }; + + let Some(peer_id) = self.peer_manager.get_peer_id_from_public_key(&public_key) else { + return Ok(()); + }; + + self.send_network_message(NetworkMessage::InstanceRequest { + peer: peer_id, + request: instance_message_request, + })?; + debug!("Sent outbound p2p `NetworkMessage` to {:?}", peer_id); + }, + None => { + let gossip_message = NetworkMessage::GossipMessage { + source: self.local_peer_id, + topic: protocol, + message: raw_payload, + }; + self.send_network_message(gossip_message)?; + debug!("Sent outbound gossip `NetworkMessage`"); } - } else { - let gossip_message = NetworkMessage::GossipMessage { - source: self.local_peer_id, - topic: message.protocol, - message: raw_payload, - }; - self.send_network_message(gossip_message)?; - debug!("Sent outbound gossip `NetworkMessage`"); } Ok(()) diff --git a/crates/networking/src/tests/blueprint_protocol.rs b/crates/networking/src/tests/blueprint_protocol.rs index 0c163f9bc..040ee064d 100644 --- a/crates/networking/src/tests/blueprint_protocol.rs +++ b/crates/networking/src/tests/blueprint_protocol.rs @@ -1,563 +1,564 @@ -use gadget_crypto::KeyType; -use libp2p::PeerId; -use serde::{Deserialize, Serialize}; -use std::{collections::HashSet, time::Duration}; -use tokio::time::timeout; -use tracing::info; - -use crate::{ - blueprint_protocol::{InstanceMessageRequest, InstanceMessageResponse}, - key_types::{Curve, InstanceMsgKeyPair, InstanceMsgPublicKey}, - service::NetworkMessage, - service_handle::NetworkServiceHandle, - tests::{ - create_whitelisted_nodes, wait_for_all_handshakes, wait_for_handshake_completion, TestNode, - }, - types::{MessageRouting, ParticipantId, ParticipantInfo, ProtocolMessage}, -}; - -const TEST_TIMEOUT: Duration = Duration::from_secs(10); -const PROTOCOL_NAME: &str = "summation/1.0.0"; - -// Protocol message types -#[derive(Debug, Clone, Serialize, Deserialize)] -enum SummationMessage { - Number(u64), - Verification { sum: u64 }, -} - -// Helper to create a protocol message -fn create_protocol_message( - protocol: &str, - msg: SummationMessage, - sender: &NetworkServiceHandle, - recipient: Option, -) -> ProtocolMessage { - ProtocolMessage { - protocol: protocol.to_string(), - routing: MessageRouting { - message_id: 0, - round_id: 0, - sender: ParticipantInfo { - id: ParticipantId(0), - public_key: None, - }, - recipient: recipient.map(|peer_id| ParticipantInfo { - id: ParticipantId(0), - public_key: None, - }), - }, - payload: bincode::serialize(&msg).expect("Failed to serialize message"), - } -} - -// Helper to extract number from message -fn extract_number_from_message(msg: &ProtocolMessage) -> u64 { - match bincode::deserialize::(&msg.payload).expect("Failed to deserialize") { - SummationMessage::Number(n) => n, - _ => panic!("Expected number message"), - } -} - -// Helper to extract sum from verification message -fn extract_sum_from_verification(msg: &ProtocolMessage) -> u64 { - match bincode::deserialize::(&msg.payload).expect("Failed to deserialize") { - SummationMessage::Verification { sum } => sum, - _ => panic!("Expected verification message"), - } -} - -#[tokio::test] -async fn test_summation_protocol_basic() { - super::init_tracing(); - info!("Starting summation protocol test"); - - // Create nodes with whitelisted keys - let instance_key_pair2 = Curve::generate_with_seed(None).unwrap(); - let mut allowed_keys1 = HashSet::new(); - allowed_keys1.insert(instance_key_pair2.public()); - - let mut node1 = TestNode::new("test-net", "sum-test", allowed_keys1, vec![]).await; - - let mut allowed_keys2 = HashSet::new(); - allowed_keys2.insert(node1.instance_key_pair.public()); - let mut node2 = TestNode::new_with_keys( - "test-net", - "sum-test", - allowed_keys2, - vec![], - Some(instance_key_pair2), - None, - ) - .await; - - info!("Starting nodes"); - let mut handle1 = node1.start().await.expect("Failed to start node1"); - let mut handle2 = node2.start().await.expect("Failed to start node2"); - - info!("Waiting for handshake completion"); - wait_for_handshake_completion(&handle1, &handle2, TEST_TIMEOUT).await; - - // Generate test numbers - let num1 = 42; - let num2 = 58; - let expected_sum = num1 + num2; - - info!("Sending numbers via gossip"); - // Send numbers via gossip - handle1 - .send_protocol_message(create_protocol_message( - PROTOCOL_NAME, - SummationMessage::Number(num1), - &handle1, - None, - )) - .expect("Failed to send number from node1"); - - handle2 - .send_protocol_message(create_protocol_message( - PROTOCOL_NAME, - SummationMessage::Number(num2), - &handle2, - None, - )) - .expect("Failed to send number from node2"); - - info!("Waiting for messages to be processed"); - // Wait for messages and compute sums - let mut sum1 = 0; - let mut sum2 = 0; - let mut node1_received = false; - let mut node2_received = false; - - timeout(TEST_TIMEOUT, async { - loop { - // Process incoming messages - if let Some(msg) = handle1.next_protocol_message() { - if !node1_received { - sum1 += extract_number_from_message(&msg); - node1_received = true; - } - } - if let Some(msg) = handle2.next_protocol_message() { - if !node2_received { - sum2 += extract_number_from_message(&msg); - node2_received = true; - } - } - - // Check if both nodes have received messages - if node1_received && node2_received { - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Timeout waiting for summation completion"); - - info!("Verifying sums via P2P messages"); - // Verify sums via P2P messages - handle1 - .send_protocol_message(create_protocol_message( - PROTOCOL_NAME, - SummationMessage::Verification { sum: sum1 }, - &handle1, - Some(node2.peer_id), - )) - .expect("Failed to send verification from node1"); - - handle2 - .send_protocol_message(create_protocol_message( - PROTOCOL_NAME, - SummationMessage::Verification { sum: sum2 }, - &handle2, - Some(node1.peer_id), - )) - .expect("Failed to send verification from node2"); - - info!("Waiting for verification messages"); - // Wait for verification messages - timeout(TEST_TIMEOUT, async { - let mut node1_verified = false; - let mut node2_verified = false; - - loop { - // Process verification messages - if let Some(msg) = handle1.next_protocol_message() { - if !node1_verified { - assert_eq!(extract_sum_from_verification(&msg), expected_sum); - node1_verified = true; - } - } - if let Some(msg) = handle2.next_protocol_message() { - if !node2_verified { - assert_eq!(extract_sum_from_verification(&msg), expected_sum); - node2_verified = true; - } - } - - if node1_verified && node2_verified { - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Timeout waiting for verification completion"); - - info!("Summation protocol test completed successfully"); -} - -#[tokio::test] -async fn test_summation_protocol_multi_node() { - super::init_tracing(); - info!("Starting multi-node summation protocol test"); - - // Create 3 nodes with whitelisted keys - info!("Creating whitelisted nodes"); - let mut nodes = create_whitelisted_nodes(3).await; - info!("Created {} nodes successfully", nodes.len()); - - // Start all nodes - info!("Starting all nodes"); - let mut handles = Vec::new(); - for (i, node) in nodes.iter_mut().enumerate() { - info!("Starting node {}", i); - handles.push(node.start().await.expect("Failed to start node")); - info!("Node {} started successfully", i); - } - - // Convert handles to mutable references - info!("Converting handles to mutable references"); - let mut handles: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); - let handles_len = handles.len(); - info!("Converted {} handles", handles_len); - - // Wait for all handshakes to complete - info!( - "Waiting for handshake completion between {} nodes", - handles_len - ); - wait_for_all_handshakes(&handles, TEST_TIMEOUT).await; - info!("All handshakes completed successfully"); - - // Generate test numbers - let numbers = vec![42, 58, 100]; - let expected_sum: u64 = numbers.iter().sum(); - info!( - "Generated test numbers: {:?}, expected sum: {}", - numbers, expected_sum - ); - - info!("Sending numbers via gossip"); - // Each node broadcasts its number - for (i, handle) in handles.iter().enumerate() { - info!("Node {} broadcasting number {}", i, numbers[i]); - handle - .send_protocol_message(create_protocol_message( - PROTOCOL_NAME, - SummationMessage::Number(numbers[i]), - handle, - None, - )) - .expect("Failed to send number"); - info!("Node {} successfully broadcast its number", i); - // Add a small delay between broadcasts to avoid message collisions - tokio::time::sleep(Duration::from_millis(100)).await; - } - - info!("Waiting for messages to be processed"); - // Wait for all nodes to receive all numbers - let mut sums = vec![0; handles_len]; - let mut received = vec![0; handles_len]; - - timeout(TEST_TIMEOUT, async { - loop { - for (i, handle) in handles.iter_mut().enumerate() { - if let Some(msg) = handle.next_protocol_message() { - if received[i] < handles_len - 1 { - let num = extract_number_from_message(&msg); - sums[i] += num; - received[i] += 1; - info!( - "Node {} received number {}, total sum: {}, received count: {}", - i, num, sums[i], received[i] - ); - } - } - } - - let all_received = received.iter().all(|&r| r == handles_len - 1); - info!( - "Current received counts: {:?}, target count: {}", - received, - handles_len - 1 - ); - if all_received { - info!("All nodes have received all numbers"); - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Timeout waiting for summation completion"); - - info!("Verifying sums via P2P messages"); - info!("Final sums: {:?}", sums); - // Each node verifies with every other node - for (i, sender) in handles.iter().enumerate() { - for (j, recipient) in handles.iter().enumerate() { - if i != j { - info!( - "Node {} sending verification sum {} to node {}", - i, sums[i], j - ); - sender - .send_protocol_message(create_protocol_message( - PROTOCOL_NAME, - SummationMessage::Verification { sum: sums[i] }, - sender, - Some(recipient.local_peer_id), - )) - .expect("Failed to send verification"); - } - } - } - - info!("Waiting for verification messages"); - // Wait for all verifications - timeout(TEST_TIMEOUT, async { - let mut verified = vec![0; handles_len]; - loop { - for (i, handle) in handles.iter_mut().enumerate() { - if let Some(msg) = handle.next_protocol_message() { - if verified[i] < handles_len - 1 { - let sum = extract_sum_from_verification(&msg); - info!( - "Node {} received verification sum {}, expected {}", - i, sum, expected_sum - ); - assert_eq!(sum, expected_sum); - verified[i] += 1; - info!("Node {} verification count: {}", i, verified[i]); - } - } - } - - let all_verified = verified.iter().all(|&v| v == handles_len - 1); - info!("Current verification counts: {:?}", verified); - if all_verified { - info!("All nodes have verified all sums"); - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Timeout waiting for verification completion"); - - info!("Multi-node summation protocol test completed successfully"); -} - -#[tokio::test] -async fn test_summation_protocol_late_join() { - super::init_tracing(); - info!("Starting late join summation protocol test"); - - // Create 3 nodes but only start 2 initially - let mut nodes = create_whitelisted_nodes(3).await; - - // Start first two nodes - let mut handles = Vec::new(); - for node in nodes[..2].iter_mut() { - handles.push(node.start().await.expect("Failed to start node")); - } - - // Convert handles to mutable references - let handles_refs: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); - - // Wait for initial handshakes - info!("Waiting for initial handshake completion"); - wait_for_all_handshakes(&handles_refs, TEST_TIMEOUT).await; - - // Initial nodes send their numbers - let numbers = vec![42, 58, 100]; - let _expected_sum: u64 = numbers.iter().sum(); - - info!("Initial nodes sending numbers"); - for (i, handle) in handles.iter().enumerate() { - handle - .send_protocol_message(create_protocol_message( - PROTOCOL_NAME, - SummationMessage::Number(numbers[i]), - handle, - None, - )) - .expect("Failed to send number"); - } - - // Wait for initial nodes to process messages - timeout(TEST_TIMEOUT, async { - let mut received = vec![false; 2]; - loop { - for (i, handle) in handles.iter_mut().enumerate() { - if let Some(msg) = handle.next_protocol_message() { - if !received[i] { - assert_eq!(extract_number_from_message(&msg), numbers[1 - i]); - received[i] = true; - } - } - } - if received.iter().all(|&r| r) { - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Timeout waiting for initial summation"); - - // Start the late joining node - info!("Starting late joining node"); - handles.push(nodes[2].start().await.expect("Failed to start late node")); - let all_handles: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); - - // Wait for the new node to complete handshakes - wait_for_all_handshakes(&all_handles, TEST_TIMEOUT).await; - - // Late node sends its number and receives history - info!("Late node sending number and receiving history"); - handles[2] - .send_protocol_message(create_protocol_message( - PROTOCOL_NAME, - SummationMessage::Number(numbers[2]), - &handles[2], - None, - )) - .expect("Failed to send number from late node"); - - // Verify final state - timeout(TEST_TIMEOUT, async { - let mut verified = vec![false; handles.len()]; - loop { - for (i, handle) in handles.iter_mut().enumerate() { - if let Some(msg) = handle.next_protocol_message() { - let num = extract_number_from_message(&msg); - if !verified[i] && (num == numbers[2] || numbers.contains(&num)) { - verified[i] = true; - } - } - } - if verified.iter().all(|&v| v) { - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Timeout waiting for late node synchronization"); - - info!("Late join test completed successfully"); -} - -#[tokio::test] -async fn test_summation_protocol_node_disconnect() { - super::init_tracing(); - info!("Starting node disconnect test"); - - // Create 3 nodes - let mut nodes = create_whitelisted_nodes(3).await; - - // Start all nodes - let mut handles = Vec::new(); - for node in nodes.iter_mut() { - handles.push(node.start().await.expect("Failed to start node")); - } - - // Convert handles to mutable references - let handles_refs: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); - - // Wait for all handshakes - wait_for_all_handshakes(&handles_refs, TEST_TIMEOUT).await; - - // Send initial numbers - let numbers = vec![42, 58, 100]; - for (i, handle) in handles.iter().enumerate() { - handle - .send_protocol_message(create_protocol_message( - PROTOCOL_NAME, - SummationMessage::Number(numbers[i]), - handle, - None, - )) - .expect("Failed to send number"); - } - - // Wait for initial processing - timeout(TEST_TIMEOUT, async { - let mut received = vec![0; handles.len()]; - loop { - for (i, handle) in handles.iter_mut().enumerate() { - if let Some(msg) = handle.next_protocol_message() { - if received[i] < 2 { - extract_number_from_message(&msg); - received[i] += 1; - } - } - } - if received.iter().all(|&r| r == 2) { - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Timeout waiting for initial messages"); - - // Disconnect one node - info!("Disconnecting node"); - drop(handles.pop()); - let mut remaining_handles: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); - - // Verify remaining nodes can still communicate - info!("Verifying remaining nodes can communicate"); - for (i, sender) in remaining_handles.iter().enumerate() { - for (j, recipient) in remaining_handles.iter().enumerate() { - if i != j { - sender - .send_protocol_message(create_protocol_message( - PROTOCOL_NAME, - SummationMessage::Verification { sum: numbers[i] }, - sender, - Some(recipient.local_peer_id), - )) - .expect("Failed to send verification"); - } - } - } - - // Wait for verification messages - timeout(TEST_TIMEOUT, async { - let mut verified = vec![false; remaining_handles.len()]; - loop { - for (i, handle) in remaining_handles.iter_mut().enumerate() { - if let Some(msg) = handle.next_protocol_message() { - if !verified[i] { - extract_sum_from_verification(&msg); - verified[i] = true; - } - } - } - if verified.iter().all(|&v| v) { - break; - } - tokio::time::sleep(Duration::from_millis(100)).await; - } - }) - .await - .expect("Timeout waiting for verification after disconnect"); - - info!("Node disconnect test completed successfully"); -} +// TODO +// use gadget_crypto::KeyType; +// use libp2p::PeerId; +// use serde::{Deserialize, Serialize}; +// use std::{collections::HashSet, time::Duration}; +// use tokio::time::timeout; +// use tracing::info; +// +// use crate::{ +// blueprint_protocol::{InstanceMessageRequest, InstanceMessageResponse}, +// key_types::{Curve, InstanceMsgKeyPair, InstanceMsgPublicKey}, +// service::NetworkMessage, +// service_handle::NetworkServiceHandle, +// tests::{ +// create_whitelisted_nodes, wait_for_all_handshakes, wait_for_handshake_completion, TestNode, +// }, +// types::{MessageRouting, ParticipantId, ParticipantInfo, ProtocolMessage}, +// }; +// +// const TEST_TIMEOUT: Duration = Duration::from_secs(10); +// const PROTOCOL_NAME: &str = "summation/1.0.0"; +// +// // Protocol message types +// #[derive(Debug, Clone, Serialize, Deserialize)] +// enum SummationMessage { +// Number(u64), +// Verification { sum: u64 }, +// } +// +// // Helper to create a protocol message +// fn create_protocol_message( +// protocol: &str, +// msg: SummationMessage, +// sender: &NetworkServiceHandle, +// recipient: Option, +// ) -> ProtocolMessage { +// ProtocolMessage { +// protocol: protocol.to_string(), +// routing: MessageRouting { +// message_id: 0, +// round_id: 0, +// sender: ParticipantInfo { +// id: ParticipantId(0), +// public_key: None, +// }, +// recipient: recipient.map(|peer_id| ParticipantInfo { +// id: ParticipantId(0), +// public_key: None, +// }), +// }, +// payload: bincode::serialize(&msg).expect("Failed to serialize message"), +// } +// } +// +// // Helper to extract number from message +// fn extract_number_from_message(msg: &ProtocolMessage) -> u64 { +// match bincode::deserialize::(&msg.payload).expect("Failed to deserialize") { +// SummationMessage::Number(n) => n, +// _ => panic!("Expected number message"), +// } +// } +// +// // Helper to extract sum from verification message +// fn extract_sum_from_verification(msg: &ProtocolMessage) -> u64 { +// match bincode::deserialize::(&msg.payload).expect("Failed to deserialize") { +// SummationMessage::Verification { sum } => sum, +// _ => panic!("Expected verification message"), +// } +// } +// +// #[tokio::test] +// async fn test_summation_protocol_basic() { +// super::init_tracing(); +// info!("Starting summation protocol test"); +// +// // Create nodes with whitelisted keys +// let instance_key_pair2 = Curve::generate_with_seed(None).unwrap(); +// let mut allowed_keys1 = HashSet::new(); +// allowed_keys1.insert(instance_key_pair2.public()); +// +// let mut node1 = TestNode::new("test-net", "sum-test", allowed_keys1, vec![]).await; +// +// let mut allowed_keys2 = HashSet::new(); +// allowed_keys2.insert(node1.instance_key_pair.public()); +// let mut node2 = TestNode::new_with_keys( +// "test-net", +// "sum-test", +// allowed_keys2, +// vec![], +// Some(instance_key_pair2), +// None, +// ) +// .await; +// +// info!("Starting nodes"); +// let mut handle1 = node1.start().await.expect("Failed to start node1"); +// let mut handle2 = node2.start().await.expect("Failed to start node2"); +// +// info!("Waiting for handshake completion"); +// wait_for_handshake_completion(&handle1, &handle2, TEST_TIMEOUT).await; +// +// // Generate test numbers +// let num1 = 42; +// let num2 = 58; +// let expected_sum = num1 + num2; +// +// info!("Sending numbers via gossip"); +// // Send numbers via gossip +// handle1 +// .send(create_protocol_message( +// PROTOCOL_NAME, +// SummationMessage::Number(num1), +// &handle1, +// None, +// )) +// .expect("Failed to send number from node1"); +// +// handle2 +// .send(create_protocol_message( +// PROTOCOL_NAME, +// SummationMessage::Number(num2), +// &handle2, +// None, +// )) +// .expect("Failed to send number from node2"); +// +// info!("Waiting for messages to be processed"); +// // Wait for messages and compute sums +// let mut sum1 = 0; +// let mut sum2 = 0; +// let mut node1_received = false; +// let mut node2_received = false; +// +// timeout(TEST_TIMEOUT, async { +// loop { +// // Process incoming messages +// if let Some(msg) = handle1.next_protocol_message() { +// if !node1_received { +// sum1 += extract_number_from_message(&msg); +// node1_received = true; +// } +// } +// if let Some(msg) = handle2.next_protocol_message() { +// if !node2_received { +// sum2 += extract_number_from_message(&msg); +// node2_received = true; +// } +// } +// +// // Check if both nodes have received messages +// if node1_received && node2_received { +// break; +// } +// tokio::time::sleep(Duration::from_millis(100)).await; +// } +// }) +// .await +// .expect("Timeout waiting for summation completion"); +// +// info!("Verifying sums via P2P messages"); +// // Verify sums via P2P messages +// handle1 +// .send(create_protocol_message( +// PROTOCOL_NAME, +// SummationMessage::Verification { sum: sum1 }, +// &handle1, +// Some(node2.peer_id), +// )) +// .expect("Failed to send verification from node1"); +// +// handle2 +// .send(create_protocol_message( +// PROTOCOL_NAME, +// SummationMessage::Verification { sum: sum2 }, +// &handle2, +// Some(node1.peer_id), +// )) +// .expect("Failed to send verification from node2"); +// +// info!("Waiting for verification messages"); +// // Wait for verification messages +// timeout(TEST_TIMEOUT, async { +// let mut node1_verified = false; +// let mut node2_verified = false; +// +// loop { +// // Process verification messages +// if let Some(msg) = handle1.next_protocol_message() { +// if !node1_verified { +// assert_eq!(extract_sum_from_verification(&msg), expected_sum); +// node1_verified = true; +// } +// } +// if let Some(msg) = handle2.next_protocol_message() { +// if !node2_verified { +// assert_eq!(extract_sum_from_verification(&msg), expected_sum); +// node2_verified = true; +// } +// } +// +// if node1_verified && node2_verified { +// break; +// } +// tokio::time::sleep(Duration::from_millis(100)).await; +// } +// }) +// .await +// .expect("Timeout waiting for verification completion"); +// +// info!("Summation protocol test completed successfully"); +// } +// +// #[tokio::test] +// async fn test_summation_protocol_multi_node() { +// super::init_tracing(); +// info!("Starting multi-node summation protocol test"); +// +// // Create 3 nodes with whitelisted keys +// info!("Creating whitelisted nodes"); +// let mut nodes = create_whitelisted_nodes(3).await; +// info!("Created {} nodes successfully", nodes.len()); +// +// // Start all nodes +// info!("Starting all nodes"); +// let mut handles = Vec::new(); +// for (i, node) in nodes.iter_mut().enumerate() { +// info!("Starting node {}", i); +// handles.push(node.start().await.expect("Failed to start node")); +// info!("Node {} started successfully", i); +// } +// +// // Convert handles to mutable references +// info!("Converting handles to mutable references"); +// let mut handles: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); +// let handles_len = handles.len(); +// info!("Converted {} handles", handles_len); +// +// // Wait for all handshakes to complete +// info!( +// "Waiting for handshake completion between {} nodes", +// handles_len +// ); +// wait_for_all_handshakes(&handles, TEST_TIMEOUT).await; +// info!("All handshakes completed successfully"); +// +// // Generate test numbers +// let numbers = vec![42, 58, 100]; +// let expected_sum: u64 = numbers.iter().sum(); +// info!( +// "Generated test numbers: {:?}, expected sum: {}", +// numbers, expected_sum +// ); +// +// info!("Sending numbers via gossip"); +// // Each node broadcasts its number +// for (i, handle) in handles.iter().enumerate() { +// info!("Node {} broadcasting number {}", i, numbers[i]); +// handle +// .send(create_protocol_message( +// PROTOCOL_NAME, +// SummationMessage::Number(numbers[i]), +// handle, +// None, +// )) +// .expect("Failed to send number"); +// info!("Node {} successfully broadcast its number", i); +// // Add a small delay between broadcasts to avoid message collisions +// tokio::time::sleep(Duration::from_millis(100)).await; +// } +// +// info!("Waiting for messages to be processed"); +// // Wait for all nodes to receive all numbers +// let mut sums = vec![0; handles_len]; +// let mut received = vec![0; handles_len]; +// +// timeout(TEST_TIMEOUT, async { +// loop { +// for (i, handle) in handles.iter_mut().enumerate() { +// if let Some(msg) = handle.next_protocol_message() { +// if received[i] < handles_len - 1 { +// let num = extract_number_from_message(&msg); +// sums[i] += num; +// received[i] += 1; +// info!( +// "Node {} received number {}, total sum: {}, received count: {}", +// i, num, sums[i], received[i] +// ); +// } +// } +// } +// +// let all_received = received.iter().all(|&r| r == handles_len - 1); +// info!( +// "Current received counts: {:?}, target count: {}", +// received, +// handles_len - 1 +// ); +// if all_received { +// info!("All nodes have received all numbers"); +// break; +// } +// tokio::time::sleep(Duration::from_millis(100)).await; +// } +// }) +// .await +// .expect("Timeout waiting for summation completion"); +// +// info!("Verifying sums via P2P messages"); +// info!("Final sums: {:?}", sums); +// // Each node verifies with every other node +// for (i, sender) in handles.iter().enumerate() { +// for (j, recipient) in handles.iter().enumerate() { +// if i != j { +// info!( +// "Node {} sending verification sum {} to node {}", +// i, sums[i], j +// ); +// sender +// .send(create_protocol_message( +// PROTOCOL_NAME, +// SummationMessage::Verification { sum: sums[i] }, +// sender, +// Some(recipient.local_peer_id), +// )) +// .expect("Failed to send verification"); +// } +// } +// } +// +// info!("Waiting for verification messages"); +// // Wait for all verifications +// timeout(TEST_TIMEOUT, async { +// let mut verified = vec![0; handles_len]; +// loop { +// for (i, handle) in handles.iter_mut().enumerate() { +// if let Some(msg) = handle.next_protocol_message() { +// if verified[i] < handles_len - 1 { +// let sum = extract_sum_from_verification(&msg); +// info!( +// "Node {} received verification sum {}, expected {}", +// i, sum, expected_sum +// ); +// assert_eq!(sum, expected_sum); +// verified[i] += 1; +// info!("Node {} verification count: {}", i, verified[i]); +// } +// } +// } +// +// let all_verified = verified.iter().all(|&v| v == handles_len - 1); +// info!("Current verification counts: {:?}", verified); +// if all_verified { +// info!("All nodes have verified all sums"); +// break; +// } +// tokio::time::sleep(Duration::from_millis(100)).await; +// } +// }) +// .await +// .expect("Timeout waiting for verification completion"); +// +// info!("Multi-node summation protocol test completed successfully"); +// } +// +// #[tokio::test] +// async fn test_summation_protocol_late_join() { +// super::init_tracing(); +// info!("Starting late join summation protocol test"); +// +// // Create 3 nodes but only start 2 initially +// let mut nodes = create_whitelisted_nodes(3).await; +// +// // Start first two nodes +// let mut handles = Vec::new(); +// for node in nodes[..2].iter_mut() { +// handles.push(node.start().await.expect("Failed to start node")); +// } +// +// // Convert handles to mutable references +// let handles_refs: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); +// +// // Wait for initial handshakes +// info!("Waiting for initial handshake completion"); +// wait_for_all_handshakes(&handles_refs, TEST_TIMEOUT).await; +// +// // Initial nodes send their numbers +// let numbers = vec![42, 58, 100]; +// let _expected_sum: u64 = numbers.iter().sum(); +// +// info!("Initial nodes sending numbers"); +// for (i, handle) in handles.iter().enumerate() { +// handle +// .send(create_protocol_message( +// PROTOCOL_NAME, +// SummationMessage::Number(numbers[i]), +// handle, +// None, +// )) +// .expect("Failed to send number"); +// } +// +// // Wait for initial nodes to process messages +// timeout(TEST_TIMEOUT, async { +// let mut received = vec![false; 2]; +// loop { +// for (i, handle) in handles.iter_mut().enumerate() { +// if let Some(msg) = handle.next_protocol_message() { +// if !received[i] { +// assert_eq!(extract_number_from_message(&msg), numbers[1 - i]); +// received[i] = true; +// } +// } +// } +// if received.iter().all(|&r| r) { +// break; +// } +// tokio::time::sleep(Duration::from_millis(100)).await; +// } +// }) +// .await +// .expect("Timeout waiting for initial summation"); +// +// // Start the late joining node +// info!("Starting late joining node"); +// handles.push(nodes[2].start().await.expect("Failed to start late node")); +// let all_handles: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); +// +// // Wait for the new node to complete handshakes +// wait_for_all_handshakes(&all_handles, TEST_TIMEOUT).await; +// +// // Late node sends its number and receives history +// info!("Late node sending number and receiving history"); +// handles[2] +// .send(create_protocol_message( +// PROTOCOL_NAME, +// SummationMessage::Number(numbers[2]), +// &handles[2], +// None, +// )) +// .expect("Failed to send number from late node"); +// +// // Verify final state +// timeout(TEST_TIMEOUT, async { +// let mut verified = vec![false; handles.len()]; +// loop { +// for (i, handle) in handles.iter_mut().enumerate() { +// if let Some(msg) = handle.next_protocol_message() { +// let num = extract_number_from_message(&msg); +// if !verified[i] && (num == numbers[2] || numbers.contains(&num)) { +// verified[i] = true; +// } +// } +// } +// if verified.iter().all(|&v| v) { +// break; +// } +// tokio::time::sleep(Duration::from_millis(100)).await; +// } +// }) +// .await +// .expect("Timeout waiting for late node synchronization"); +// +// info!("Late join test completed successfully"); +// } +// +// #[tokio::test] +// async fn test_summation_protocol_node_disconnect() { +// super::init_tracing(); +// info!("Starting node disconnect test"); +// +// // Create 3 nodes +// let mut nodes = create_whitelisted_nodes(3).await; +// +// // Start all nodes +// let mut handles = Vec::new(); +// for node in nodes.iter_mut() { +// handles.push(node.start().await.expect("Failed to start node")); +// } +// +// // Convert handles to mutable references +// let handles_refs: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); +// +// // Wait for all handshakes +// wait_for_all_handshakes(&handles_refs, TEST_TIMEOUT).await; +// +// // Send initial numbers +// let numbers = vec![42, 58, 100]; +// for (i, handle) in handles.iter().enumerate() { +// handle +// .send(create_protocol_message( +// PROTOCOL_NAME, +// SummationMessage::Number(numbers[i]), +// handle, +// None, +// )) +// .expect("Failed to send number"); +// } +// +// // Wait for initial processing +// timeout(TEST_TIMEOUT, async { +// let mut received = vec![0; handles.len()]; +// loop { +// for (i, handle) in handles.iter_mut().enumerate() { +// if let Some(msg) = handle.next_protocol_message() { +// if received[i] < 2 { +// extract_number_from_message(&msg); +// received[i] += 1; +// } +// } +// } +// if received.iter().all(|&r| r == 2) { +// break; +// } +// tokio::time::sleep(Duration::from_millis(100)).await; +// } +// }) +// .await +// .expect("Timeout waiting for initial messages"); +// +// // Disconnect one node +// info!("Disconnecting node"); +// drop(handles.pop()); +// let mut remaining_handles: Vec<&mut NetworkServiceHandle> = handles.iter_mut().collect(); +// +// // Verify remaining nodes can still communicate +// info!("Verifying remaining nodes can communicate"); +// for (i, sender) in remaining_handles.iter().enumerate() { +// for (j, recipient) in remaining_handles.iter().enumerate() { +// if i != j { +// sender +// .send(create_protocol_message( +// PROTOCOL_NAME, +// SummationMessage::Verification { sum: numbers[i] }, +// sender, +// Some(recipient.local_peer_id), +// )) +// .expect("Failed to send verification"); +// } +// } +// } +// +// // Wait for verification messages +// timeout(TEST_TIMEOUT, async { +// let mut verified = vec![false; remaining_handles.len()]; +// loop { +// for (i, handle) in remaining_handles.iter_mut().enumerate() { +// if let Some(msg) = handle.next_protocol_message() { +// if !verified[i] { +// extract_sum_from_verification(&msg); +// verified[i] = true; +// } +// } +// } +// if verified.iter().all(|&v| v) { +// break; +// } +// tokio::time::sleep(Duration::from_millis(100)).await; +// } +// }) +// .await +// .expect("Timeout waiting for verification after disconnect"); +// +// info!("Node disconnect test completed successfully"); +// } diff --git a/crates/networking/src/tests/gossip.rs b/crates/networking/src/tests/gossip.rs index efa6e1d4c..c3240f533 100644 --- a/crates/networking/src/tests/gossip.rs +++ b/crates/networking/src/tests/gossip.rs @@ -11,7 +11,7 @@ use tokio::time::timeout; use tracing::info; const TEST_TIMEOUT: Duration = Duration::from_secs(10); -const PROTOCOL_NAME: &str = "test-gossip"; +const PROTOCOL_NAME: &str = "/blueprint_protocol/gossip-test/1.0.0"; #[tokio::test] async fn test_gossip_between_verified_peers() { @@ -45,24 +45,21 @@ async fn test_gossip_between_verified_peers() { wait_for_handshake_completion(&handle1, &handle2, TEST_TIMEOUT).await; // Create test message + info!("Sending gossip message from node1"); + let test_payload = b"Hello, gossip network!".to_vec(); - let message = ProtocolMessage { - protocol: PROTOCOL_NAME.to_string(), - routing: MessageRouting { - message_id: 1, - round_id: 0, - sender: ParticipantInfo { - id: ParticipantId(1), - public_key: Some(node1.instance_key_pair.public()), - }, - recipient: None, // No specific recipient for gossip + let routing = MessageRouting { + message_id: 1, + round_id: 0, + sender: ParticipantInfo { + id: ParticipantId(1), + public_key: Some(node1.instance_key_pair.public()), }, - payload: test_payload.clone(), + recipient: None, // No specific recipient for gossip }; - info!("Sending gossip message from node1"); handle1 - .send_protocol_message(message) + .send(routing, test_payload.clone()) .expect("Failed to send gossip message"); info!("Waiting for node2 to receive the message"); @@ -70,6 +67,7 @@ async fn test_gossip_between_verified_peers() { let received_message = timeout(TEST_TIMEOUT, async { loop { if let Some(msg) = handle2.next_protocol_message() { + dbg!(&msg); if msg.protocol == PROTOCOL_NAME { return msg; } @@ -114,23 +112,19 @@ async fn test_multi_node_gossip() { // Create test message let test_payload = b"Multi-node gossip test".to_vec(); - let message = ProtocolMessage { - protocol: PROTOCOL_NAME.to_string(), - routing: MessageRouting { - message_id: 1, - round_id: 0, - sender: ParticipantInfo { - id: ParticipantId(0), - public_key: Some(nodes[0].instance_key_pair.public()), - }, - recipient: None, + let routing = MessageRouting { + message_id: 1, + round_id: 0, + sender: ParticipantInfo { + id: ParticipantId(0), + public_key: Some(nodes[0].instance_key_pair.public()), }, - payload: test_payload.clone(), + recipient: None, }; info!("Sending gossip message from node 0"); handles[0] - .send_protocol_message(message) + .send(routing, test_payload.clone()) .expect("Failed to send gossip message"); info!("Waiting for all nodes to receive the message"); @@ -180,23 +174,19 @@ async fn test_unverified_peer_gossip() { // Create test message let test_payload = b"This message should not be received".to_vec(); - let message = ProtocolMessage { - protocol: PROTOCOL_NAME.to_string(), - routing: MessageRouting { - message_id: 1, - round_id: 0, - sender: ParticipantInfo { - id: ParticipantId(1), - public_key: Some(node1.instance_key_pair.public()), - }, - recipient: None, + let routing = MessageRouting { + message_id: 1, + round_id: 0, + sender: ParticipantInfo { + id: ParticipantId(1), + public_key: Some(node1.instance_key_pair.public()), }, - payload: test_payload, + recipient: None, }; info!("Attempting to send gossip message from unverified node"); handle1 - .send_protocol_message(message) + .send(routing, test_payload.clone()) .expect("Failed to send gossip message"); // Wait a bit to ensure message is not received From e1d2c349a6190c1266806725e48b27dca5465fae Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 19 Feb 2025 15:59:49 -0700 Subject: [PATCH 44/52] chore: fix names --- crates/networking/src/service.rs | 11 +++++++++-- crates/networking/src/service_handle.rs | 16 ++++++++-------- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index ca9cc9cc9..156663081 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -182,7 +182,10 @@ impl NetworkService { .unwrap() .build(); - swarm.behaviour_mut().blueprint_protocol.subscribe(&blueprint_protocol_name)?; + swarm + .behaviour_mut() + .blueprint_protocol + .subscribe(&blueprint_protocol_name)?; // Start listening swarm.listen_on(listen_addr)?; @@ -215,7 +218,11 @@ impl NetworkService { // Create handle with new interface let handle = NetworkServiceHandle::new( local_peer_id, - self.swarm.behaviour().blueprint_protocol.blueprint_protocol_name.clone(), + self.swarm + .behaviour() + .blueprint_protocol + .blueprint_protocol_name + .clone(), self.peer_manager.clone(), network_sender, protocol_message_receiver, diff --git a/crates/networking/src/service_handle.rs b/crates/networking/src/service_handle.rs index 35b6660e1..14cc55b3d 100644 --- a/crates/networking/src/service_handle.rs +++ b/crates/networking/src/service_handle.rs @@ -1,3 +1,4 @@ +use crate::types::{MessageRouting, ParticipantInfo}; use crate::{ blueprint_protocol::InstanceMessageRequest, discovery::{PeerInfo, PeerManager}, @@ -5,12 +6,11 @@ use crate::{ types::ProtocolMessage, }; use crossbeam_channel::{self, Receiver, Sender}; +use libp2p::gossipsub::Sha256Topic; use libp2p::{Multiaddr, PeerId}; use std::sync::Arc; -use libp2p::gossipsub::Sha256Topic; use tokio::task::JoinHandle; use tracing::{debug, info}; -use crate::types::{MessageRouting, ParticipantInfo}; /// Handle for sending outgoing messages to the network #[derive(Clone)] @@ -114,9 +114,8 @@ impl NetworkServiceHandle { #[must_use] pub fn send(&self, routing: MessageRouting, message: impl Into>) -> Result<(), String> { - let protocol = Sha256Topic::new(&*self.blueprint_protocol_name).to_string(); let protocol_message = ProtocolMessage { - protocol: protocol.clone(), + protocol: self.blueprint_protocol_name.clone().to_string(), routing, payload: message.into(), }; @@ -125,7 +124,7 @@ impl NetworkServiceHandle { match protocol_message.routing.recipient { Some(recipient) => { let instance_message_request = InstanceMessageRequest::Protocol { - protocol, + protocol: self.blueprint_protocol_name.clone().to_string(), payload: raw_payload, metadata: None, }; @@ -134,7 +133,8 @@ impl NetworkServiceHandle { return Ok(()); }; - let Some(peer_id) = self.peer_manager.get_peer_id_from_public_key(&public_key) else { + let Some(peer_id) = self.peer_manager.get_peer_id_from_public_key(&public_key) + else { return Ok(()); }; @@ -143,11 +143,11 @@ impl NetworkServiceHandle { request: instance_message_request, })?; debug!("Sent outbound p2p `NetworkMessage` to {:?}", peer_id); - }, + } None => { let gossip_message = NetworkMessage::GossipMessage { source: self.local_peer_id, - topic: protocol, + topic: self.blueprint_protocol_name.clone().to_string(), message: raw_payload, }; self.send_network_message(gossip_message)?; From db84f812ad602105210871fc7306b1a752d68b77 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Thu, 20 Feb 2025 01:23:07 -0700 Subject: [PATCH 45/52] chore: lock --- Cargo.lock | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 48d10266c..8ed8ab724 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4397,9 +4397,9 @@ dependencies = [ [[package]] name = "cxx" -version = "1.0.140" +version = "1.0.141" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc49567e08c72902f4cbc7242ee8d874ec9cbe97fbabf77b4e0e1f447513e13a" +checksum = "8bc580dceb395cae0efdde0a88f034cfd8a276897e40c693a7b87bed17971d33" dependencies = [ "cc", "cxxbridge-cmd", @@ -4411,9 +4411,9 @@ dependencies = [ [[package]] name = "cxx-build" -version = "1.0.140" +version = "1.0.141" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe46b5309c99e9775e7a338c98e4097455f52db5b684fd793ca22848fde6e371" +checksum = "49d8c1baedad72a7efda12ad8d7ad687b3e7221dfb304a12443fd69e9de8bb30" dependencies = [ "cc", "codespan-reporting", @@ -4425,9 +4425,9 @@ dependencies = [ [[package]] name = "cxxbridge-cmd" -version = "1.0.140" +version = "1.0.141" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4315c4ce8d23c26d87f2f83698725fd5718d8e6ace4a9093da2664d23294d372" +checksum = "e43afb0e3b2ef293492a31ecd796af902112460d53e5f923f7804f348a769f9c" dependencies = [ "clap", "codespan-reporting", @@ -4438,15 +4438,15 @@ dependencies = [ [[package]] name = "cxxbridge-flags" -version = "1.0.140" +version = "1.0.141" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f55d69deb3a92f610a60ecc524a72c7374b6dc822f8fb7bb4e5d9473f10530c4" +checksum = "0257ad2096a2474fe877e9e055ab69603851c3d6b394efcc7e0443899c2492ce" [[package]] name = "cxxbridge-macro" -version = "1.0.140" +version = "1.0.141" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5bee7a1d9b5091462002c2b8de2a4ed0f0fde011d503cc272633f66075bd5141" +checksum = "b46cbd7358a46b760609f1cb5093683328e58ca50e594a308716f5403fdc03e5" dependencies = [ "proc-macro2", "quote", @@ -14452,9 +14452,9 @@ dependencies = [ [[package]] name = "proc-macro-warning" -version = "1.0.2" +version = "1.84.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "834da187cfe638ae8abb0203f0b33e5ccdb02a28e7199f2f47b3e2754f50edca" +checksum = "75eea531cfcd120e0851a3f8aed42c4841f78c889eefafd96339c72677ae42c3" dependencies = [ "proc-macro2", "quote", From 6455c608b30111b5c86119476d73f9f9677ceb4c Mon Sep 17 00:00:00 2001 From: Shady Khalifa Date: Thu, 20 Feb 2025 19:51:26 +0200 Subject: [PATCH 46/52] fix: round-based networking adapter --- .../src/lib.rs | 76 +++++++++---------- .../tests/common/mod.rs | 2 +- .../tests/rand_protocol.rs | 23 +++++- crates/networking/src/service.rs | 2 +- crates/networking/src/service_handle.rs | 5 +- 5 files changed, 61 insertions(+), 47 deletions(-) diff --git a/crates/networking-round-based-extension/src/lib.rs b/crates/networking-round-based-extension/src/lib.rs index 1a1c1ff5c..1b4224367 100644 --- a/crates/networking-round-based-extension/src/lib.rs +++ b/crates/networking-round-based-extension/src/lib.rs @@ -34,10 +34,6 @@ pub struct RoundBasedNetworkAdapter { parties: Arc>, /// Counter for message IDs next_msg_id: Arc, - /// Channel for forwarding messages - forward_tx: Sender<(PartyIndex, M)>, - /// Channel for receiving forwarded messages - forward_rx: Receiver<(PartyIndex, M)>, /// Protocol identifier protocol_id: String, _phantom: std::marker::PhantomData, @@ -55,15 +51,11 @@ where parties: HashMap, protocol_id: impl Into, ) -> Self { - let (forward_tx, forward_rx) = crossbeam_channel::unbounded(); - Self { handle, party_index, parties: Arc::new(DashMap::from_iter(parties)), next_msg_id: Arc::new(AtomicU64::new(0)), - forward_tx, - forward_rx, protocol_id: protocol_id.into(), _phantom: std::marker::PhantomData, } @@ -87,8 +79,6 @@ where party_index, parties, next_msg_id, - forward_tx, - forward_rx, protocol_id, .. } = self; @@ -99,11 +89,10 @@ where parties: parties.clone(), next_msg_id: next_msg_id.clone(), protocol_id: protocol_id.clone(), - forward_tx, _phantom: std::marker::PhantomData, }; - let receiver = RoundBasedReceiver::new(handle, party_index, forward_rx); + let receiver = RoundBasedReceiver::new(handle, party_index); (receiver, sender) } @@ -115,13 +104,12 @@ pub struct RoundBasedSender { parties: Arc>, next_msg_id: Arc, protocol_id: String, - forward_tx: Sender<(PartyIndex, M)>, _phantom: std::marker::PhantomData, } impl Sink> for RoundBasedSender where - M: Serialize + round_based::ProtocolMessage + Clone, + M: Serialize + round_based::ProtocolMessage + Clone + Unpin, { type Error = NetworkError; @@ -130,17 +118,18 @@ where } fn start_send(self: Pin<&mut Self>, outgoing: Outgoing) -> Result<(), Self::Error> { - let this = unsafe { self.get_unchecked_mut() }; + let this = self.get_mut(); let msg_id = this.next_msg_id.fetch_add(1, Ordering::Relaxed); let round = outgoing.msg.round(); - // Handle local message forwarding for self-messages - if outgoing.recipient == MessageDestination::OneParty(this.party_index) { - return this - .forward_tx - .send((this.party_index, outgoing.msg.clone())) - .map_err(|_| NetworkError::Send("Failed to forward local message".into())); - } + tracing::trace!( + i = %this.party_index, + recipient = ?outgoing.recipient, + %round, + %msg_id, + protocol_id = %this.protocol_id, + "Sending message", + ); let (recipient, recipient_key) = match outgoing.recipient { MessageDestination::AllParties => (None, None), @@ -168,8 +157,15 @@ where .map_err(|e| NetworkError::Serialization(e))?, }; + tracing::trace!( + %round, + %msg_id, + protocol_id = %this.protocol_id, + "Sending message to network", + ); + this.handle - .send(protocol_message) + .send(protocol_message.routing, protocol_message.payload) .map_err(|e| NetworkError::Send(e)) } @@ -185,20 +181,14 @@ where pub struct RoundBasedReceiver { handle: NetworkServiceHandle, party_index: PartyIndex, - forward_rx: Receiver<(PartyIndex, M)>, _phantom: std::marker::PhantomData, } impl RoundBasedReceiver { - fn new( - handle: NetworkServiceHandle, - party_index: PartyIndex, - forward_rx: Receiver<(PartyIndex, M)>, - ) -> Self { + fn new(handle: NetworkServiceHandle, party_index: PartyIndex) -> Self { Self { handle, party_index, - forward_rx, _phantom: std::marker::PhantomData, } } @@ -211,16 +201,6 @@ where type Item = Result, NetworkError>; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - // First check forwarded messages - if let Ok((sender, msg)) = self.forward_rx.try_recv() { - return Poll::Ready(Some(Ok(Incoming { - msg, - sender, - id: 0, - msg_type: MessageType::P2P, - }))); - } - // Get a mutable reference to self let this = self.get_mut(); @@ -236,6 +216,15 @@ where let sender = protocol_message.routing.sender.id.0; let id = protocol_message.routing.message_id; + tracing::trace!( + i = %this.party_index, + sender = ?sender, + %id, + protocol_id = %protocol_message.protocol, + ?msg_type, + size = %protocol_message.payload.len(), + "Received message", + ); match serde_json::from_slice(&protocol_message.payload) { Ok(msg) => Poll::Ready(Some(Ok(Incoming { msg, @@ -246,7 +235,12 @@ where Err(e) => Poll::Ready(Some(Err(NetworkError::Serialization(e)))), } } - None => Poll::Pending, + None => { + tracing::trace!(i = %this.party_index, "No message received; the waker will wake us up when there is a new message"); + // In this case, tell the waker to wake us up when there is a new message + cx.waker().wake_by_ref(); + Poll::Pending + } } } } diff --git a/crates/networking-round-based-extension/tests/common/mod.rs b/crates/networking-round-based-extension/tests/common/mod.rs index 0b4e2be25..a33372c52 100644 --- a/crates/networking-round-based-extension/tests/common/mod.rs +++ b/crates/networking-round-based-extension/tests/common/mod.rs @@ -18,7 +18,7 @@ pub fn init_tracing() { let _ = tracing_subscriber::fmt() .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) .with_target(true) - .with_thread_ids(true) + .with_thread_ids(false) .with_file(true) .with_line_number(true) .try_init(); diff --git a/crates/networking-round-based-extension/tests/rand_protocol.rs b/crates/networking-round-based-extension/tests/rand_protocol.rs index 978a5fd07..0259a0bc2 100644 --- a/crates/networking-round-based-extension/tests/rand_protocol.rs +++ b/crates/networking-round-based-extension/tests/rand_protocol.rs @@ -35,6 +35,7 @@ pub struct DecommitMsg { } /// Carries out the randomness generation protocol +#[tracing::instrument(skip(party, rng))] pub async fn protocol_of_random_generation( party: M, i: PartyIndex, @@ -60,8 +61,11 @@ where let mut local_randomness = [0u8; 32]; rng.fill_bytes(&mut local_randomness); + tracing::debug!(local_randomness = %hex::encode(&local_randomness), "Generated local randomness"); + // 2. Commit local randomness (broadcast m=sha256(randomness)) let commitment = Sha256::digest(local_randomness); + tracing::debug!(commitment = %hex::encode(&commitment), "Committed local randomness"); outgoing .send(Outgoing::broadcast(Msg::CommitMsg(CommitMsg { commitment, @@ -69,13 +73,18 @@ where .await .map_err(Error::Round1Send)?; + tracing::debug!("Sent commitment and waiting for others to send theirs"); + // 3. Receive committed randomness from other parties let commitments = rounds .complete(round1) .await .map_err(Error::Round1Receive)?; + tracing::debug!("Received commitments from all parties"); + // 4. Open local randomness + tracing::debug!("Opening local randomness"); outgoing .send(Outgoing::broadcast(Msg::DecommitMsg(DecommitMsg { randomness: local_randomness, @@ -83,12 +92,16 @@ where .await .map_err(Error::Round2Send)?; + tracing::debug!("Sent decommitment and waiting for others to send theirs"); + // 5. Receive opened local randomness from other parties, verify them, and output protocol randomness let randomness = rounds .complete(round2) .await .map_err(Error::Round2Receive)?; + tracing::debug!("Received decommitments from all parties"); + let mut guilty_parties = vec![]; let mut output = local_randomness; for ((party_i, com_msg_id, commit), (_, decom_msg_id, decommit)) in commitments @@ -112,8 +125,11 @@ where } if !guilty_parties.is_empty() { + tracing::error!(guilty_parties = ?guilty_parties, "Some parties cheated"); Err(Error::PartiesOpenedRandomnessDoesntMatchCommitment { guilty_parties }) } else { + tracing::debug!(output = %hex::encode(&output), "Generated randomness"); + tracing::info!("Randomness generation protocol completed successfully."); Ok(output) } } @@ -156,6 +172,7 @@ pub struct Blame { #[cfg(test)] mod tests { use std::collections::{HashMap, HashSet}; + use std::time::Duration; use super::common::*; use gadget_networking::{Curve, KeyType}; @@ -201,7 +218,7 @@ mod tests { std::println!("Output randomness: {}", hex::encode(randomness)); } - #[tokio::test] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn p2p_networking() { init_tracing(); let network_name = "rand-test-network"; @@ -233,6 +250,10 @@ mod tests { let handle1 = node1.start().await.expect("Failed to start node1"); let handle2 = node2.start().await.expect("Failed to start node2"); + wait_for_peer_discovery(&[&handle1, &handle2], Duration::from_secs(5)) + .await + .unwrap(); + let parties = HashMap::from_iter([ (0, node1.instance_key_pair.public()), (1, node2.instance_key_pair.public()), diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index 156663081..423d74b65 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -425,7 +425,7 @@ async fn handle_discovery_event( info!(%peer_info.peer_id, "Newly discovered peer from Kademlia"); let info = PeerInfo::default(); peer_manager.update_peer(peer_info.peer_id, info); - let addrs: Vec<_> = peer_info.addrs.iter().cloned().collect(); + let addrs: Vec<_> = peer_info.addrs.to_vec(); for addr in addrs { debug!(%peer_info.peer_id, %addr, "Dialing peer from Kademlia"); if let Err(e) = swarm.dial(DialOpts::from(addr)) { diff --git a/crates/networking/src/service_handle.rs b/crates/networking/src/service_handle.rs index 14cc55b3d..2c893542a 100644 --- a/crates/networking/src/service_handle.rs +++ b/crates/networking/src/service_handle.rs @@ -1,4 +1,4 @@ -use crate::types::{MessageRouting, ParticipantInfo}; +use crate::types::MessageRouting; use crate::{ blueprint_protocol::InstanceMessageRequest, discovery::{PeerInfo, PeerManager}, @@ -6,11 +6,10 @@ use crate::{ types::ProtocolMessage, }; use crossbeam_channel::{self, Receiver, Sender}; -use libp2p::gossipsub::Sha256Topic; use libp2p::{Multiaddr, PeerId}; use std::sync::Arc; use tokio::task::JoinHandle; -use tracing::{debug, info}; +use tracing::debug; /// Handle for sending outgoing messages to the network #[derive(Clone)] From d1cd1e7bead798865ee067aa00389b8a2fc5c104 Mon Sep 17 00:00:00 2001 From: Shady Khalifa Date: Thu, 20 Feb 2025 20:43:20 +0200 Subject: [PATCH 47/52] fix: allowed keys --- crates/config/src/lib.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 10cdc5878..d139e2abf 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -137,9 +137,8 @@ impl GadgetConfiguration { pub fn libp2p_start_network( &self, network_config: gadget_networking::NetworkConfig, + allowed_keys: gadget_std::collections::HashSet, ) -> Result { - // TODO: Add allowed keys or pass them. - let allowed_keys = Default::default(); let networking_service = gadget_networking::NetworkService::new(network_config, allowed_keys)?; From 0e73291753b0eb4de180bcc0d68483f9b2617f74 Mon Sep 17 00:00:00 2001 From: Serial <69764315+Serial-ATA@users.noreply.github.com> Date: Thu, 20 Feb 2025 14:53:16 -0500 Subject: [PATCH 48/52] fix(tangle-testing-utils): stop expecting call ID 0 --- Cargo.lock | 1 - crates/testing-utils/tangle/src/harness.rs | 2 +- crates/testing-utils/tangle/src/node/transactions.rs | 3 +-- 3 files changed, 2 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index df7257778..8ed8ab724 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7191,7 +7191,6 @@ dependencies = [ "sha2 0.10.8", "thiserror 2.0.11", "tokio", - "tokio-stream", "tracing", "tracing-subscriber 0.3.19", ] diff --git a/crates/testing-utils/tangle/src/harness.rs b/crates/testing-utils/tangle/src/harness.rs index cd877dc54..0dba8a27c 100644 --- a/crates/testing-utils/tangle/src/harness.rs +++ b/crates/testing-utils/tangle/src/harness.rs @@ -348,7 +348,7 @@ impl TangleTestHarness { service_id, Job::from(job_id), inputs, - 0, + 0, // TODO: Should this take a call ID? or leave it up to the caller to verify? ) .await .map_err(|e| Error::Setup(e.to_string()))?; diff --git a/crates/testing-utils/tangle/src/node/transactions.rs b/crates/testing-utils/tangle/src/node/transactions.rs index 9a4713fa5..2d3479076 100644 --- a/crates/testing-utils/tangle/src/node/transactions.rs +++ b/crates/testing-utils/tangle/src/node/transactions.rs @@ -198,7 +198,7 @@ pub async fn submit_job>( service_id: u64, job_id: Job, job_params: Args, - call_id: u64, + _call_id: u64, // TODO: Actually verify this ) -> Result { let call = api::tx().services().call(service_id, job_id, job_params); let events = client @@ -215,7 +215,6 @@ pub async fn submit_job>( if job_called.service_id == service_id && job_called.job == job_id && user.account_id() == job_called.caller - && job_called.call_id == call_id { return Ok(job_called); } From 8d76767373caad7e26190e24d4c0bf42796721c0 Mon Sep 17 00:00:00 2001 From: Shady Khalifa Date: Thu, 20 Feb 2025 21:53:37 +0200 Subject: [PATCH 49/52] fix: use 32 bit for target peers --- crates/config/src/context_config.rs | 2 +- crates/config/src/lib.rs | 2 +- crates/networking/src/behaviours.rs | 22 ++++++-------------- crates/networking/src/discovery/behaviour.rs | 6 +++--- crates/networking/src/discovery/config.rs | 4 ++-- crates/networking/src/service.rs | 4 ++-- 6 files changed, 15 insertions(+), 25 deletions(-) diff --git a/crates/config/src/context_config.rs b/crates/config/src/context_config.rs index 2ec4e66ba..c508dc0b8 100644 --- a/crates/config/src/context_config.rs +++ b/crates/config/src/context_config.rs @@ -48,7 +48,7 @@ pub enum GadgetCLICoreSettings { #[cfg(feature = "networking")] #[arg(long, env)] #[serde(default)] - target_peer_count: Option, + target_peer_count: Option, #[arg(long, short = 'd', env)] keystore_uri: String, #[arg(long, value_enum, env)] diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index d139e2abf..2f9a5683c 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -129,7 +129,7 @@ pub struct GadgetConfiguration { pub enable_kademlia: bool, /// The target number of peers to connect to #[cfg(feature = "networking")] - pub target_peer_count: u64, + pub target_peer_count: u32, } impl GadgetConfiguration { diff --git a/crates/networking/src/behaviours.rs b/crates/networking/src/behaviours.rs index ab3b6a347..dcefec29e 100644 --- a/crates/networking/src/behaviours.rs +++ b/crates/networking/src/behaviours.rs @@ -25,8 +25,6 @@ use std::{ }; use tracing::{debug, info}; -const MAX_ESTABLISHED_PER_PEER: u32 = 4; - /// Events that can be emitted by the `GadgetBehavior` #[derive(Debug)] pub enum GadgetEvent { @@ -57,25 +55,17 @@ impl GadgetBehaviour { blueprint_protocol_name: &str, local_key: &Keypair, instance_key_pair: &InstanceMsgKeyPair, - target_peer_count: u64, + target_peer_count: u32, peer_manager: Arc, protocol_message_sender: Sender, ) -> Self { let connection_limits = connection_limits::Behaviour::new( ConnectionLimits::default() - .with_max_pending_incoming(Some( - target_peer_count as u32 * MAX_ESTABLISHED_PER_PEER, - )) - .with_max_pending_outgoing(Some( - target_peer_count as u32 * MAX_ESTABLISHED_PER_PEER, - )) - .with_max_established_incoming(Some( - target_peer_count as u32 * MAX_ESTABLISHED_PER_PEER, - )) - .with_max_established_outgoing(Some( - target_peer_count as u32 * MAX_ESTABLISHED_PER_PEER, - )) - .with_max_established_per_peer(Some(MAX_ESTABLISHED_PER_PEER)), + .with_max_pending_incoming(Some(target_peer_count)) + .with_max_pending_outgoing(Some(target_peer_count)) + .with_max_established_incoming(Some(target_peer_count)) + .with_max_established_outgoing(Some(target_peer_count)) + .with_max_established_per_peer(Some(target_peer_count)), ); let ping = ping::Behaviour::new(ping::Config::new().with_interval(Duration::from_secs(30))); diff --git a/crates/networking/src/discovery/behaviour.rs b/crates/networking/src/discovery/behaviour.rs index 515d2b345..95d4d11eb 100644 --- a/crates/networking/src/discovery/behaviour.rs +++ b/crates/networking/src/discovery/behaviour.rs @@ -11,7 +11,7 @@ use libp2p::{ core::Multiaddr, identify, identity::PeerId, - kad::{self, store::MemoryStore, Event as KademliaEvent}, + kad::{self, store::MemoryStore}, mdns::{self, Event as MdnsEvent}, relay, swarm::{ @@ -69,13 +69,13 @@ pub struct DiscoveryBehaviour { /// Events to return in priority when polled. pub pending_events: VecDeque, /// Number of nodes we're currently connected to. - pub n_node_connected: u64, + pub n_node_connected: u32, /// Peers pub peers: HashSet, /// Peer info pub peer_info: HashMap, /// Target peer count - pub target_peer_count: u64, + pub target_peer_count: u32, /// Options to configure dials to known peers. pub pending_dial_opts: VecDeque, } diff --git a/crates/networking/src/discovery/config.rs b/crates/networking/src/discovery/config.rs index 3a5ff1c74..aa06b573a 100644 --- a/crates/networking/src/discovery/config.rs +++ b/crates/networking/src/discovery/config.rs @@ -21,7 +21,7 @@ pub struct DiscoveryConfig { /// The relay nodes. relay_nodes: Vec<(PeerId, Multiaddr)>, /// The number of peers to connect to. - target_peer_count: u64, + target_peer_count: u32, /// Enable mDNS discovery. enable_mdns: bool, /// Enable Kademlia discovery. @@ -74,7 +74,7 @@ impl DiscoveryConfig { self } - pub fn target_peer_count(mut self, count: u64) -> Self { + pub fn target_peer_count(mut self, count: u32) -> Self { self.target_peer_count = count; self } diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index 423d74b65..632ee3f73 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -97,7 +97,7 @@ pub struct NetworkConfig { /// Address to listen on pub listen_addr: Multiaddr, /// Target number of peers to maintain - pub target_peer_count: u64, + pub target_peer_count: u32, /// Bootstrap peers to connect to pub bootstrap_peers: Vec, /// Whether to enable mDNS discovery @@ -149,7 +149,7 @@ impl NetworkService { } = config; let peer_manager = Arc::new(PeerManager::new(allowed_keys)); - let blueprint_protocol_name = format!("/blueprint_protocol/{}/1.0.0", instance_id); + let blueprint_protocol_name = format!("{network_name}/{instance_id}"); let (network_sender, network_receiver) = crossbeam_channel::unbounded(); let (protocol_message_sender, protocol_message_receiver) = crossbeam_channel::unbounded(); From 885c72373c0d5f021a2746546f8bbf6dc8bc41e9 Mon Sep 17 00:00:00 2001 From: Serial <69764315+Serial-ATA@users.noreply.github.com> Date: Thu, 20 Feb 2025 15:35:11 -0500 Subject: [PATCH 50/52] chore: clippy --- Cargo.toml | 1 + .../src/lib.rs | 2 +- crates/networking/src/behaviours.rs | 9 ++-- .../src/blueprint_protocol/behaviour.rs | 4 +- .../src/blueprint_protocol/handler.rs | 9 ++-- crates/networking/src/discovery/behaviour.rs | 6 +-- crates/networking/src/discovery/config.rs | 23 ++++++++-- crates/networking/src/discovery/mod.rs | 3 +- crates/networking/src/discovery/peers.rs | 12 +++--- crates/networking/src/error.rs | 3 ++ crates/networking/src/service.rs | 42 +++++++++---------- crates/networking/src/service_handle.rs | 1 - 12 files changed, 63 insertions(+), 52 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index adc070995..529a06b36 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,6 +22,7 @@ all = { level = "deny", priority = -1 } single_match_else = "allow" uninlined_format_args = "allow" needless_late_init = "allow" +struct_excessive_bools = "allow" [workspace.lints.rustdoc] broken_intra_doc_links = "deny" diff --git a/crates/networking-round-based-extension/src/lib.rs b/crates/networking-round-based-extension/src/lib.rs index 1b4224367..75159769a 100644 --- a/crates/networking-round-based-extension/src/lib.rs +++ b/crates/networking-round-based-extension/src/lib.rs @@ -236,7 +236,7 @@ where } } None => { - tracing::trace!(i = %this.party_index, "No message received; the waker will wake us up when there is a new message"); + //tracing::trace!(i = %this.party_index, "No message received; the waker will wake us up when there is a new message"); // In this case, tell the waker to wake us up when there is a new message cx.waker().wake_by_ref(); Poll::Pending diff --git a/crates/networking/src/behaviours.rs b/crates/networking/src/behaviours.rs index dcefec29e..40b3b737e 100644 --- a/crates/networking/src/behaviours.rs +++ b/crates/networking/src/behaviours.rs @@ -58,7 +58,7 @@ impl GadgetBehaviour { target_peer_count: u32, peer_manager: Arc, protocol_message_sender: Sender, - ) -> Self { + ) -> NetworkingResult { let connection_limits = connection_limits::Behaviour::new( ConnectionLimits::default() .with_max_pending_incoming(Some(target_peer_count)) @@ -78,8 +78,7 @@ impl GadgetBehaviour { .mdns(true) .kademlia(true) .target_peer_count(target_peer_count) - .build() - .unwrap(); + .build()?; info!( "Setting up blueprint protocol with name: {}", @@ -94,12 +93,12 @@ impl GadgetBehaviour { ); debug!("Created GadgetBehaviour with all components initialized"); - Self { + Ok(Self { connection_limits, discovery, blueprint_protocol, ping, - } + }) } /// Bootstrap Kademlia network diff --git a/crates/networking/src/blueprint_protocol/behaviour.rs b/crates/networking/src/blueprint_protocol/behaviour.rs index 4b033946f..52f3104ed 100644 --- a/crates/networking/src/blueprint_protocol/behaviour.rs +++ b/crates/networking/src/blueprint_protocol/behaviour.rs @@ -11,7 +11,7 @@ use dashmap::DashMap; use gadget_crypto::KeyType; use libp2p::{ core::transport::PortUse, - gossipsub::{self, IdentTopic, MessageAuthenticity, MessageId, Sha256Topic}, + gossipsub::{self, IdentTopic, MessageId, Sha256Topic}, identity::Keypair, request_response::{self, OutboundRequestId, ResponseChannel}, swarm::{ @@ -69,7 +69,7 @@ pub struct BlueprintProtocolBehaviour { pub(crate) peer_manager: Arc, /// Libp2p peer ID pub(crate) local_peer_id: PeerId, - /// Instance key pair for handshakes and blueprint_protocol + /// Instance key pair for handshakes and blueprint protocol pub(crate) instance_key_pair: InstanceMsgKeyPair, /// Peers with pending inbound handshakes pub(crate) inbound_handshakes: DashMap, diff --git a/crates/networking/src/blueprint_protocol/handler.rs b/crates/networking/src/blueprint_protocol/handler.rs index fbe7d86bf..5c970b2a6 100644 --- a/crates/networking/src/blueprint_protocol/handler.rs +++ b/crates/networking/src/blueprint_protocol/handler.rs @@ -1,9 +1,6 @@ use std::time::{Duration, Instant}; -use libp2p::{ - gossipsub::{self}, - request_response, PeerId, -}; +use libp2p::{request_response, PeerId}; use tracing::{debug, warn}; use crate::blueprint_protocol::HandshakeMessage; @@ -141,7 +138,7 @@ impl BlueprintProtocolBehaviour { InstanceMessageRequest::Protocol { protocol, payload, - metadata, + metadata: _, }, channel, .. @@ -204,7 +201,7 @@ impl BlueprintProtocolBehaviour { peer, message: request_response::Message::Response { - response: InstanceMessageResponse::Success { protocol, data }, + response: InstanceMessageResponse::Success { protocol, data: _ }, .. }, .. diff --git a/crates/networking/src/discovery/behaviour.rs b/crates/networking/src/discovery/behaviour.rs index 95d4d11eb..10f872618 100644 --- a/crates/networking/src/discovery/behaviour.rs +++ b/crates/networking/src/discovery/behaviour.rs @@ -306,9 +306,9 @@ impl NetworkBehaviour for DiscoveryBehaviour { DerivedDiscoveryBehaviourEvent::Kademlia(ev) => match ev { // Adding to Kademlia buckets is automatic with our config, // no need to do manually. - kad::Event::RoutingUpdated { .. } => {} - kad::Event::RoutablePeer { .. } => {} - kad::Event::PendingRoutablePeer { .. } => { + kad::Event::RoutingUpdated { .. } + | kad::Event::RoutablePeer { .. } + | kad::Event::PendingRoutablePeer { .. } => { // Intentionally ignore } other => { diff --git a/crates/networking/src/discovery/config.rs b/crates/networking/src/discovery/config.rs index aa06b573a..a90522a14 100644 --- a/crates/networking/src/discovery/config.rs +++ b/crates/networking/src/discovery/config.rs @@ -2,6 +2,7 @@ use super::{ behaviour::{DerivedDiscoveryBehaviour, DiscoveryBehaviour}, new_kademlia, }; +use crate::error::Result; use libp2p::{ autonat, identify, identity::PublicKey, mdns, relay, upnp, Multiaddr, PeerId, StreamProtocol, }; @@ -40,6 +41,7 @@ pub struct DiscoveryConfig { } impl DiscoveryConfig { + #[must_use] pub fn new(local_public_key: PublicKey, network_name: impl Into) -> Self { Self { local_peer_id: local_public_key.to_peer_id(), @@ -59,47 +61,60 @@ impl DiscoveryConfig { /// Set the protocol version that uniquely identifies your P2P service. /// This should be unique to your application to avoid conflicts with other P2P networks. /// Format recommendation: "/" + #[must_use] pub fn protocol_version(mut self, version: impl Into) -> Self { self.protocol_version = version.into(); self } + #[must_use] pub fn bootstrap_peers(mut self, peers: Vec<(PeerId, Multiaddr)>) -> Self { self.bootstrap_peers = peers; self } + #[must_use] pub fn relay_nodes(mut self, nodes: Vec<(PeerId, Multiaddr)>) -> Self { self.relay_nodes = nodes; self } + #[must_use] pub fn target_peer_count(mut self, count: u32) -> Self { self.target_peer_count = count; self } + #[must_use] pub fn mdns(mut self, enable: bool) -> Self { self.enable_mdns = enable; self } + #[must_use] pub fn kademlia(mut self, enable: bool) -> Self { self.enable_kademlia = enable; self } + #[must_use] pub fn upnp(mut self, enable: bool) -> Self { self.enable_upnp = enable; self } + #[must_use] pub fn relay(mut self, enable: bool) -> Self { self.enable_relay = enable; self } - pub fn build(self) -> anyhow::Result { + /// Construct this [`DiscoveryConfig`] into a [`DiscoveryBehaviour`] + /// + /// # Errors + /// + /// If `mdns` is enabled, see [mdns::Behaviour::new] + pub fn build(self) -> Result { let kademlia_opt = if self.enable_kademlia { let protocol = StreamProtocol::try_from_owned(format!( "/gadget/kad/{}/kad/1.0.0", @@ -125,7 +140,7 @@ impl DiscoveryConfig { let mdns_opt = if self.enable_mdns { Some(mdns::Behaviour::new( - Default::default(), + mdns::Config::default(), self.local_peer_id, )?) } else { @@ -139,7 +154,7 @@ impl DiscoveryConfig { }; let relay_opt = if self.enable_relay { - let relay = relay::Behaviour::new(self.local_peer_id, Default::default()); + let relay = relay::Behaviour::new(self.local_peer_id, relay::Config::default()); Some(relay) } else { None @@ -153,7 +168,7 @@ impl DiscoveryConfig { .with_agent_version(format!("gadget-{}", env!("CARGO_PKG_VERSION"))) .with_push_listen_addr_updates(true), ), - autonat: autonat::Behaviour::new(self.local_peer_id, Default::default()), + autonat: autonat::Behaviour::new(self.local_peer_id, autonat::Config::default()), upnp: upnp_opt.into(), relay: relay_opt.into(), }; diff --git a/crates/networking/src/discovery/mod.rs b/crates/networking/src/discovery/mod.rs index 54b3e7f4e..2b63a01e7 100644 --- a/crates/networking/src/discovery/mod.rs +++ b/crates/networking/src/discovery/mod.rs @@ -11,9 +11,8 @@ pub mod peers; pub use peers::{PeerEvent, PeerInfo, PeerManager}; -const MAX_ESTABLISHED_PER_PEER: u32 = 4; - #[must_use] +#[allow(clippy::missing_panics_doc)] pub fn new_kademlia(peer_id: PeerId, protocol: StreamProtocol) -> kad::Behaviour { let store = kad::store::MemoryStore::new(peer_id); let mut config = kad::Config::new(protocol); diff --git a/crates/networking/src/discovery/peers.rs b/crates/networking/src/discovery/peers.rs index b07d78b35..db83fb7e9 100644 --- a/crates/networking/src/discovery/peers.rs +++ b/crates/networking/src/discovery/peers.rs @@ -76,7 +76,7 @@ pub struct PeerManager { impl Default for PeerManager { fn default() -> Self { - Self::new(Default::default()) + Self::new(HashSet::default()) } } @@ -85,10 +85,10 @@ impl PeerManager { pub fn new(whitelisted_keys: HashSet) -> Self { let (event_tx, _) = broadcast::channel(100); Self { - peers: Default::default(), - banned_peers: Default::default(), - verified_peers: Default::default(), - public_keys_to_peer_ids: Default::default(), + peers: DashMap::default(), + banned_peers: DashMap::default(), + verified_peers: DashSet::default(), + public_keys_to_peer_ids: Arc::new(DashMap::default()), whitelisted_keys: DashSet::from_iter(whitelisted_keys), event_tx, } @@ -172,7 +172,7 @@ impl PeerManager { } /// Bans a peer with the default duration(`1h`) - pub async fn ban_peer_with_default_duration(&self, peer: PeerId, reason: impl Into) { + pub fn ban_peer_with_default_duration(&self, peer: PeerId, reason: impl Into) { const BAN_PEER_DURATION: Duration = Duration::from_secs(60 * 60); //1h self.ban_peer(peer, reason, Some(BAN_PEER_DURATION)); } diff --git a/crates/networking/src/error.rs b/crates/networking/src/error.rs index 005673143..b424d28f0 100644 --- a/crates/networking/src/error.rs +++ b/crates/networking/src/error.rs @@ -47,6 +47,9 @@ pub enum Error { Io(#[from] std::io::Error), // libp2p compat + #[error(transparent)] + InvalidProtocol(#[from] libp2p::swarm::InvalidProtocol), + #[error(transparent)] NoKnownPeers(#[from] libp2p::kad::NoKnownPeers), diff --git a/crates/networking/src/service.rs b/crates/networking/src/service.rs index 632ee3f73..1b2455713 100644 --- a/crates/networking/src/service.rs +++ b/crates/networking/src/service.rs @@ -1,8 +1,4 @@ -use std::{ - collections::{HashMap, HashSet}, - sync::Arc, - time::Duration, -}; +use std::{collections::HashSet, sync::Arc, time::Duration}; use crate::{ behaviours::{GadgetBehaviour, GadgetBehaviourEvent}, @@ -12,7 +8,7 @@ use crate::{ PeerInfo, PeerManager, }, error::Error, - key_types::{InstanceMsgKeyPair, InstanceMsgPublicKey, InstanceSignedMsgSignature}, + key_types::{InstanceMsgKeyPair, InstanceMsgPublicKey}, service_handle::NetworkServiceHandle, types::ProtocolMessage, }; @@ -131,6 +127,12 @@ pub struct NetworkService { impl NetworkService { /// Create a new network service + /// + /// # Errors + /// + /// * See [`GadgetBehaviour::new`] + /// * Bad `listen_addr` in the provided [`NetworkConfig`] + #[allow(clippy::missing_panics_doc)] // Unwrapping an Infallible pub fn new( config: NetworkConfig, allowed_keys: HashSet, @@ -144,8 +146,8 @@ impl NetworkService { listen_addr, target_peer_count, bootstrap_peers, - enable_mdns, - enable_kademlia, + enable_mdns: _, + enable_kademlia: _, } = config; let peer_manager = Arc::new(PeerManager::new(allowed_keys)); @@ -164,7 +166,7 @@ impl NetworkService { target_peer_count, peer_manager.clone(), protocol_message_sender.clone(), - ); + )?; let mut swarm = SwarmBuilder::with_existing_identity(local_key) .with_tokio() @@ -294,7 +296,6 @@ impl NetworkService { &self.peer_manager, &self.event_sender, ) - .await { warn!("Failed to handle network message: {}", e); } @@ -336,14 +337,13 @@ async fn handle_behaviour_event( match event { GadgetBehaviourEvent::ConnectionLimits(_) => {} GadgetBehaviourEvent::Discovery(discovery_event) => { - handle_discovery_event(swarm, peer_manager, discovery_event, event_sender).await?; + handle_discovery_event(swarm, peer_manager, discovery_event, event_sender)?; } GadgetBehaviourEvent::BlueprintProtocol(blueprint_event) => { - handle_blueprint_protocol_event(swarm, peer_manager, blueprint_event, event_sender) - .await?; + handle_blueprint_protocol_event(swarm, peer_manager, blueprint_event, event_sender)?; } GadgetBehaviourEvent::Ping(ping_event) => { - handle_ping_event(swarm, peer_manager, ping_event, event_sender).await?; + handle_ping_event(swarm, peer_manager, ping_event, event_sender)?; } } @@ -351,7 +351,7 @@ async fn handle_behaviour_event( } /// Handle a discovery event -async fn handle_discovery_event( +fn handle_discovery_event( swarm: &mut Swarm, peer_manager: &Arc, event: DiscoveryEvent, @@ -390,9 +390,7 @@ async fn handle_discovery_event( &swarm.behaviour().blueprint_protocol.blueprint_protocol_name; if !protocols.contains(blueprint_protocol_name) { warn!(%peer_id, %blueprint_protocol_name, "Peer does not support required protocol"); - peer_manager - .ban_peer_with_default_duration(*peer_id, "protocol unsupported") - .await; + peer_manager.ban_peer_with_default_duration(*peer_id, "protocol unsupported"); return Ok(()); } @@ -425,7 +423,7 @@ async fn handle_discovery_event( info!(%peer_info.peer_id, "Newly discovered peer from Kademlia"); let info = PeerInfo::default(); peer_manager.update_peer(peer_info.peer_id, info); - let addrs: Vec<_> = peer_info.addrs.to_vec(); + let addrs: Vec<_> = peer_info.addrs.clone(); for addr in addrs { debug!(%peer_info.peer_id, %addr, "Dialing peer from Kademlia"); if let Err(e) = swarm.dial(DialOpts::from(addr)) { @@ -458,7 +456,7 @@ async fn handle_discovery_event( } /// Handle a blueprint event -async fn handle_blueprint_protocol_event( +fn handle_blueprint_protocol_event( _swarm: &mut Swarm, _peer_manager: &Arc, event: BlueprintProtocolEvent, @@ -490,7 +488,7 @@ async fn handle_blueprint_protocol_event( } /// Handle a ping event -async fn handle_ping_event( +fn handle_ping_event( _swarm: &mut Swarm, _peer_manager: &Arc, event: ping::Event, @@ -519,7 +517,7 @@ async fn handle_ping_event( } /// Handle a network message -async fn handle_network_message( +fn handle_network_message( swarm: &mut Swarm, msg: NetworkMessage, peer_manager: &Arc, diff --git a/crates/networking/src/service_handle.rs b/crates/networking/src/service_handle.rs index 2c893542a..6f62832a7 100644 --- a/crates/networking/src/service_handle.rs +++ b/crates/networking/src/service_handle.rs @@ -111,7 +111,6 @@ impl NetworkServiceHandle { self.peer_manager.get_peer_info(peer_id) } - #[must_use] pub fn send(&self, routing: MessageRouting, message: impl Into>) -> Result<(), String> { let protocol_message = ProtocolMessage { protocol: self.blueprint_protocol_name.clone().to_string(), From 46a7a7bb2d5d932d0ee212c1c92f4d4900f5591f Mon Sep 17 00:00:00 2001 From: Serial <69764315+Serial-ATA@users.noreply.github.com> Date: Thu, 20 Feb 2025 15:38:45 -0500 Subject: [PATCH 51/52] refactor(networking): move round-based extension crate --- Cargo.toml | 2 +- .../extensions/round-based}/Cargo.toml | 2 +- .../extensions/round-based}/src/lib.rs | 0 .../extensions/round-based}/src/tests.rs | 0 .../extensions/round-based}/tests/common/mod.rs | 0 .../extensions/round-based}/tests/rand_protocol.rs | 0 6 files changed, 2 insertions(+), 2 deletions(-) rename crates/{networking-round-based-extension => networking/extensions/round-based}/Cargo.toml (96%) rename crates/{networking-round-based-extension => networking/extensions/round-based}/src/lib.rs (100%) rename crates/{networking-round-based-extension => networking/extensions/round-based}/src/tests.rs (100%) rename crates/{networking-round-based-extension => networking/extensions/round-based}/tests/common/mod.rs (100%) rename crates/{networking-round-based-extension => networking/extensions/round-based}/tests/rand_protocol.rs (100%) diff --git a/Cargo.toml b/Cargo.toml index 529a06b36..b9d76f626 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -102,7 +102,7 @@ gadget-std = { version = "0.1.0", path = "./crates/std", default-features = fals # P2P gadget-networking = { version = "0.1.0", path = "./crates/networking", default-features = false } -gadget-networking-round-based-extension = { version = "0.1.0", path = "./crates/networking-round-based-extension", default-features = false } +gadget-networking-round-based-extension = { version = "0.1.0", path = "./crates/networking/extensions/round-based", default-features = false } # Utilities gadget-utils = { version = "0.1.0", path = "./crates/utils", default-features = false } diff --git a/crates/networking-round-based-extension/Cargo.toml b/crates/networking/extensions/round-based/Cargo.toml similarity index 96% rename from crates/networking-round-based-extension/Cargo.toml rename to crates/networking/extensions/round-based/Cargo.toml index 491addc40..006905b4d 100644 --- a/crates/networking-round-based-extension/Cargo.toml +++ b/crates/networking/extensions/round-based/Cargo.toml @@ -8,7 +8,7 @@ homepage.workspace = true repository.workspace = true [dependencies] -gadget-networking = { path = "../networking" } +gadget-networking = { workspace = true } round-based = { workspace = true } tokio = { workspace = true } futures = { workspace = true } diff --git a/crates/networking-round-based-extension/src/lib.rs b/crates/networking/extensions/round-based/src/lib.rs similarity index 100% rename from crates/networking-round-based-extension/src/lib.rs rename to crates/networking/extensions/round-based/src/lib.rs diff --git a/crates/networking-round-based-extension/src/tests.rs b/crates/networking/extensions/round-based/src/tests.rs similarity index 100% rename from crates/networking-round-based-extension/src/tests.rs rename to crates/networking/extensions/round-based/src/tests.rs diff --git a/crates/networking-round-based-extension/tests/common/mod.rs b/crates/networking/extensions/round-based/tests/common/mod.rs similarity index 100% rename from crates/networking-round-based-extension/tests/common/mod.rs rename to crates/networking/extensions/round-based/tests/common/mod.rs diff --git a/crates/networking-round-based-extension/tests/rand_protocol.rs b/crates/networking/extensions/round-based/tests/rand_protocol.rs similarity index 100% rename from crates/networking-round-based-extension/tests/rand_protocol.rs rename to crates/networking/extensions/round-based/tests/rand_protocol.rs From 741e794a4319c9696e7fcd4b69be7cd2f8e8e182 Mon Sep 17 00:00:00 2001 From: Serial <69764315+Serial-ATA@users.noreply.github.com> Date: Thu, 20 Feb 2025 15:44:07 -0500 Subject: [PATCH 52/52] chore: remove round_based compat feature from networking --- crates/networking/src/lib.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/crates/networking/src/lib.rs b/crates/networking/src/lib.rs index 119dc6aee..3656cbd36 100644 --- a/crates/networking/src/lib.rs +++ b/crates/networking/src/lib.rs @@ -12,9 +12,6 @@ pub mod types; #[cfg(test)] mod tests; -#[cfg(feature = "round-based-compat")] -pub use gadget_networking_round_based_extension as round_based_compat; - pub use gadget_crypto::KeyType; pub use key_types::*; pub use service::{NetworkConfig, NetworkEvent, NetworkService};