use crate::rpc::{RPCEvent, RPCMessage, Rpc}; use crate::NetworkConfig; use futures::prelude::*; use libp2p::{ core::{ swarm::{NetworkBehaviourAction, NetworkBehaviourEventProcess}, PublicKey, }, gossipsub::{Gossipsub, GossipsubEvent}, identify::{protocol::IdentifyInfo, Identify, IdentifyEvent}, ping::{Ping, PingEvent}, tokio_io::{AsyncRead, AsyncWrite}, NetworkBehaviour, PeerId, }; use slog::{debug, o, trace, warn}; use ssz::{ssz_encode, Decode, DecodeError, Encode}; use types::{Attestation, BeaconBlock}; use types::{Topic, TopicHash}; /// Builds the network behaviour for the libp2p Swarm. /// Implements gossipsub message routing. #[derive(NetworkBehaviour)] #[behaviour(out_event = "BehaviourEvent", poll_method = "poll")] pub struct Behaviour { /// The routing pub-sub mechanism for eth2. gossipsub: Gossipsub, // TODO: Add Kademlia for peer discovery /// The events generated by this behaviour to be consumed in the swarm poll. serenity_rpc: Rpc, /// Allows discovery of IP addresses for peers on the network. identify: Identify, /// Keep regular connection to peers and disconnect if absent. // TODO: Keepalive, likely remove this later. // TODO: Make the ping time customizeable. ping: Ping, #[behaviour(ignore)] events: Vec, /// Logger for behaviour actions. #[behaviour(ignore)] log: slog::Logger, } // Implement the NetworkBehaviourEventProcess trait so that we can derive NetworkBehaviour for Behaviour impl NetworkBehaviourEventProcess for Behaviour { fn inject_event(&mut self, event: GossipsubEvent) { match event { GossipsubEvent::Message(gs_msg) => { trace!(self.log, "Received GossipEvent"; "msg" => format!("{:?}", gs_msg)); let pubsub_message = match PubsubMessage::from_ssz_bytes(&gs_msg.data) { //TODO: Punish peer on error Err(e) => { warn!( self.log, "Received undecodable message from Peer {:?} error", gs_msg.source; "error" => format!("{:?}", e) ); return; } Ok(msg) => msg, }; self.events.push(BehaviourEvent::GossipMessage { source: gs_msg.source, topics: gs_msg.topics, message: Box::new(pubsub_message), }); } GossipsubEvent::Subscribed { .. } => {} GossipsubEvent::Unsubscribed { .. } => {} } } } impl NetworkBehaviourEventProcess for Behaviour { fn inject_event(&mut self, event: RPCMessage) { match event { RPCMessage::PeerDialed(peer_id) => { self.events.push(BehaviourEvent::PeerDialed(peer_id)) } RPCMessage::RPC(peer_id, rpc_event) => { self.events.push(BehaviourEvent::RPC(peer_id, rpc_event)) } } } } impl NetworkBehaviourEventProcess for Behaviour { fn inject_event(&mut self, event: IdentifyEvent) { match event { IdentifyEvent::Identified { peer_id, mut info, .. } => { if info.listen_addrs.len() > 20 { debug!( self.log, "More than 20 peers have been identified, truncating" ); info.listen_addrs.truncate(20); } self.events .push(BehaviourEvent::Identified(peer_id, Box::new(info))); } IdentifyEvent::Error { .. } => {} IdentifyEvent::SendBack { .. } => {} } } } impl NetworkBehaviourEventProcess for Behaviour { fn inject_event(&mut self, _event: PingEvent) { // not interested in ping responses at the moment. } } impl Behaviour { pub fn new(local_public_key: PublicKey, net_conf: &NetworkConfig, log: &slog::Logger) -> Self { let local_peer_id = local_public_key.clone().into_peer_id(); let identify_config = net_conf.identify_config.clone(); let behaviour_log = log.new(o!()); Behaviour { gossipsub: Gossipsub::new(local_peer_id, net_conf.gs_config.clone()), serenity_rpc: Rpc::new(log), identify: Identify::new( identify_config.version, identify_config.user_agent, local_public_key, ), ping: Ping::new(), events: Vec::new(), log: behaviour_log, } } /// Consumes the events list when polled. fn poll( &mut self, ) -> Async> { if !self.events.is_empty() { return Async::Ready(NetworkBehaviourAction::GenerateEvent(self.events.remove(0))); } Async::NotReady } } /// Implements the combined behaviour for the libp2p service. impl Behaviour { /// Subscribes to a gossipsub topic. pub fn subscribe(&mut self, topic: Topic) -> bool { self.gossipsub.subscribe(topic) } /// Sends an RPC Request/Response via the RPC protocol. pub fn send_rpc(&mut self, peer_id: PeerId, rpc_event: RPCEvent) { self.serenity_rpc.send_rpc(peer_id, rpc_event); } /// Publishes a message on the pubsub (gossipsub) behaviour. pub fn publish(&mut self, topics: Vec, message: PubsubMessage) { let message_bytes = ssz_encode(&message); for topic in topics { self.gossipsub.publish(topic, message_bytes.clone()); } } } /// The types of events than can be obtained from polling the behaviour. pub enum BehaviourEvent { RPC(PeerId, RPCEvent), PeerDialed(PeerId), Identified(PeerId, Box), // TODO: This is a stub at the moment GossipMessage { source: PeerId, topics: Vec, message: Box, }, } /// Messages that are passed to and from the pubsub (Gossipsub) behaviour. #[derive(Debug, Clone, PartialEq)] pub enum PubsubMessage { /// Gossipsub message providing notification of a new block. Block(BeaconBlock), /// Gossipsub message providing notification of a new attestation. Attestation(Attestation), } //TODO: Correctly encode/decode enums. Prefixing with integer for now. impl Encode for PubsubMessage { fn is_ssz_fixed_len() -> bool { false } fn ssz_append(&self, buf: &mut Vec) { let offset = ::ssz_fixed_len() + as Encode>::ssz_fixed_len(); let mut encoder = ssz::SszEncoder::container(buf, offset); match self { PubsubMessage::Block(block_gossip) => { encoder.append(&0_u32); // Encode the gossip as a Vec; encoder.append(&block_gossip.as_ssz_bytes()); } PubsubMessage::Attestation(attestation_gossip) => { encoder.append(&1_u32); // Encode the gossip as a Vec; encoder.append(&attestation_gossip.as_ssz_bytes()); } } encoder.finalize(); } } impl Decode for PubsubMessage { fn is_ssz_fixed_len() -> bool { false } fn from_ssz_bytes(bytes: &[u8]) -> Result { let mut builder = ssz::SszDecoderBuilder::new(&bytes); builder.register_type::()?; builder.register_type::>()?; let mut decoder = builder.build()?; let id: u32 = decoder.decode_next()?; let body: Vec = decoder.decode_next()?; match id { 0 => Ok(PubsubMessage::Block(BeaconBlock::from_ssz_bytes(&body)?)), 1 => Ok(PubsubMessage::Attestation(Attestation::from_ssz_bytes( &body, )?)), _ => Err(DecodeError::BytesInvalid( "Invalid PubsubMessage id".to_string(), )), } } } #[cfg(test)] mod test { use super::*; use types::*; #[test] fn ssz_encoding() { let original = PubsubMessage::Block(BeaconBlock::empty(&FoundationEthSpec::spec())); let encoded = ssz_encode(&original); let decoded = PubsubMessage::from_ssz_bytes(&encoded).unwrap(); assert_eq!(original, decoded); } }