p2p_chat/net/
discovery.rs

1//! This module defines the peer discovery mechanism for the network.
2//!
3//! It combines mDNS for local peer discovery and Kademlia for decentralized
4//! peer discovery in the wider network.
5use anyhow::Result;
6use libp2p::{kad, mdns, PeerId};
7
8/// The `libp2p` network behaviour for peer discovery.
9#[derive(libp2p::swarm::NetworkBehaviour)]
10pub struct DiscoveryBehaviour {
11    /// The mDNS behaviour for local peer discovery.
12    pub mdns: mdns::tokio::Behaviour,
13    /// The Kademlia behaviour for decentralized peer discovery.
14    pub kademlia: kad::Behaviour<kad::store::MemoryStore>,
15}
16
17impl DiscoveryBehaviour {
18    /// Creates a new `DiscoveryBehaviour`.
19    ///
20    /// # Arguments
21    ///
22    /// * `local_peer_id` - The `PeerId` of the local node.
23    ///
24    /// # Errors
25    ///
26    /// This function will return an error if the mDNS behaviour cannot be created.
27    pub fn new(local_peer_id: PeerId) -> Result<Self> {
28        // Initialize mDNS for local discovery.
29        let mdns = mdns::tokio::Behaviour::new(mdns::Config::default(), local_peer_id)?;
30
31        // Initialize Kademlia DHT.
32        let store = kad::store::MemoryStore::new(local_peer_id);
33        let mut kademlia = kad::Behaviour::new(local_peer_id, store);
34
35        // Set Kademlia to server mode to participate in the DHT.
36        kademlia.set_mode(Some(kad::Mode::Server));
37
38        Ok(Self { mdns, kademlia })
39    }
40
41    /// Bootstraps the Kademlia DHT.
42    ///
43    /// This will start the process of finding other peers in the network.
44    ///
45    /// # Errors
46    ///
47    /// This function will return an error if the bootstrap process fails.
48    pub fn bootstrap(&mut self) -> Result<()> {
49        if let Err(e) = self.kademlia.bootstrap() {
50            tracing::warn!("Failed to bootstrap Kademlia: {}", e);
51        }
52        Ok(())
53    }
54
55    /// Starts providing a key in the Kademlia DHT.
56    ///
57    /// This announces to the network that the local node can provide information
58    /// about the given key.
59    ///
60    /// # Arguments
61    ///
62    /// * `key` - The key to start providing.
63    ///
64    /// # Errors
65    ///
66    /// This function will return an error if the providing process fails to start.
67    pub fn start_providing(&mut self, key: kad::RecordKey) -> Result<()> {
68        self.kademlia.start_providing(key)?;
69        Ok(())
70    }
71
72    /// Gets the providers for a given key from the Kademlia DHT.
73    ///
74    /// # Arguments
75    ///
76    /// * `key` - The key to get providers for.
77    ///
78    /// # Returns
79    ///
80    /// A `QueryId` for the get providers query.
81    pub fn get_providers(&mut self, key: kad::RecordKey) -> kad::QueryId {
82        self.kademlia.get_providers(key)
83    }
84
85    /// Adds a known address for a peer to the Kademlia DHT.
86    ///
87    /// # Arguments
88    ///
89    /// * `peer_id` - The `PeerId` of the peer.
90    /// * `address` - The `Multiaddr` of the peer.
91    pub fn add_peer_address(&mut self, peer_id: PeerId, address: libp2p::Multiaddr) {
92        self.kademlia.add_address(&peer_id, address);
93    }
94}