Skip to main content

karyon_p2p/
node.rs

1use std::{
2    collections::{HashMap, HashSet},
3    sync::Arc,
4};
5
6use log::{debug, info};
7use parking_lot::RwLock as SyncRwLock;
8
9use karyon_core::{
10    async_runtime::Executor,
11    async_util::{TaskGroup, TaskResult},
12    crypto::KeyPair,
13};
14use karyon_eventemitter::EventListener;
15use karyon_net::Endpoint;
16
17use crate::{
18    access_control::{Action, Subject},
19    bloom::{Bloom, BloomRef},
20    codec::PeerNetMsgCodec,
21    config::Config,
22    conn_queue::ConnQueue,
23    connector::Connector,
24    discovery::{kademlia::KademliaDiscovery, DiscoveredPeer, Discovery, PeerConnectionEvent},
25    listener::Listener,
26    message::{pick_endpoint, Protocol},
27    monitor::{Monitor, PoolEvent},
28    peer::ConnDirection,
29    peer_pool::{PeerEvent, PeerEventTopic, PeerPool},
30    protocol::{PeerConn, Protocol as ProtocolTrait, ProtocolID, ProtocolKind},
31    protocols::PingProtocol,
32    slots::ConnectionSlots,
33    PeerID, Result,
34};
35
36/// Central entry point for the p2p network.
37///
38/// Manages peer connections, discovery, and protocol registration.
39///
40/// # Example
41///
42/// ```no_run
43/// use karyon_core::async_runtime::global_executor;
44/// use karyon_p2p::{Node, Config, keypair::{KeyPair, KeyPairType}};
45///
46/// let key_pair = KeyPair::generate(&KeyPairType::Ed25519);
47/// let config = Config {
48///     listen_endpoints: vec![
49///         "tcp://0.0.0.0:8000".parse().unwrap(),
50///     ],
51///     ..Config::default()
52/// };
53///
54/// let node = Node::new(&key_pair, config, global_executor());
55///
56/// // node.run().await.unwrap();
57/// // node.attach_protocol::<MyProto>(|peer| MyProto::new(peer)).await;
58/// // node.shutdown().await;
59/// ```
60pub struct Node {
61    /// The Configuration for the P2P network.
62    config: Arc<Config>,
63
64    /// Identity Key pair
65    key_pair: KeyPair,
66
67    /// Peer ID
68    peer_id: PeerID,
69
70    /// Responsible for network and system monitoring.
71    monitor: Arc<Monitor>,
72
73    /// Discovery instance.
74    discovery: Arc<dyn Discovery>,
75
76    /// PeerPool instance.
77    peer_pool: Arc<PeerPool>,
78
79    /// Connector for outbound connections.
80    connector: Arc<Connector<PeerNetMsgCodec>>,
81
82    /// Listener for inbound connections.
83    listener: Arc<Listener<PeerNetMsgCodec>>,
84
85    /// Managing spawned tasks.
86    task_group: TaskGroup,
87
88    /// Local bloom advertising what items (protocols, swarm keys, ...)
89    /// this node supports. The mandatory side is filtered with `covers`,
90    /// the optional side with `intersects`. Updated by `attach_protocol`
91    /// (via `Protocol::kind()`) and by application layers like Swarm.
92    bloom: BloomRef,
93}
94
95impl Node {
96    /// Creates a new Node with the default Kademlia discovery.
97    pub fn new(key_pair: &KeyPair, config: Config, ex: Executor) -> Arc<Node> {
98        let config = Arc::new(config);
99        let monitor = Arc::new(Monitor::new(config.clone()));
100        let peer_id = PeerID::try_from(key_pair.public())
101            .expect("Derive a peer id from the provided key pair.");
102        info!("PeerID: {peer_id}");
103
104        let conn_queue = ConnQueue::new();
105        let peer_pool = PeerPool::new(
106            &peer_id,
107            conn_queue.clone(),
108            config.clone(),
109            monitor.clone(),
110            ex.clone(),
111        );
112
113        let bloom: BloomRef = Arc::new(SyncRwLock::new(Bloom::empty()));
114
115        let discovery: Arc<dyn Discovery> = KademliaDiscovery::new(
116            key_pair,
117            &peer_id,
118            config.clone(),
119            monitor.clone(),
120            bloom.clone(),
121            ex.clone(),
122        );
123
124        let outbound_slots = Arc::new(ConnectionSlots::new(config.outbound_slots));
125        let connector = Connector::new_with_queue(
126            key_pair,
127            config.max_connect_retries,
128            outbound_slots,
129            conn_queue.clone(),
130            monitor.clone(),
131            ex.clone(),
132        );
133
134        let inbound_slots = Arc::new(ConnectionSlots::new(config.inbound_slots));
135        let listener = Listener::new_with_queue(
136            key_pair,
137            inbound_slots,
138            conn_queue,
139            monitor.clone(),
140            config.handshake_timeout,
141            config.access_control.clone(),
142            ex.clone(),
143        );
144
145        let task_group = TaskGroup::with_executor(ex);
146
147        Arc::new(Self {
148            key_pair: key_pair.clone(),
149            peer_id,
150            monitor,
151            discovery,
152            config,
153            peer_pool,
154            connector,
155            listener,
156            task_group,
157            bloom,
158        })
159    }
160
161    /// Creates a new Node with a custom discovery implementation.
162    /// The caller is responsible for wiring its own bloom_provider into
163    /// the discovery; Node's `bloom_add_*` methods will not affect
164    /// it unless the discovery reads from the same source.
165    pub fn with_discovery(
166        key_pair: &KeyPair,
167        config: Config,
168        discovery: Arc<dyn Discovery>,
169        ex: Executor,
170    ) -> Arc<Node> {
171        let config = Arc::new(config);
172        let monitor = Arc::new(Monitor::new(config.clone()));
173        let peer_id = PeerID::try_from(key_pair.public())
174            .expect("Derive a peer id from the provided key pair.");
175        info!("PeerID: {peer_id}");
176
177        let conn_queue = ConnQueue::new();
178        let peer_pool = PeerPool::new(
179            &peer_id,
180            conn_queue.clone(),
181            config.clone(),
182            monitor.clone(),
183            ex.clone(),
184        );
185
186        let bloom: BloomRef = Arc::new(SyncRwLock::new(Bloom::empty()));
187
188        let outbound_slots = Arc::new(ConnectionSlots::new(config.outbound_slots));
189        let connector = Connector::new_with_queue(
190            key_pair,
191            config.max_connect_retries,
192            outbound_slots,
193            conn_queue.clone(),
194            monitor.clone(),
195            ex.clone(),
196        );
197
198        let inbound_slots = Arc::new(ConnectionSlots::new(config.inbound_slots));
199        let listener = Listener::new_with_queue(
200            key_pair,
201            inbound_slots,
202            conn_queue,
203            monitor.clone(),
204            config.handshake_timeout,
205            config.access_control.clone(),
206            ex.clone(),
207        );
208
209        let task_group = TaskGroup::with_executor(ex);
210
211        Arc::new(Self {
212            key_pair: key_pair.clone(),
213            peer_id,
214            monitor,
215            discovery,
216            config,
217            peer_pool,
218            connector,
219            listener,
220            task_group,
221            bloom,
222        })
223    }
224
225    /// Run the Node, starting listeners, PeerPool, and Discovery.
226    pub async fn run(self: &Arc<Self>) -> Result<()> {
227        // Core protocols (PING) are attached before the pool starts so
228        // they're advertised on every handshake.
229        self.attach_core_protocols().await?;
230
231        self.peer_pool.start().await?;
232
233        // Start data listeners.
234        for endpoint in &self.config.listen_endpoints {
235            let resolved = self.listener.start(endpoint.clone()).await?;
236            info!("Listening on {resolved}");
237        }
238
239        // Start discovery.
240        self.discovery.clone().start().await?;
241
242        // Forward peer lifecycle events to discovery.
243        self.task_group.spawn_then(
244            {
245                let this = self.clone();
246                async move { this.forward_peer_events().await }
247            },
248            |res: TaskResult<()>| async move {
249                debug!("forward_peer_events task ended: {res}");
250            },
251        );
252
253        // Spawn task to connect discovered peers.
254        self.task_group.spawn_then(
255            {
256                let this = self.clone();
257                async move { this.connect_discovered_peers().await }
258            },
259            |res: TaskResult<()>| async move {
260                debug!("connect_discovered_peers task ended: {res}");
261            },
262        );
263
264        Ok(())
265    }
266
267    /// Forward peer-pool lifecycle events to discovery so it can
268    /// update its routing state. Each registered listener (Node's
269    /// here, plus any external Swarm subscriber) gets every event
270    /// independently - no MPMC stealing.
271    async fn forward_peer_events(self: Arc<Self>) {
272        let listener = self.peer_pool.register_peer_events();
273        while let Ok(event) = listener.recv().await {
274            let mapped = match event {
275                PeerEvent::Added(pid) => PeerConnectionEvent::Connected(pid),
276                PeerEvent::Removed(pid) => PeerConnectionEvent::Disconnected(pid),
277                PeerEvent::HandshakeFailed(pid) => PeerConnectionEvent::ConnectFailed(pid),
278            };
279            self.discovery.on_event(mapped);
280        }
281    }
282
283    /// Consume discovered peers from the discovery service and connect to them.
284    /// Runs forever; the task_group cancels it on Node::shutdown.
285    async fn connect_discovered_peers(self: Arc<Self>) {
286        let supported = [Protocol::Tcp, Protocol::Tls, Protocol::Quic];
287
288        loop {
289            let discovered = self.discovery.recv().await;
290
291            let endpoint = match pick_endpoint(&discovered.addrs, &supported) {
292                Some(ep) => ep,
293                None => continue,
294            };
295
296            if !self.config.access_control.allow(
297                &Subject::Endpoint(&endpoint),
298                Action::Connect(ConnDirection::Outbound),
299            ) {
300                debug!("Skipped dialing denied endpoint {endpoint}");
301                continue;
302            }
303
304            let peer_id = discovered.peer_id.clone();
305
306            if self
307                .connector
308                .connect_and_queue(&endpoint, &peer_id)
309                .await
310                .is_err()
311            {
312                self.monitor
313                    .notify(PoolEvent::ConnectFailed(peer_id.clone(), endpoint))
314                    .await;
315                self.discovery
316                    .on_event(PeerConnectionEvent::ConnectFailed(peer_id));
317            }
318        }
319    }
320
321    /// Attach a custom protocol. karyon runs the constructor closure
322    /// once per connected peer with a typed `PeerConn` scoped to this
323    /// protocol. Bloom advertises the protocol id according to
324    /// `P::kind()`.
325    pub async fn attach_protocol<P: ProtocolTrait>(
326        &self,
327        c: impl Fn(PeerConn) -> Result<Arc<dyn ProtocolTrait>> + Send + Sync + 'static,
328    ) -> Result<()> {
329        self.peer_pool.attach_protocol::<P>(Box::new(c)).await?;
330        let id = P::id();
331        match P::kind() {
332            ProtocolKind::Mandatory => self.bloom_add_mandatory(&id),
333            ProtocolKind::Optional => self.bloom_add_optional(&id),
334        }
335        Ok(())
336    }
337
338    /// Attach the core protocols (PING). Called once during `run`.
339    async fn attach_core_protocols(self: &Arc<Self>) -> Result<()> {
340        self.attach_protocol::<PingProtocol>(|conn| {
341            Ok(PingProtocol::new(conn) as Arc<dyn ProtocolTrait>)
342        })
343        .await
344    }
345
346    /// Add an item the local node REQUIRES peers to also support.
347    /// Reflected in the next bloom snapshot advertised in PeerMsg.
348    pub fn bloom_add_mandatory(&self, item: impl AsRef<[u8]>) {
349        self.bloom.write().add_mandatory(item);
350    }
351
352    /// Add an item the local node would LIKE peers to also support but
353    /// doesn't require. Used by Swarm and other layers for fuzzy
354    /// protocol-aware discovery without rejecting non-matches.
355    pub fn bloom_add_optional(&self, item: impl AsRef<[u8]>) {
356        self.bloom.write().add_optional(item);
357    }
358
359    /// Snapshot of the local bloom (mandatory + optional sides).
360    pub fn bloom_snapshot(&self) -> Bloom {
361        *self.bloom.read()
362    }
363
364    /// Find peers in the routing table whose advertised bloom may
365    /// contain `item`. Useful for swarm-targeted lookups (e.g.
366    /// "peers in this room") without changing handshake semantics.
367    pub fn find_peers_with(&self, item: impl AsRef<[u8]>) -> Vec<DiscoveredPeer> {
368        self.discovery.find_peers_with(item.as_ref())
369    }
370
371    /// Returns the number of currently connected peers.
372    pub async fn peers(&self) -> usize {
373        self.peer_pool.peers_len().await
374    }
375
376    /// Returns the `Config`.
377    pub fn config(&self) -> Arc<Config> {
378        self.config.clone()
379    }
380
381    /// Returns the `PeerID`.
382    pub fn peer_id(&self) -> &PeerID {
383        &self.peer_id
384    }
385
386    /// Returns the `KeyPair`.
387    pub fn key_pair(&self) -> &KeyPair {
388        &self.key_pair
389    }
390
391    /// Returns a map of inbound connected peers with their endpoints.
392    pub async fn inbound_peers(&self) -> HashMap<PeerID, Endpoint> {
393        self.peer_pool.inbound_peers().await
394    }
395
396    /// Returns a map of outbound connected peers with their endpoints.
397    pub async fn outbound_peers(&self) -> HashMap<PeerID, Endpoint> {
398        self.peer_pool.outbound_peers().await
399    }
400
401    /// Returns the monitor to receive system events.
402    pub fn monitor(&self) -> Arc<Monitor> {
403        self.monitor.clone()
404    }
405
406    /// Register a listener for peer lifecycle events. Each call returns
407    /// a fresh listener that receives every event (true broadcast).
408    pub fn register_peer_events(&self) -> EventListener<PeerEventTopic, PeerEvent> {
409        self.peer_pool.register_peer_events()
410    }
411
412    /// Broadcast a message to a specific set of peers on a given protocol.
413    /// Used by Swarm and other layers to scope broadcasts.
414    pub async fn broadcast_to(
415        &self,
416        proto_id: &ProtocolID,
417        msg: Vec<u8>,
418        targets: &HashSet<PeerID>,
419    ) {
420        self.peer_pool.broadcast_to(proto_id, msg, targets).await;
421    }
422
423    /// Send a message to a specific peer on the given protocol.
424    /// Returns `PeerNotFound` if the peer is not currently connected.
425    pub async fn send_to(
426        &self,
427        peer_id: &PeerID,
428        proto_id: &ProtocolID,
429        msg: Vec<u8>,
430    ) -> Result<()> {
431        self.peer_pool.send_to(peer_id, proto_id, msg).await
432    }
433
434    /// Returns the negotiated protocol set for a connected peer, or
435    /// `None` if no peer with that id is currently in the pool.
436    pub async fn peer_protocol_set(&self, pid: &PeerID) -> Option<HashSet<ProtocolID>> {
437        self.peer_pool.peer_protocol_set(pid).await
438    }
439
440    /// Shuts down the Node.
441    pub async fn shutdown(&self) {
442        self.discovery.shutdown().await;
443        self.peer_pool.shutdown().await;
444        self.connector.shutdown().await;
445        self.listener.shutdown().await;
446        self.task_group.cancel().await;
447    }
448}