lighthouse/beacon_node/libp2p/src/behaviour.rs

use futures::prelude::*;
use libp2p::{
    core::swarm::{NetworkBehaviourAction, NetworkBehaviourEventProcess},
    gossipsub::{Gossipsub, GossipsubConfig, GossipsubEvent, GossipsubRpc},
    tokio_io::{AsyncRead, AsyncWrite},
    NetworkBehaviour, PeerId,
};

/// Builds the network behaviour for the libp2p Swarm.
/// Implements gossipsub message routing.
#[derive(NetworkBehaviour)]
pub struct Behaviour<TSubstream: AsyncRead + AsyncWrite> {
    gossipsub: Gossipsub<TSubstream>,
    // TODO: Add Kademlia for peer discovery
    /// The events generated by this behaviour to be consumed in the swarm poll.
    // We use gossipsub events for now, generalise later.
    #[behaviour(ignore)]
    events: Vec<GossipsubEvent>,
}

// Implement the NetworkBehaviourEventProcess trait so that we can derive NetworkBehaviour for Behaviour
impl<TSubstream: AsyncRead + AsyncWrite> NetworkBehaviourEventProcess<GossipsubEvent>
    for Behaviour<TSubstream>
{
    fn inject_event(&mut self, event: GossipsubEvent) {
        self.events.push(event);
    }
}

impl<TSubstream: AsyncRead + AsyncWrite> Behaviour<TSubstream> {
    pub fn new(local_peer_id: PeerId, gs_config: GossipsubConfig) -> Self {
        Behaviour {
            gossipsub: Gossipsub::new(local_peer_id, gs_config),
            events: Vec::new(),
        }
    }

    /// Consume the events list when polled.
    fn poll(&mut self) -> Async<NetworkBehaviourAction<GossipsubRpc, GossipsubEvent>> {
        if !self.events.is_empty() {
            return Async::Ready(NetworkBehaviourAction::GenerateEvent(self.events.remove(0)));
        }

        Async::NotReady
    }
}