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
36pub struct Node {
61 config: Arc<Config>,
63
64 key_pair: KeyPair,
66
67 peer_id: PeerID,
69
70 monitor: Arc<Monitor>,
72
73 discovery: Arc<dyn Discovery>,
75
76 peer_pool: Arc<PeerPool>,
78
79 connector: Arc<Connector<PeerNetMsgCodec>>,
81
82 listener: Arc<Listener<PeerNetMsgCodec>>,
84
85 task_group: TaskGroup,
87
88 bloom: BloomRef,
93}
94
95impl Node {
96 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 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 pub async fn run(self: &Arc<Self>) -> Result<()> {
227 self.attach_core_protocols().await?;
230
231 self.peer_pool.start().await?;
232
233 for endpoint in &self.config.listen_endpoints {
235 let resolved = self.listener.start(endpoint.clone()).await?;
236 info!("Listening on {resolved}");
237 }
238
239 self.discovery.clone().start().await?;
241
242 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 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 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 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 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 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 pub fn bloom_add_mandatory(&self, item: impl AsRef<[u8]>) {
349 self.bloom.write().add_mandatory(item);
350 }
351
352 pub fn bloom_add_optional(&self, item: impl AsRef<[u8]>) {
356 self.bloom.write().add_optional(item);
357 }
358
359 pub fn bloom_snapshot(&self) -> Bloom {
361 *self.bloom.read()
362 }
363
364 pub fn find_peers_with(&self, item: impl AsRef<[u8]>) -> Vec<DiscoveredPeer> {
368 self.discovery.find_peers_with(item.as_ref())
369 }
370
371 pub async fn peers(&self) -> usize {
373 self.peer_pool.peers_len().await
374 }
375
376 pub fn config(&self) -> Arc<Config> {
378 self.config.clone()
379 }
380
381 pub fn peer_id(&self) -> &PeerID {
383 &self.peer_id
384 }
385
386 pub fn key_pair(&self) -> &KeyPair {
388 &self.key_pair
389 }
390
391 pub async fn inbound_peers(&self) -> HashMap<PeerID, Endpoint> {
393 self.peer_pool.inbound_peers().await
394 }
395
396 pub async fn outbound_peers(&self) -> HashMap<PeerID, Endpoint> {
398 self.peer_pool.outbound_peers().await
399 }
400
401 pub fn monitor(&self) -> Arc<Monitor> {
403 self.monitor.clone()
404 }
405
406 pub fn register_peer_events(&self) -> EventListener<PeerEventTopic, PeerEvent> {
409 self.peer_pool.register_peer_events()
410 }
411
412 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 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 pub async fn peer_protocol_set(&self, pid: &PeerID) -> Option<HashSet<ProtocolID>> {
437 self.peer_pool.peer_protocol_set(pid).await
438 }
439
440 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}