karyon_p2p/
peer_pool.rs

1use std::{
2    collections::{HashMap, HashSet},
3    sync::Arc,
4};
5
6use log::{error, info, warn};
7
8use karyon_core::{
9    async_runtime::{lock::RwLock, Executor},
10    async_util::{TaskGroup, TaskResult},
11};
12
13use karyon_eventemitter::{EventEmitter, EventListener, EventTopic, EventValue};
14
15use karyon_net::Endpoint;
16
17use crate::{
18    access_control::{Action, PeerCandidate, Subject},
19    config::Config,
20    conn_queue::{ConnQueue, QueuedConn},
21    handshake::{handshake, HandshakeParams},
22    monitor::{Monitor, PoolEvent},
23    peer::{ConnDirection, Peer},
24    protocol::{Protocol, ProtocolConstructor, ProtocolID, ProtocolMeta},
25    Error, PeerID, Result,
26};
27
28/// Topic key for the peer-lifecycle event emitter.
29#[derive(Hash, PartialEq, Eq, Debug, Clone)]
30pub enum PeerEventTopic {
31    Lifecycle,
32}
33
34/// Peer-lifecycle events. Each registered listener receives every
35/// event independently.
36#[derive(Debug, Clone, EventValue)]
37pub enum PeerEvent {
38    /// Peer added after a successful handshake.
39    Added(PeerID),
40    /// Previously-added peer removed.
41    Removed(PeerID),
42    /// Handshake failed before the peer was added.
43    HandshakeFailed(Option<PeerID>),
44}
45
46impl EventTopic for PeerEvent {
47    type Topic = PeerEventTopic;
48    fn topic() -> Self::Topic {
49        PeerEventTopic::Lifecycle
50    }
51}
52
53pub struct PeerPool {
54    /// Peer's ID
55    pub id: PeerID,
56
57    /// Connection queue
58    conn_queue: Arc<ConnQueue>,
59
60    /// Holds the running peers.
61    peers: RwLock<HashMap<PeerID, Arc<Peer>>>,
62
63    /// Hashmap contains protocol constructors.
64    pub(crate) protocols: RwLock<HashMap<ProtocolID, Box<ProtocolConstructor>>>,
65
66    /// Per-protocol metadata (version + kind, extensible). Keyed by
67    /// protocol id. Source of truth for the handshake's mandatory check
68    /// and version negotiation.
69    pub(crate) protocol_meta: RwLock<HashMap<ProtocolID, ProtocolMeta>>,
70
71    /// Peer-lifecycle event emitter. Each registered listener gets its
72    /// own copy of every event.
73    peer_emitter: Arc<EventEmitter<PeerEventTopic>>,
74
75    /// Managing spawned tasks.
76    task_group: TaskGroup,
77
78    /// A global Executor
79    pub(crate) executor: Executor,
80
81    /// The Configuration for the P2P network.
82    pub(crate) config: Arc<Config>,
83
84    /// Responsible for network and system monitoring.
85    monitor: Arc<Monitor>,
86}
87
88impl PeerPool {
89    /// Creates a new PeerPool
90    pub fn new(
91        id: &PeerID,
92        conn_queue: Arc<ConnQueue>,
93        config: Arc<Config>,
94        monitor: Arc<Monitor>,
95        executor: Executor,
96    ) -> Arc<Self> {
97        Arc::new(Self {
98            id: id.clone(),
99            conn_queue,
100            peers: RwLock::new(HashMap::new()),
101            protocols: RwLock::new(HashMap::new()),
102            protocol_meta: RwLock::new(HashMap::new()),
103            peer_emitter: EventEmitter::new(),
104            task_group: TaskGroup::with_executor(executor.clone()),
105            executor,
106            monitor,
107            config,
108        })
109    }
110
111    /// Register a listener for the peer-lifecycle events.
112    pub fn register_peer_events(&self) -> EventListener<PeerEventTopic, PeerEvent> {
113        self.peer_emitter.register(&PeerEventTopic::Lifecycle)
114    }
115
116    /// Starts the [`PeerPool`]
117    pub async fn start(self: &Arc<Self>) -> Result<()> {
118        self.task_group.spawn(self.clone().run());
119        Ok(())
120    }
121
122    /// Shuts down
123    pub async fn shutdown(&self) {
124        for peer in self.peers.read().await.values() {
125            let _ = peer.shutdown().await;
126        }
127
128        self.task_group.cancel().await;
129    }
130
131    /// Register a protocol's user-supplied constructor and metadata.
132    /// Bloom advertising is handled by `Node::attach_protocol`.
133    pub async fn attach_protocol<P: Protocol>(&self, c: Box<ProtocolConstructor>) -> Result<()> {
134        let id = P::id();
135        self.protocols.write().await.insert(id.clone(), c);
136        self.protocol_meta.write().await.insert(
137            id,
138            ProtocolMeta {
139                version: P::version()?,
140                kind: P::kind(),
141            },
142        );
143        Ok(())
144    }
145
146    /// Broadcast a message to all connected peers.
147    pub async fn broadcast(&self, proto_id: &ProtocolID, msg: Vec<u8>) {
148        for (pid, peer) in self.peers.read().await.iter() {
149            if let Err(err) = peer.send(proto_id.to_string(), msg.clone()).await {
150                error!("failed to send msg to {pid}: {err}");
151                continue;
152            }
153        }
154    }
155
156    /// Broadcast a message to a specific set of peers.
157    pub async fn broadcast_to(
158        &self,
159        proto_id: &ProtocolID,
160        msg: Vec<u8>,
161        targets: &HashSet<PeerID>,
162    ) {
163        for (pid, peer) in self.peers.read().await.iter() {
164            if !targets.contains(pid) {
165                continue;
166            }
167            if let Err(err) = peer.send(proto_id.to_string(), msg.clone()).await {
168                error!("failed to send msg to {pid}: {err}");
169            }
170        }
171    }
172
173    /// Send a message to a specific peer on the given protocol. Returns
174    /// `PeerNotFound` if the peer is not currently in the pool.
175    pub async fn send_to(
176        &self,
177        peer_id: &PeerID,
178        proto_id: &ProtocolID,
179        msg: Vec<u8>,
180    ) -> Result<()> {
181        let peers = self.peers.read().await;
182        let peer = peers
183            .get(peer_id)
184            .ok_or_else(|| Error::PeerNotFound(peer_id.to_string()))?;
185        peer.send(proto_id.to_string(), msg).await
186    }
187
188    /// Returns the negotiated protocol set for a peer.
189    pub async fn peer_protocol_set(&self, pid: &PeerID) -> Option<HashSet<ProtocolID>> {
190        self.peers
191            .read()
192            .await
193            .get(pid)
194            .map(|p| p.negotiated_protocols().clone())
195    }
196
197    /// Checks if the peer list contains a peer with the given peer id
198    pub async fn contains_peer(&self, pid: &PeerID) -> bool {
199        self.peers.read().await.contains_key(pid)
200    }
201
202    /// Returns the number of currently connected peers.
203    pub async fn peers_len(&self) -> usize {
204        self.peers.read().await.len()
205    }
206
207    /// Returns a map of inbound peers with their endpoints.
208    pub async fn inbound_peers(&self) -> HashMap<PeerID, Endpoint> {
209        let mut peers = HashMap::new();
210        for (id, peer) in self.peers.read().await.iter() {
211            if peer.is_inbound() {
212                peers.insert(id.clone(), peer.remote_endpoint().clone());
213            }
214        }
215        peers
216    }
217
218    /// Returns a map of outbound peers with their endpoints.
219    pub async fn outbound_peers(&self) -> HashMap<PeerID, Endpoint> {
220        let mut peers = HashMap::new();
221        for (id, peer) in self.peers.read().await.iter() {
222            if !peer.is_inbound() {
223                peers.insert(id.clone(), peer.remote_endpoint().clone());
224            }
225        }
226        peers
227    }
228
229    async fn run(self: Arc<Self>) {
230        loop {
231            let mut queued = self.conn_queue.next().await;
232
233            // Snapshot the protocol metadata so we don't hold a lock
234            // across the handshake. Drives both version negotiation
235            // and the mandatory-subset check.
236            let meta = self.protocol_meta.read().await.clone();
237
238            let params = HandshakeParams {
239                own_id: &self.id,
240                is_inbound: matches!(queued.direction, ConnDirection::Inbound),
241                config_version: &self.config.version,
242                protocols: &meta,
243                timeout_secs: self.config.handshake_timeout,
244                verified_peer_id: queued.verified_peer_id.as_ref(),
245            };
246            let handshake = handshake(&mut queued.reader, &mut queued.writer, &params).await;
247
248            let (pid, negotiated) = match handshake {
249                Ok(v) => v,
250                Err(err) => {
251                    let pid = queued.verified_peer_id.clone();
252                    self.monitor
253                        .notify(PoolEvent::HandshakeFailed(pid.clone()))
254                        .await;
255                    let _ = self
256                        .peer_emitter
257                        .emit(&PeerEvent::HandshakeFailed(pid))
258                        .await;
259                    let _ = queued.disconnect_signal.send(Err(err)).await;
260                    continue;
261                }
262            };
263
264            // Admission gate. Runs before the Peer is built, so a denied
265            // peer never gets tasks spawned or QUIC streams opened.
266            let candidate = PeerCandidate {
267                peer_id: &pid,
268                remote_endpoint: &queued.remote_endpoint,
269                negotiated_protocols: &negotiated,
270            };
271            let allowed = self.config.access_control.allow(
272                &Subject::Peer(&candidate),
273                Action::Connect(queued.direction.clone()),
274            );
275
276            if !allowed {
277                warn!("Access denied for peer {pid}");
278                self.monitor
279                    .notify(PoolEvent::HandshakeFailed(Some(pid.clone())))
280                    .await;
281                let _ = self
282                    .peer_emitter
283                    .emit(&PeerEvent::HandshakeFailed(Some(pid)))
284                    .await;
285                let _ = queued
286                    .disconnect_signal
287                    .send(Err(Error::AccessDenied))
288                    .await;
289                continue;
290            }
291
292            if let Err(err) = self.new_peer(queued, pid, negotiated).await {
293                error!("new_peer failed: {err}");
294            }
295        }
296    }
297
298    /// Build a Peer from a post-handshake `QueuedConn` and run it.
299    async fn new_peer(
300        self: &Arc<Self>,
301        queued: QueuedConn,
302        pid: PeerID,
303        negotiated: Vec<ProtocolID>,
304    ) -> Result<()> {
305        if self.contains_peer(&pid).await {
306            self.monitor
307                .notify(PoolEvent::PeerAlreadyConnected(pid.clone()))
308                .await;
309            let _ = queued
310                .disconnect_signal
311                .send(Err(Error::PeerAlreadyConnected))
312                .await;
313            return Err(Error::PeerAlreadyConnected);
314        }
315
316        let protocol_ids: Vec<ProtocolID> = self.protocols.read().await.keys().cloned().collect();
317        let negotiated: HashSet<ProtocolID> = negotiated.into_iter().collect();
318
319        let peer = Peer::new(self.clone(), queued, pid.clone(), negotiated, protocol_ids).await?;
320
321        self.peers.write().await.insert(pid.clone(), peer.clone());
322
323        let on_disconnect = {
324            let this = self.clone();
325            let pid = pid.clone();
326            |result| async move {
327                if let TaskResult::Completed(_) = result {
328                    if let Err(err) = this.remove_peer(&pid).await {
329                        error!("Failed to remove peer {pid}: {err}");
330                    }
331                }
332            }
333        };
334
335        self.task_group
336            .spawn_then(peer.clone().run(), on_disconnect);
337
338        info!("Add new peer {pid}");
339        self.monitor.notify(PoolEvent::NewPeer(pid.clone())).await;
340        let _ = self.peer_emitter.emit(&PeerEvent::Added(pid)).await;
341
342        Ok(())
343    }
344
345    /// Shuts down the peer and remove it from the peer list.
346    async fn remove_peer(&self, pid: &PeerID) -> Result<()> {
347        let result = self.peers.write().await.remove(pid);
348
349        let peer = match result {
350            Some(p) => p,
351            None => return Ok(()),
352        };
353
354        let _ = peer.shutdown().await;
355
356        self.monitor
357            .notify(PoolEvent::RemovePeer(pid.clone()))
358            .await;
359        let _ = self
360            .peer_emitter
361            .emit(&PeerEvent::Removed(pid.clone()))
362            .await;
363
364        warn!("Peer {pid} removed",);
365        Ok(())
366    }
367}
karyon_p2p/peer_pool.rs

karyon_p2p/
peer_pool.rs
