karyon_p2p/discovery/kademlia/
lookup.rs

1use std::{sync::Arc, time::Duration};
2
3use futures_util::stream::{FuturesUnordered, StreamExt};
4use log::{error, trace};
5use rand::{rngs::OsRng, seq::IndexedRandom, TryRngCore};
6
7use karyon_core::{async_runtime::Executor, async_util::timeout, crypto::KeyPair};
8
9use karyon_net::Endpoint;
10
11use crate::{
12    access_control::{Action, Subject},
13    bloom::BloomRef,
14    connector::Connector,
15    discovery::kademlia::{
16        messages::{
17            FindPeerMsg, KadNetCmd, KadNetMsg, KadNetMsgCodec, PeerMsg, PeersMsg, PingMsg, PongMsg,
18        },
19        routing_table::RoutingTable,
20        SUPPORTED_LOOKUP_PROTOCOLS,
21    },
22    listener::Listener,
23    message::{pick_endpoint, PeerAddr, Protocol, ShutdownMsg},
24    monitor::{ConnectionKind, DiscoveryKind, Monitor},
25    peer::ConnDirection,
26    slots::ConnectionSlots,
27    util::decode,
28    version::version_match,
29    Config, Error, PeerID, Result,
30};
31
32/// Framed lookup-plane connection.
33type KadConnRef = karyon_net::FramedConn<KadNetMsgCodec>;
34
35/// Maximum number of peers that can be returned in a PeersMsg.
36pub const MAX_PEERS_IN_PEERSMSG: usize = 10;
37
38/// Maximum data-plane addresses a peer may advertise about itself.
39/// Legitimate peers only need one per transport (tcp/tls/quic).
40pub const MAX_ADDRS_PER_PEER: usize = 4;
41
42/// Maximum discovery addresses a peer may advertise about itself.
43/// Legitimate peers only need lookup + refresh.
44pub const MAX_DISCOVERY_ADDRS_PER_PEER: usize = 3;
45
46/// Endpoints the lookup service advertises and binds to.
47pub struct LookupEndpoints {
48    /// Data-plane listen addrs (advertised in PeerMsg.addrs).
49    pub listen: Vec<Endpoint>,
50    /// Local bind for the lookup listener (also advertised).
51    pub lookup: Option<Endpoint>,
52    /// UDP refresh addr to advertise (if any).
53    pub refresh: Option<Endpoint>,
54}
55
56pub struct LookupService {
57    /// Peer's ID
58    id: PeerID,
59
60    /// Routing Table
61    table: Arc<RoutingTable>,
62
63    /// Listener
64    listener: Arc<Listener<KadNetMsgCodec>>,
65    /// Connector
66    connector: Arc<Connector<KadNetMsgCodec>>,
67
68    /// Outbound slots.
69    outbound_slots: Arc<ConnectionSlots>,
70
71    /// Endpoints this service advertises and binds to.
72    endpoints: LookupEndpoints,
73
74    /// Holds the configuration for the P2P network.
75    config: Arc<Config>,
76
77    /// Responsible for network and system monitoring.
78    monitor: Arc<Monitor>,
79
80    /// Shared local bloom. Snapshotted on every outgoing PeerMsg.
81    bloom: BloomRef,
82}
83
84impl LookupService {
85    /// Creates a new lookup service.
86    pub fn new(
87        key_pair: &KeyPair,
88        table: Arc<RoutingTable>,
89        config: Arc<Config>,
90        monitor: Arc<Monitor>,
91        bloom: BloomRef,
92        endpoints: LookupEndpoints,
93        ex: Executor,
94    ) -> Self {
95        let inbound_slots = Arc::new(ConnectionSlots::new(config.lookup_inbound_slots));
96        let outbound_slots = Arc::new(ConnectionSlots::new(config.lookup_outbound_slots));
97
98        let listener = Listener::new(
99            key_pair,
100            inbound_slots.clone(),
101            monitor.clone(),
102            config.handshake_timeout,
103            config.access_control.clone(),
104            ex.clone(),
105        );
106
107        let connector = Connector::new(
108            key_pair,
109            config.lookup_connect_retries,
110            outbound_slots.clone(),
111            monitor.clone(),
112            ex,
113        );
114
115        let id = key_pair
116            .public()
117            .try_into()
118            .expect("Get PeerID from KeyPair");
119        Self {
120            id,
121            table,
122            listener,
123            connector,
124            outbound_slots,
125            endpoints,
126            config,
127            monitor,
128            bloom,
129        }
130    }
131
132    /// Start the lookup service.
133    pub async fn start(self: &Arc<Self>) -> Result<()> {
134        self.start_listener().await?;
135        Ok(())
136    }
137
138    pub fn lookup_endpoint(&self) -> Option<&Endpoint> {
139        self.endpoints.lookup.as_ref()
140    }
141
142    /// Shuts down the lookup service.
143    pub async fn shutdown(&self) {
144        self.connector.shutdown().await;
145        self.listener.shutdown().await;
146    }
147
148    /// Starts iterative lookup and populate the routing table.
149    ///
150    /// This method begins by generating a random peer ID and connecting to the
151    /// provided endpoint. It then sends a FindPeer message containing the
152    /// randomly generated peer ID. Upon receiving peers from the initial lookup,
153    /// it starts connecting to these received peers and sends them a FindPeer
154    /// message that contains our own peer ID.
155    pub async fn start_lookup(&self, endpoint: &Endpoint, peer_id: Option<PeerID>) -> Result<()> {
156        trace!("Lookup started {endpoint}");
157        self.monitor
158            .notify(DiscoveryKind::LookupStarted(endpoint.clone()))
159            .await;
160
161        let mut random_peers = vec![];
162        if let Err(err) = self
163            .random_lookup(endpoint, peer_id, &mut random_peers)
164            .await
165        {
166            self.monitor
167                .notify(DiscoveryKind::LookupFailed(endpoint.clone()))
168                .await;
169            return Err(err);
170        };
171
172        let mut peer_buffer = vec![];
173        if let Err(err) = self.self_lookup(&random_peers, &mut peer_buffer).await {
174            self.monitor
175                .notify(DiscoveryKind::LookupFailed(endpoint.clone()))
176                .await;
177            return Err(err);
178        }
179
180        while peer_buffer.len() < MAX_PEERS_IN_PEERSMSG {
181            match random_peers.pop() {
182                Some(p) => peer_buffer.push(p),
183                None => break,
184            }
185        }
186
187        for peer in peer_buffer.iter() {
188            let result = self.table.add_entry(peer.clone().into());
189            trace!("Add entry {result:?}");
190        }
191
192        self.monitor
193            .notify(DiscoveryKind::LookupSucceeded(
194                endpoint.clone(),
195                peer_buffer.len(),
196            ))
197            .await;
198
199        Ok(())
200    }
201
202    /// Starts a random lookup
203    ///
204    /// This will perfom lookup on a random generated PeerID
205    async fn random_lookup(
206        &self,
207        endpoint: &Endpoint,
208        peer_id: Option<PeerID>,
209        random_peers: &mut Vec<PeerMsg>,
210    ) -> Result<()> {
211        for _ in 0..2 {
212            let random_peer_id = PeerID::random()?;
213            let peers = self
214                .connect(endpoint.clone(), peer_id.clone(), &random_peer_id)
215                .await?;
216
217            for peer in peers {
218                if random_peers.contains(&peer)
219                    || peer.peer_id == self.id
220                    || self.table.contains_key(&peer.peer_id.0)
221                {
222                    continue;
223                }
224
225                random_peers.push(peer);
226            }
227        }
228
229        Ok(())
230    }
231
232    /// Starts a self lookup
233    async fn self_lookup(
234        &self,
235        random_peers: &[PeerMsg],
236        peer_buffer: &mut Vec<PeerMsg>,
237    ) -> Result<()> {
238        let mut results = FuturesUnordered::new();
239        for peer in random_peers.choose_multiple(&mut rand::rng(), random_peers.len()) {
240            let endpoint = match pick_endpoint(&peer.discovery_addrs, SUPPORTED_LOOKUP_PROTOCOLS) {
241                Some(ep) => ep,
242                None => continue,
243            };
244            results.push(self.connect(endpoint, Some(peer.peer_id.clone()), &self.id))
245        }
246
247        while let Some(result) = results.next().await {
248            match result {
249                Ok(peers) => peer_buffer.extend(peers),
250                Err(err) => {
251                    error!("Failed to do self lookup: {err}");
252                }
253            }
254        }
255
256        Ok(())
257    }
258
259    /// Connects to the given endpoint and initiates a lookup process for the
260    /// provided peer ID.
261    async fn connect(
262        &self,
263        endpoint: Endpoint,
264        peer_id: Option<PeerID>,
265        target_peer_id: &PeerID,
266    ) -> Result<Vec<PeerMsg>> {
267        if !self.config.access_control.allow(
268            &Subject::Endpoint(&endpoint),
269            Action::Lookup(ConnDirection::Outbound),
270        ) {
271            return Err(Error::AccessDenied);
272        }
273
274        let conn = self.connector.connect(&endpoint, &peer_id).await?;
275        let result = self.handle_outbound(conn, target_peer_id).await;
276
277        self.monitor
278            .notify(ConnectionKind::Disconnected(endpoint))
279            .await;
280        self.outbound_slots.remove().await;
281
282        result
283    }
284
285    /// Handles outbound connection
286    async fn handle_outbound(
287        &self,
288        mut conn: KadConnRef,
289        target_peer_id: &PeerID,
290    ) -> Result<Vec<PeerMsg>> {
291        trace!("Send Ping msg");
292        let mut peers;
293
294        let ping_msg = self.send_ping_msg(&mut conn).await?;
295
296        loop {
297            let t = Duration::from_secs(self.config.lookup_response_timeout);
298            let msg: KadNetMsg = timeout(t, conn.recv_msg()).await??;
299            match msg.header.command {
300                KadNetCmd::Pong => {
301                    let (pong_msg, _) = decode::<PongMsg>(&msg.payload)?;
302                    if ping_msg.nonce != pong_msg.0 {
303                        return Err(Error::InvalidPongMsg);
304                    }
305                    trace!("Send FindPeer msg");
306                    self.send_findpeer_msg(&mut conn, target_peer_id).await?;
307                }
308                KadNetCmd::Peers => {
309                    peers = decode::<PeersMsg>(&msg.payload)?.0.peers;
310                    if peers.len() > MAX_PEERS_IN_PEERSMSG {
311                        return Err(Error::Lookup(
312                            "Received too many peers in PeersMsg".to_string(),
313                        ));
314                    }
315                    for p in &mut peers {
316                        validate_peer_msg(p)?;
317                    }
318                    break;
319                }
320                c => return Err(Error::InvalidMsg(format!("Unexpected msg: {c:?}"))),
321            };
322        }
323
324        trace!("Send Peer msg");
325        self.send_peer_msg(&mut conn).await?;
326
327        trace!("Send Shutdown msg");
328        self.send_shutdown_msg(&mut conn).await?;
329
330        Ok(peers)
331    }
332
333    /// Start a listener.
334    async fn start_listener(self: &Arc<Self>) -> Result<()> {
335        let endpoint = match self.lookup_endpoint() {
336            Some(e) => e.clone(),
337            None => return Ok(()),
338        };
339
340        if !endpoint.is_tcp() {
341            return Err(Error::Config(format!(
342                "lookup endpoint must be tcp://..., got {endpoint}"
343            )));
344        }
345
346        let callback = {
347            let this = self.clone();
348            |conn: KadConnRef| async move {
349                let t = Duration::from_secs(this.config.lookup_connection_lifespan);
350                timeout(t, this.handle_inbound(conn)).await??;
351                Ok(())
352            }
353        };
354
355        self.listener
356            .start_with_callback(endpoint, callback)
357            .await?;
358        Ok(())
359    }
360
361    /// Handles inbound connection
362    async fn handle_inbound(self: &Arc<Self>, mut conn: KadConnRef) -> Result<()> {
363        loop {
364            let msg: KadNetMsg = conn.recv_msg().await?;
365            trace!("Receive msg {:?}", msg.header.command);
366
367            if let KadNetCmd::Shutdown = msg.header.command {
368                return Ok(());
369            }
370
371            match &msg.header.command {
372                KadNetCmd::Ping => {
373                    let (ping_msg, _) = decode::<PingMsg>(&msg.payload)?;
374                    if !version_match(&self.config.version.req, &ping_msg.version) {
375                        return Err(Error::IncompatibleVersion("system: {}".into()));
376                    }
377                    self.send_pong_msg(ping_msg.nonce, &mut conn).await?;
378                }
379                KadNetCmd::FindPeer => {
380                    let (findpeer_msg, _) = decode::<FindPeerMsg>(&msg.payload)?;
381                    let peer_id = findpeer_msg.0;
382                    self.send_peers_msg(&peer_id, &mut conn).await?;
383                }
384                KadNetCmd::Peer => {
385                    let (mut peer, _) = decode::<PeerMsg>(&msg.payload)?;
386                    validate_peer_msg(&mut peer)?;
387                    let result = self.table.add_entry(peer.into());
388                    trace!("Add entry result: {result:?}");
389                }
390                c => return Err(Error::InvalidMsg(format!("Unexpected msg: {c:?}"))),
391            }
392        }
393    }
394
395    /// Sends a Ping msg.
396    async fn send_ping_msg(&self, conn: &mut KadConnRef) -> Result<PingMsg> {
397        trace!("Send Pong msg");
398        let mut nonce: [u8; 32] = [0; 32];
399        OsRng.try_fill_bytes(&mut nonce)?;
400
401        let ping_msg = PingMsg {
402            version: self.config.version.v.clone(),
403            nonce,
404        };
405        conn.send_msg(KadNetMsg::new(KadNetCmd::Ping, &ping_msg)?)
406            .await?;
407        Ok(ping_msg)
408    }
409
410    /// Sends a Pong msg
411    async fn send_pong_msg(&self, nonce: [u8; 32], conn: &mut KadConnRef) -> Result<()> {
412        trace!("Send Pong msg");
413        conn.send_msg(KadNetMsg::new(KadNetCmd::Pong, PongMsg(nonce))?)
414            .await?;
415        Ok(())
416    }
417
418    /// Sends a FindPeer msg
419    async fn send_findpeer_msg(&self, conn: &mut KadConnRef, peer_id: &PeerID) -> Result<()> {
420        trace!("Send FindPeer msg");
421        conn.send_msg(KadNetMsg::new(
422            KadNetCmd::FindPeer,
423            FindPeerMsg(peer_id.clone()),
424        )?)
425        .await?;
426        Ok(())
427    }
428
429    /// Sends a Peers msg.
430    async fn send_peers_msg(&self, peer_id: &PeerID, conn: &mut KadConnRef) -> Result<()> {
431        trace!("Send Peers msg");
432        let entries = self
433            .table
434            .closest_entries(&peer_id.0, MAX_PEERS_IN_PEERSMSG);
435
436        let peers: Vec<PeerMsg> = entries.into_iter().map(|e| e.into()).collect();
437        conn.send_msg(KadNetMsg::new(KadNetCmd::Peers, PeersMsg { peers })?)
438            .await?;
439        Ok(())
440    }
441
442    /// Sends a Peer msg advertising our listen and discovery addresses.
443    /// `addrs` carries every data-plane listen endpoint; `discovery_addrs`
444    /// carries the lookup endpoint and (when set) the udp refresh endpoint.
445    async fn send_peer_msg(&self, conn: &mut KadConnRef) -> Result<()> {
446        trace!("Send Peer msg");
447
448        let mut addrs = Vec::new();
449        for ep in &self.endpoints.listen {
450            if let Some(pa) = PeerAddr::from_endpoint(ep, 0) {
451                addrs.push(pa);
452            }
453        }
454
455        let mut discovery_addrs = Vec::new();
456        if let Some(ep) = self.endpoints.lookup.as_ref() {
457            if let Some(pa) = PeerAddr::from_endpoint(ep, 0) {
458                discovery_addrs.push(pa);
459            }
460        }
461        if let Some(ep) = self.endpoints.refresh.as_ref() {
462            if let Some(pa) = PeerAddr::from_endpoint(ep, 0) {
463                discovery_addrs.push(pa);
464            }
465        }
466
467        let peer_msg = PeerMsg {
468            peer_id: self.id.clone(),
469            addrs,
470            discovery_addrs,
471            protocols: *self.bloom.read(),
472        };
473        conn.send_msg(KadNetMsg::new(KadNetCmd::Peer, &peer_msg)?)
474            .await?;
475        Ok(())
476    }
477
478    /// Sends a Shutdown msg.
479    async fn send_shutdown_msg(&self, conn: &mut KadConnRef) -> Result<()> {
480        trace!("Send Shutdown msg");
481        conn.send_msg(KadNetMsg::new(KadNetCmd::Shutdown, ShutdownMsg(0))?)
482            .await?;
483        Ok(())
484    }
485}
486
487/// Reject PeerMsgs that advertise more addresses than a legitimate
488/// peer would. Caps memory blow-up from malicious or buggy peers.
489fn validate_peer_msg(p: &mut PeerMsg) -> Result<()> {
490    if p.addrs.len() > MAX_ADDRS_PER_PEER {
491        return Err(Error::InvalidMsg(format!(
492            "PeerMsg.addrs has {} entries, max {MAX_ADDRS_PER_PEER}",
493            p.addrs.len()
494        )));
495    }
496    p.discovery_addrs
497        .retain(|a| !matches!(a.protocol, Protocol::Tls));
498    if p.discovery_addrs.len() > MAX_DISCOVERY_ADDRS_PER_PEER {
499        return Err(Error::InvalidMsg(format!(
500            "PeerMsg.discovery_addrs has {} entries, max {MAX_DISCOVERY_ADDRS_PER_PEER}",
501            p.discovery_addrs.len()
502        )));
503    }
504    Ok(())
505}
karyon_p2p/discovery/kademlia/lookup.rs

karyon_p2p/discovery/kademlia/
lookup.rs
