libp2p_peer_store/
behaviour.rs

1use std::{collections::VecDeque, task::Poll};
2
3use libp2p_core::{Multiaddr, PeerId};
4use libp2p_swarm::{dummy, NetworkBehaviour};
5
6use crate::store::Store;
7
8/// Events generated by [`Behaviour`] and emitted back to [`Swarm`](libp2p_swarm::Swarm).
9#[derive(Debug, Clone)]
10pub enum Event<T> {
11    /// The peer's record has been updated.  
12    /// Manually updating a record will always emit this event
13    /// even if it provides no new information.
14    RecordUpdated {
15        /// The peer that has an update.
16        peer: PeerId,
17    },
18    /// Event from the internal store.
19    Store(T),
20}
21
22/// Behaviour that maintains a peer address book.
23///
24/// Usage:
25/// ```
26/// use libp2p::swarm::NetworkBehaviour;
27/// use libp2p_peer_store::{memory_store::MemoryStore, Behaviour};
28///
29/// // `identify::Behaviour` broadcasts listen addresses of the peer,
30/// // `peer_store::Behaviour` will then capture the resulting
31/// // `FromSwarm::NewExternalAddrOfPeer` and add the addresses
32/// // to address book.
33/// #[derive(NetworkBehaviour)]
34/// struct ComposedBehaviour {
35///     peer_store: Behaviour<MemoryStore>,
36///     identify: libp2p::identify::Behaviour,
37/// }
38/// ```
39pub struct Behaviour<S: Store> {
40    /// The internal store.
41    store: S,
42    /// Pending Events to be emitted back to [`Swarm`](libp2p_swarm::Swarm).
43    pending_events: VecDeque<Event<S::FromStore>>,
44}
45
46impl<'a, S> Behaviour<S>
47where
48    S: Store + 'static,
49{
50    /// Build a new [`Behaviour`] with the given store.
51    pub fn new(store: S) -> Self {
52        Self {
53            store,
54            pending_events: VecDeque::new(),
55        }
56    }
57
58    /// Try to get all observed address of the given peer.  
59    /// Returns `None` when the peer is not in the store.
60    pub fn address_of_peer<'b>(
61        &'a self,
62        peer: &'b PeerId,
63    ) -> Option<impl Iterator<Item = &'a Multiaddr> + use<'a, 'b, S>> {
64        self.store.addresses_of_peer(peer)
65    }
66
67    /// Get an immutable reference to the internal store.
68    pub fn store(&self) -> &S {
69        &self.store
70    }
71
72    /// Get a mutable reference to the internal store.
73    pub fn store_mut(&mut self) -> &mut S {
74        &mut self.store
75    }
76
77    fn handle_store_event(&mut self, event: crate::store::Event<<S as Store>::FromStore>) {
78        use crate::store::Event::*;
79        match event {
80            RecordUpdated(peer) => self.pending_events.push_back(Event::RecordUpdated { peer }),
81            Store(ev) => self.pending_events.push_back(Event::Store(ev)),
82        }
83    }
84}
85
86impl<S> NetworkBehaviour for Behaviour<S>
87where
88    S: Store + 'static,
89    <S as Store>::FromStore: Send + Sync,
90{
91    type ConnectionHandler = dummy::ConnectionHandler;
92
93    type ToSwarm = Event<S::FromStore>;
94
95    fn handle_established_inbound_connection(
96        &mut self,
97        _connection_id: libp2p_swarm::ConnectionId,
98        _peer: libp2p_core::PeerId,
99        _local_addr: &libp2p_core::Multiaddr,
100        _remote_addr: &libp2p_core::Multiaddr,
101    ) -> Result<libp2p_swarm::THandler<Self>, libp2p_swarm::ConnectionDenied> {
102        Ok(dummy::ConnectionHandler)
103    }
104
105    fn handle_pending_outbound_connection(
106        &mut self,
107        _connection_id: libp2p_swarm::ConnectionId,
108        maybe_peer: Option<PeerId>,
109        _addresses: &[Multiaddr],
110        _effective_role: libp2p_core::Endpoint,
111    ) -> Result<Vec<Multiaddr>, libp2p_swarm::ConnectionDenied> {
112        if maybe_peer.is_none() {
113            return Ok(Vec::new());
114        }
115        let peer = maybe_peer.expect("already handled");
116        Ok(self
117            .store
118            .addresses_of_peer(&peer)
119            .map(|i| i.cloned().collect())
120            .unwrap_or_default())
121    }
122
123    fn handle_established_outbound_connection(
124        &mut self,
125        _connection_id: libp2p_swarm::ConnectionId,
126        _peer: libp2p_core::PeerId,
127        _addr: &libp2p_core::Multiaddr,
128        _role_override: libp2p_core::Endpoint,
129        _port_use: libp2p_core::transport::PortUse,
130    ) -> Result<libp2p_swarm::THandler<Self>, libp2p_swarm::ConnectionDenied> {
131        Ok(dummy::ConnectionHandler)
132    }
133
134    fn on_swarm_event(&mut self, event: libp2p_swarm::FromSwarm) {
135        self.store.on_swarm_event(&event);
136    }
137
138    fn on_connection_handler_event(
139        &mut self,
140        _peer_id: libp2p_core::PeerId,
141        _connection_id: libp2p_swarm::ConnectionId,
142        _event: libp2p_swarm::THandlerOutEvent<Self>,
143    ) {
144        unreachable!("No event will be produced by a dummy handler.")
145    }
146
147    fn poll(
148        &mut self,
149        cx: &mut std::task::Context<'_>,
150    ) -> std::task::Poll<libp2p_swarm::ToSwarm<Self::ToSwarm, libp2p_swarm::THandlerInEvent<Self>>>
151    {
152        if let Some(ev) = self.store.poll(cx) {
153            self.handle_store_event(ev);
154        };
155
156        if let Some(ev) = self.pending_events.pop_front() {
157            return Poll::Ready(libp2p_swarm::ToSwarm::GenerateEvent(ev));
158        }
159        Poll::Pending
160    }
161}