Skip to main content

karyon_p2p/
node.rs

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