1use super::*;
2
3impl Driver {
4 pub(crate) fn upsert_known_destination(
5 &mut self,
6 dest_hash: [u8; 16],
7 announced: crate::destination::AnnouncedIdentity,
8 ) {
9 if let Some(existing) = self.known_destinations.get_mut(&dest_hash) {
10 existing.announced = announced;
11 return;
12 }
13
14 self.enforce_known_destination_cap(true);
15 self.known_destinations.insert(
16 dest_hash,
17 KnownDestinationState {
18 announced,
19 was_used: false,
20 last_used_at: None,
21 retained: false,
22 },
23 );
24 }
25
26 pub(crate) fn known_destination_entry(
27 dest_hash: [u8; 16],
28 state: &KnownDestinationState,
29 ) -> KnownDestinationEntry {
30 KnownDestinationEntry {
31 dest_hash,
32 identity_hash: state.announced.identity_hash.0,
33 public_key: state.announced.public_key,
34 app_data: state.announced.app_data.clone(),
35 hops: state.announced.hops,
36 received_at: state.announced.received_at,
37 receiving_interface: state.announced.receiving_interface,
38 was_used: state.was_used,
39 last_used_at: state.last_used_at,
40 retained: state.retained,
41 }
42 }
43
44 pub(crate) fn known_destination_entries(&self) -> Vec<KnownDestinationEntry> {
45 let mut entries: Vec<_> = self
46 .known_destinations
47 .iter()
48 .map(|(dest_hash, state)| Self::known_destination_entry(*dest_hash, state))
49 .collect();
50 entries.sort_by_key(|a| a.dest_hash);
51 entries
52 }
53
54 pub(crate) fn mark_known_destination_used(&mut self, dest_hash: &[u8; 16]) -> bool {
55 let Some(state) = self.known_destinations.get_mut(dest_hash) else {
56 return false;
57 };
58 state.was_used = true;
59 state.last_used_at = Some(time::now());
60 true
61 }
62
63 pub(crate) fn retain_known_destination(&mut self, dest_hash: &[u8; 16]) -> bool {
64 let Some(state) = self.known_destinations.get_mut(dest_hash) else {
65 return false;
66 };
67 state.retained = true;
68 true
69 }
70
71 pub(crate) fn retain_identity(&mut self, identity_hash: &[u8; 16]) -> bool {
72 let mut retained = false;
73 for state in self.known_destinations.values_mut() {
74 if &state.announced.identity_hash.0 == identity_hash {
75 state.retained = true;
76 retained = true;
77 }
78 }
79 retained
80 }
81
82 pub(crate) fn unretain_known_destination(&mut self, dest_hash: &[u8; 16]) -> bool {
83 let Some(state) = self.known_destinations.get_mut(dest_hash) else {
84 return false;
85 };
86 state.retained = false;
87 true
88 }
89
90 pub(crate) fn known_destination_announced(
91 &self,
92 dest_hash: &[u8; 16],
93 ) -> Option<crate::destination::AnnouncedIdentity> {
94 self.known_destinations
95 .get(dest_hash)
96 .map(|state| state.announced.clone())
97 }
98
99 pub(crate) fn known_destination_relevance_time(state: &KnownDestinationState) -> f64 {
100 state.last_used_at.unwrap_or(state.announced.received_at)
101 }
102
103 pub(crate) fn begin_drain(&mut self, timeout: Duration) {
104 let now = Instant::now();
105 let deadline = now + timeout;
106 match self.lifecycle_state {
107 LifecycleState::Active => {
108 self.lifecycle_state = LifecycleState::Draining;
109 self.drain_started_at = Some(now);
110 self.drain_deadline = Some(deadline);
111 log::info!(
112 "driver entering drain mode with {:.3}s timeout",
113 timeout.as_secs_f64()
114 );
115 self.stop_listener_accepts();
116 }
117 LifecycleState::Draining => {
118 self.drain_deadline = Some(deadline);
119 log::info!(
120 "driver drain deadline updated to {:.3}s from now",
121 timeout.as_secs_f64()
122 );
123 self.stop_listener_accepts();
124 }
125 LifecycleState::Stopping | LifecycleState::Stopped => {
126 log::debug!(
127 "ignoring BeginDrain while lifecycle state is {:?}",
128 self.lifecycle_state
129 );
130 }
131 }
132 }
133
134 pub(crate) fn is_draining(&self) -> bool {
135 matches!(self.lifecycle_state, LifecycleState::Draining)
136 }
137
138 pub fn register_listener_control(&mut self, control: crate::interface::ListenerControl) {
139 self.listener_controls.push(control);
140 }
141
142 pub(crate) fn stop_listener_accepts(&mut self) {
143 for control in &self.listener_controls {
144 control.request_stop();
145 }
146 #[cfg(feature = "hooks")]
147 if let Some(bridge) = self.provider_bridge.as_ref() {
148 bridge.stop_accepting();
149 }
150 }
151
152 pub(crate) fn reject_new_work(&self, op: &str) {
153 log::info!("rejecting {} while node is draining", op);
154 }
155
156 fn interface_writer_queued_frames(&self) -> usize {
157 self.interfaces
158 .values()
159 .map(|entry| {
160 entry
161 .async_writer_metrics
162 .as_ref()
163 .map(|metrics| metrics.queued_frames())
164 .unwrap_or(0)
165 })
166 .sum()
167 }
168
169 fn wait_for_writer_flush(&self, timeout: Duration) {
170 let deadline = Instant::now() + timeout;
171 while Instant::now() < deadline {
172 std::thread::sleep(Duration::from_millis(5));
173 }
174 }
175
176 pub(crate) fn graceful_shutdown(&mut self) {
177 if matches!(self.lifecycle_state, LifecycleState::Stopped) {
178 return;
179 }
180
181 self.lifecycle_state = LifecycleState::Stopping;
182 self.stop_listener_accepts();
183
184 let resource_actions = self.link_manager.cancel_all_resources(&mut self.rng);
185 let resource_flush = resource_actions
186 .iter()
187 .any(|action| matches!(action, LinkManagerAction::SendPacket { .. }));
188 self.dispatch_link_actions(resource_actions);
189
190 let link_actions = self.link_manager.teardown_all_links();
191 let link_flush = link_actions
192 .iter()
193 .any(|action| matches!(action, LinkManagerAction::SendPacket { .. }));
194 self.dispatch_link_actions(link_actions);
195
196 let cleanup_actions = self.link_manager.tick(&mut self.rng);
197 let cleanup_flush = cleanup_actions
198 .iter()
199 .any(|action| matches!(action, LinkManagerAction::SendPacket { .. }));
200 self.dispatch_link_actions(cleanup_actions);
201 self.holepunch_manager.abort_all_sessions();
202
203 if resource_flush
204 || link_flush
205 || cleanup_flush
206 || self.interface_writer_queued_frames() > 0
207 {
208 self.wait_for_writer_flush(DEFAULT_LINK_TEARDOWN_FLUSH);
209 }
210
211 self.engine.void_queues();
212 self.announce_verify_queue
213 .lock()
214 .unwrap_or_else(|poisoned| poisoned.into_inner())
215 .clear();
216 self.sent_packets.clear();
217 self.completed_proofs.clear();
218 if let Some(handle) = self.ratchet_cleanup_handle.take() {
219 let _ = handle.join();
220 }
221
222 self.lifecycle_state = LifecycleState::Stopped;
223 }
224
225 pub(crate) fn drain_error(&self, op: &str) -> String {
226 format!("cannot {} while node is draining", op)
227 }
228
229 pub(crate) fn drain_status(&self) -> DrainStatus {
230 let now = Instant::now();
231 let active_links = self.link_manager.link_count();
232 let active_resource_transfers = self.link_manager.resource_transfer_count();
233 let active_holepunch_sessions = self.holepunch_manager.session_count();
234 let interface_writer_queued_frames = self.interface_writer_queued_frames();
235 #[cfg(feature = "hooks")]
236 let (provider_backlog_events, provider_consumer_queued_events) = self
237 .provider_bridge
238 .as_ref()
239 .map(|bridge| {
240 let stats = bridge.stats();
241 (
242 stats.backlog_len,
243 stats
244 .consumers
245 .iter()
246 .map(|consumer| consumer.queue_len)
247 .sum(),
248 )
249 })
250 .unwrap_or((0, 0));
251 #[cfg(not(feature = "hooks"))]
252 let (provider_backlog_events, provider_consumer_queued_events) = (0, 0);
253 let drain_age_seconds = self
254 .drain_started_at
255 .map(|started| started.elapsed().as_secs_f64());
256 let deadline_remaining_seconds = self.drain_deadline.map(|deadline| {
257 deadline
258 .checked_duration_since(now)
259 .map(|remaining| remaining.as_secs_f64())
260 .unwrap_or(0.0)
261 });
262 let detail = match self.lifecycle_state {
263 LifecycleState::Active => Some("node is accepting normal work".into()),
264 LifecycleState::Draining => {
265 let mut remaining = Vec::new();
266 if active_links > 0 {
267 remaining.push(format!("{active_links} link(s)"));
268 }
269 if active_resource_transfers > 0 {
270 remaining.push(format!("{active_resource_transfers} resource transfer(s)"));
271 }
272 if active_holepunch_sessions > 0 {
273 remaining.push(format!("{active_holepunch_sessions} hole-punch session(s)"));
274 }
275 if interface_writer_queued_frames > 0 {
276 remaining.push(format!(
277 "{interface_writer_queued_frames} queued interface writer frame(s)"
278 ));
279 }
280 if provider_backlog_events > 0 {
281 remaining.push(format!(
282 "{provider_backlog_events} provider backlog event(s)"
283 ));
284 }
285 if provider_consumer_queued_events > 0 {
286 remaining.push(format!(
287 "{provider_consumer_queued_events} queued provider consumer event(s)"
288 ));
289 }
290 Some(if remaining.is_empty() {
291 "node is draining existing work; no active links, resource transfers, hole-punch sessions, or queued writer/provider work remain".into()
292 } else {
293 format!(
294 "node is draining existing work; {} still active",
295 remaining.join(", ")
296 )
297 })
298 }
299 LifecycleState::Stopping => Some("node is tearing down remaining work".into()),
300 LifecycleState::Stopped => Some("node is stopped".into()),
301 };
302
303 DrainStatus {
304 state: self.lifecycle_state,
305 drain_age_seconds,
306 deadline_remaining_seconds,
307 drain_complete: !matches!(self.lifecycle_state, LifecycleState::Draining)
308 || (active_links == 0
309 && active_resource_transfers == 0
310 && active_holepunch_sessions == 0
311 && interface_writer_queued_frames == 0
312 && provider_backlog_events == 0
313 && provider_consumer_queued_events == 0),
314 interface_writer_queued_frames,
315 provider_backlog_events,
316 provider_consumer_queued_events,
317 detail,
318 }
319 }
320
321 pub(crate) fn enforce_drain_deadline(&mut self) {
322 if !matches!(self.lifecycle_state, LifecycleState::Draining) {
323 return;
324 }
325 let Some(deadline) = self.drain_deadline else {
326 return;
327 };
328 if Instant::now() < deadline {
329 return;
330 }
331
332 log::info!("driver drain deadline reached; tearing down remaining links");
333 self.lifecycle_state = LifecycleState::Stopping;
334 let resource_actions = self.link_manager.cancel_all_resources(&mut self.rng);
335 self.dispatch_link_actions(resource_actions);
336 let link_actions = self.link_manager.teardown_all_links();
337 self.dispatch_link_actions(link_actions);
338 let cleanup_actions = self.link_manager.tick(&mut self.rng);
339 self.dispatch_link_actions(cleanup_actions);
340 self.holepunch_manager.abort_all_sessions();
341 }
342
343 pub(crate) fn enforce_known_destination_cap(&mut self, for_insert: bool) -> usize {
344 if self.known_destinations_max_entries == usize::MAX {
345 return 0;
346 }
347
348 let mut evicted = 0usize;
349 while if for_insert {
350 self.known_destinations.len() >= self.known_destinations_max_entries
351 } else {
352 self.known_destinations.len() > self.known_destinations_max_entries
353 } {
354 let active_dests = self.engine.active_destination_hashes();
355 let candidate = self
356 .oldest_known_destination(false, &active_dests)
357 .or_else(|| self.oldest_known_destination(true, &active_dests));
358 let Some(dest_hash) = candidate else {
359 break;
360 };
361 if self.known_destinations.remove(&dest_hash).is_some() {
362 evicted += 1;
363 self.known_destinations_cap_evict_count += 1;
364 } else {
365 break;
366 }
367 }
368 evicted
369 }
370
371 pub(crate) fn oldest_known_destination(
372 &self,
373 include_protected: bool,
374 active_dests: &std::collections::BTreeSet<[u8; 16]>,
375 ) -> Option<[u8; 16]> {
376 self.known_destinations
377 .iter()
378 .filter(|(dest_hash, state)| {
379 include_protected
380 || (!active_dests.contains(*dest_hash)
381 && !self.local_destinations.contains_key(*dest_hash)
382 && !state.retained)
383 })
384 .min_by(|a, b| {
385 Self::known_destination_relevance_time(a.1)
386 .partial_cmp(&Self::known_destination_relevance_time(b.1))
387 .unwrap_or(std::cmp::Ordering::Equal)
388 .then_with(|| a.0.cmp(b.0))
389 })
390 .map(|(dest_hash, _)| *dest_hash)
391 }
392
393 fn build_shared_announce_raw(
394 &mut self,
395 dest_hash: &[u8; 16],
396 record: &SharedAnnounceRecord,
397 path_response: bool,
398 ) -> Option<Vec<u8>> {
399 let identity = rns_crypto::identity::Identity::from_private_key(&record.identity_prv_key);
400
401 let mut random_hash = [0u8; 10];
402 self.rng.fill_bytes(&mut random_hash[..5]);
403 let now_secs = std::time::SystemTime::now()
404 .duration_since(std::time::UNIX_EPOCH)
405 .ok()?
406 .as_secs();
407 random_hash[5..10].copy_from_slice(&now_secs.to_be_bytes()[3..8]);
408
409 let (announce_data, _has_ratchet) = rns_core::announce::AnnounceData::pack(
410 &identity,
411 dest_hash,
412 &record.name_hash,
413 &random_hash,
414 None,
415 record.app_data.as_deref(),
416 )
417 .ok()?;
418
419 let flags = rns_core::packet::PacketFlags {
420 header_type: rns_core::constants::HEADER_1,
421 context_flag: rns_core::constants::FLAG_UNSET,
422 transport_type: rns_core::constants::TRANSPORT_BROADCAST,
423 destination_type: rns_core::constants::DESTINATION_SINGLE,
424 packet_type: rns_core::constants::PACKET_TYPE_ANNOUNCE,
425 };
426 let context = if path_response {
427 rns_core::constants::CONTEXT_PATH_RESPONSE
428 } else {
429 rns_core::constants::CONTEXT_NONE
430 };
431
432 rns_core::packet::RawPacket::pack(flags, 0, dest_hash, None, context, &announce_data)
433 .ok()
434 .map(|packet| packet.raw)
435 }
436
437 pub(crate) fn replay_shared_announces(&mut self) {
438 let records: Vec<([u8; 16], SharedAnnounceRecord)> = self
439 .shared_announces
440 .iter()
441 .map(|(dest_hash, record)| (*dest_hash, record.clone()))
442 .collect();
443 for (dest_hash, record) in records {
444 if let Some(raw) = self.build_shared_announce_raw(&dest_hash, &record, true) {
445 let event = Event::SendOutbound {
446 raw,
447 dest_type: rns_core::constants::DESTINATION_SINGLE,
448 attached_interface: None,
449 };
450 match event {
451 Event::SendOutbound {
452 raw,
453 dest_type,
454 attached_interface,
455 } => match RawPacket::unpack(&raw) {
456 Ok(packet) => {
457 let actions = self.engine.handle_outbound(
458 &packet,
459 dest_type,
460 attached_interface,
461 time::now(),
462 );
463 self.dispatch_all(actions);
464 }
465 Err(e) => {
466 log::warn!(
467 "Shared announce replay failed for {:02x?}: {:?}",
468 &dest_hash[..4],
469 e
470 );
471 }
472 },
473 other => {
474 log::warn!(
475 "shared announce replay returned unexpected response: {:?}",
476 other
477 );
478 }
479 }
480 }
481 }
482 }
483
484 pub(crate) fn handle_shared_interface_down(&mut self, id: InterfaceId) {
485 let dropped_paths = self.engine.drop_paths_for_interface(id);
486 let dropped_reverse = self.engine.drop_reverse_for_interface(id);
487 let dropped_links = self.engine.drop_links_for_interface(id);
488 self.engine.drop_announce_queues();
489 let link_actions = self.link_manager.teardown_all_links();
490 self.dispatch_link_actions(link_actions);
491 self.shared_reconnect_pending.insert(id, true);
492 log::info!(
493 "[{}] cleared shared state: {} paths, {} reverse entries, {} transport links",
494 id.0,
495 dropped_paths,
496 dropped_reverse,
497 dropped_links
498 );
499 }
500}