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
34pub struct Node {
59 config: Arc<Config>,
61
62 key_pair: KeyPair,
64
65 peer_id: PeerID,
67
68 monitor: Arc<Monitor>,
70
71 discovery: Arc<dyn Discovery>,
73
74 peer_pool: Arc<PeerPool>,
76
77 connector: Arc<Connector<PeerNetMsgCodec>>,
79
80 listener: Arc<Listener<PeerNetMsgCodec>>,
82
83 task_group: TaskGroup,
85}
86
87impl Node {
88 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 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 pub async fn run(self: &Arc<Self>) -> Result<()> {
209 self.attach_core_protocols().await?;
212
213 self.peer_pool.start().await?;
214
215 for endpoint in &self.config.listen_endpoints {
217 let resolved = self.listener.start(endpoint.clone()).await?;
218 info!("Listening on {resolved}");
219 }
220
221 self.discovery.clone().start().await?;
223
224 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 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 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 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 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 async fn attach_core_protocols(self: &Arc<Self>) -> Result<()> {
324 self.attach_protocol(PingProtocol::new).await
325 }
326
327 pub fn advertise(&self, item: impl AsRef<[u8]>, flags: ProtocolFlags) {
331 self.discovery.advertise(item.as_ref(), flags);
332 }
333
334 pub fn find_peers_with(&self, item: impl AsRef<[u8]>) -> Vec<DiscoveredPeer> {
338 self.discovery.find_peers_with(item.as_ref())
339 }
340
341 pub async fn peers(&self) -> usize {
343 self.peer_pool.peers_len().await
344 }
345
346 pub fn config(&self) -> Arc<Config> {
348 self.config.clone()
349 }
350
351 pub fn peer_id(&self) -> &PeerID {
353 &self.peer_id
354 }
355
356 pub fn key_pair(&self) -> &KeyPair {
358 &self.key_pair
359 }
360
361 pub async fn inbound_peers(&self) -> HashMap<PeerID, Endpoint> {
363 self.peer_pool.inbound_peers().await
364 }
365
366 pub async fn outbound_peers(&self) -> HashMap<PeerID, Endpoint> {
368 self.peer_pool.outbound_peers().await
369 }
370
371 pub fn monitor(&self) -> Arc<Monitor> {
373 self.monitor.clone()
374 }
375
376 pub fn register_peer_events(&self) -> EventListener<PeerEventTopic, PeerEvent> {
379 self.peer_pool.register_peer_events()
380 }
381
382 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 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 pub async fn peer_protocol_set(&self, pid: &PeerID) -> Option<HashSet<ProtocolID>> {
407 self.peer_pool.peer_protocol_set(pid).await
408 }
409
410 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}