Skip to main content

ant_node/
node.rs

1//! Node implementation - thin wrapper around saorsa-core's `P2PNode`.
2
3use crate::ant_protocol::CHUNK_PROTOCOL_ID;
4use crate::config::{
5    default_nodes_dir, default_root_dir, NetworkMode, NodeConfig, NODE_IDENTITY_FILENAME,
6};
7use crate::error::{Error, Result};
8use crate::event::{create_event_channel, NodeEvent, NodeEventsChannel, NodeEventsSender};
9use crate::logging::{debug, error, info, warn};
10use crate::payment::metrics::QuotingMetricsTracker;
11use crate::payment::wallet::parse_rewards_address;
12use crate::payment::{EvmVerifierConfig, PaymentVerifier, PaymentVerifierConfig, QuoteGenerator};
13use crate::replication::config::ReplicationConfig;
14use crate::replication::ReplicationEngine;
15use crate::storage::lmdb::MIB;
16use crate::storage::{AntProtocol, LmdbStorage, LmdbStorageConfig};
17use crate::upgrade::{
18    upgrade_cache_dir, AutoApplyUpgrader, BinaryCache, ReleaseCache, UpgradeMonitor, UpgradeResult,
19};
20use rand::Rng;
21use saorsa_core::identity::NodeIdentity;
22use saorsa_core::{
23    IPDiversityConfig as CoreDiversityConfig, MultiAddr, NodeConfig as CoreNodeConfig, P2PEvent,
24    P2PNode,
25};
26use std::path::PathBuf;
27use std::sync::atomic::{AtomicI32, Ordering};
28use std::sync::Arc;
29use tokio::sync::Semaphore;
30use tokio::task::JoinHandle;
31use tokio_util::sync::CancellationToken;
32
33#[cfg(unix)]
34use tokio::signal::unix::{signal, SignalKind};
35
36/// Builder for constructing an Ant node.
37pub struct NodeBuilder {
38    config: NodeConfig,
39}
40
41impl NodeBuilder {
42    /// Create a new node builder with the given configuration.
43    #[must_use]
44    pub fn new(config: NodeConfig) -> Self {
45        Self { config }
46    }
47
48    /// Build and start the node.
49    ///
50    /// # Errors
51    ///
52    /// Returns an error if the node fails to start.
53    pub async fn build(mut self) -> Result<RunningNode> {
54        info!("Building ant-node with config: {:?}", self.config);
55
56        // Validate rewards address in production
57        if self.config.network_mode == NetworkMode::Production {
58            match self.config.payment.rewards_address {
59                None => {
60                    return Err(Error::Config(
61                        "CRITICAL: Rewards address is not configured. \
62                         Set payment.rewards_address in config to your Arbitrum wallet address."
63                            .to_string(),
64                    ));
65                }
66                Some(ref addr) if addr == "0xYOUR_ARBITRUM_ADDRESS_HERE" || addr.is_empty() => {
67                    return Err(Error::Config(
68                        "CRITICAL: Rewards address is not configured. \
69                         Set payment.rewards_address in config to your Arbitrum wallet address."
70                            .to_string(),
71                    ));
72                }
73                Some(_) => {}
74            }
75        }
76
77        // Resolve identity and root_dir (may update self.config.root_dir)
78        let identity = Arc::new(Self::resolve_identity(&mut self.config).await?);
79        let peer_id = identity.peer_id().to_hex();
80
81        info!(peer_id = %peer_id, root_dir = %self.config.root_dir.display(), "Node identity resolved");
82
83        // Ensure root directory exists
84        std::fs::create_dir_all(&self.config.root_dir)?;
85
86        // Create shutdown token
87        let shutdown = CancellationToken::new();
88
89        // Create event channel
90        let (events_tx, events_rx) = create_event_channel();
91
92        // Convert our config to saorsa-core's config
93        let mut core_config = Self::build_core_config(&self.config)?;
94        // Inject the ML-DSA identity so the P2PNode's transport peer ID
95        // matches the pub_key embedded in payment quotes.
96        core_config.node_identity = Some(Arc::clone(&identity));
97        debug!("Core config: {:?}", core_config);
98
99        // Initialize saorsa-core's P2PNode
100        let p2p_node = P2PNode::new(core_config)
101            .await
102            .map_err(|e| Error::Startup(format!("Failed to create P2P node: {e}")))?;
103
104        // Create upgrade monitor
105        let upgrade_monitor = {
106            let node_id_seed = p2p_node.peer_id().as_bytes();
107            Some(Self::build_upgrade_monitor(&self.config, node_id_seed))
108        };
109
110        let repl_config = ReplicationConfig::default();
111
112        // Initialize ANT protocol handler for chunk storage and
113        // wire the fresh-write channel so PUTs trigger replication.
114        let (ant_protocol, fresh_write_rx) = if self.config.storage.enabled {
115            let (fresh_write_tx, fresh_write_rx) = tokio::sync::mpsc::unbounded_channel();
116            let mut protocol =
117                Self::build_ant_protocol(&self.config, &identity, repl_config.close_group_size)
118                    .await?;
119            protocol.set_fresh_write_sender(fresh_write_tx);
120            (Some(Arc::new(protocol)), Some(fresh_write_rx))
121        } else {
122            info!("Chunk storage disabled");
123            (None, None)
124        };
125
126        let p2p_arc = Arc::new(p2p_node);
127
128        // Wire the P2PNode handle into AntProtocol so payment proofs can query
129        // live-DHT closeness.
130        if let Some(ref protocol) = ant_protocol {
131            protocol.attach_p2p_node(Arc::clone(&p2p_arc));
132        }
133
134        // Initialize replication engine (if storage is enabled)
135        let replication_engine =
136            if let (Some(ref protocol), Some(fresh_rx)) = (&ant_protocol, fresh_write_rx) {
137                let storage_arc = protocol.storage();
138                let payment_verifier_arc = protocol.payment_verifier_arc();
139                match ReplicationEngine::new(
140                    repl_config,
141                    Arc::clone(&p2p_arc),
142                    storage_arc,
143                    payment_verifier_arc,
144                    &self.config.root_dir,
145                    fresh_rx,
146                    shutdown.clone(),
147                )
148                .await
149                {
150                    Ok(engine) => Some(engine),
151                    Err(e) => {
152                        warn!("Failed to initialize replication engine: {e}");
153                        None
154                    }
155                }
156            } else {
157                None
158            };
159
160        let node = RunningNode {
161            config: self.config,
162            p2p_node: p2p_arc,
163            shutdown,
164            events_tx,
165            events_rx: Some(events_rx),
166            upgrade_monitor,
167            ant_protocol,
168            replication_engine,
169            protocol_task: None,
170            upgrade_exit_code: Arc::new(AtomicI32::new(-1)),
171        };
172
173        Ok(node)
174    }
175
176    /// Build the saorsa-core `NodeConfig` from our config.
177    fn build_core_config(config: &NodeConfig) -> Result<CoreNodeConfig> {
178        let local = matches!(config.network_mode, NetworkMode::Development);
179
180        let mut core_config = CoreNodeConfig::builder()
181            .port(config.port)
182            .ipv6(!config.ipv4_only)
183            .local(local)
184            .max_message_size(config.max_message_size)
185            .build()
186            .map_err(|e| Error::Config(format!("Failed to create core config: {e}")))?;
187
188        // Add bootstrap peers.
189        core_config.bootstrap_peers = config
190            .bootstrap
191            .iter()
192            .map(|addr| MultiAddr::quic(*addr))
193            .collect();
194
195        // Propagate network-mode tuning into saorsa-core where supported.
196        match config.network_mode {
197            NetworkMode::Production => {
198                core_config.diversity_config = Some(CoreDiversityConfig::default());
199            }
200            NetworkMode::Testnet => {
201                // Testnet allows loopback so nodes can be co-located on one machine.
202                core_config.allow_loopback = true;
203                core_config.diversity_config = Some(CoreDiversityConfig {
204                    max_per_ip: config.testnet.max_per_ip,
205                    max_per_subnet: config.testnet.max_per_subnet,
206                });
207            }
208            NetworkMode::Development => {
209                core_config.diversity_config = Some(CoreDiversityConfig::permissive());
210            }
211        }
212
213        // Persist close group peers + trust scores across restarts.
214        // Default to root_dir (alongside node_identity.key) when not explicitly set.
215        core_config.close_group_cache_dir = Some(
216            config
217                .close_group_cache_dir
218                .clone()
219                .unwrap_or_else(|| config.root_dir.clone()),
220        );
221
222        Ok(core_config)
223    }
224
225    /// Resolve the node identity from disk or generate a new one.
226    ///
227    /// **When `root_dir` differs from the platform default** (set via `--root-dir`
228    /// or loaded from `config.toml`):
229    ///   - Use `root_dir` directly: load existing identity or generate a new one.
230    ///
231    /// **When `root_dir` is the platform default** (first run, no config file):
232    ///   1. Scan `{default_root_dir}/nodes/` for subdirectories containing
233    ///      `node_identity.key`.
234    ///   2. **None found** — first run: generate identity, create
235    ///      `nodes/{full_peer_id}/`, save identity there, update `config.root_dir`.
236    ///   3. **Exactly one found** — load it and update `config.root_dir`.
237    ///   4. **Multiple found** — return an error asking for `--root-dir`.
238    async fn resolve_identity(config: &mut NodeConfig) -> Result<NodeIdentity> {
239        if config.root_dir != default_root_dir() {
240            return Self::load_or_generate_identity(&config.root_dir).await;
241        }
242
243        let nodes_dir = default_nodes_dir();
244        let identity_dirs = Self::scan_identity_dirs(&nodes_dir)?;
245
246        match identity_dirs.len() {
247            0 => {
248                // First run: generate new identity and create a peer-id-scoped subdirectory
249                let identity = NodeIdentity::generate().map_err(|e| {
250                    Error::Startup(format!("Failed to generate node identity: {e}"))
251                })?;
252                let peer_id = identity.peer_id().to_hex();
253                let peer_dir = nodes_dir.join(&peer_id);
254                std::fs::create_dir_all(&peer_dir)?;
255                identity
256                    .save_to_file(&peer_dir.join(NODE_IDENTITY_FILENAME))
257                    .await
258                    .map_err(|e| Error::Startup(format!("Failed to save node identity: {e}")))?;
259                config.root_dir = peer_dir;
260                Ok(identity)
261            }
262            1 => {
263                let dir = identity_dirs
264                    .first()
265                    .ok_or_else(|| Error::Config("No identity dirs found".to_string()))?;
266                let identity = NodeIdentity::load_from_file(&dir.join(NODE_IDENTITY_FILENAME))
267                    .await
268                    .map_err(|e| Error::Startup(format!("Failed to load node identity: {e}")))?;
269                config.root_dir.clone_from(dir);
270                Ok(identity)
271            }
272            _ => {
273                let dirs: Vec<String> = identity_dirs
274                    .iter()
275                    .filter_map(|d| d.file_name().map(|n| n.to_string_lossy().into_owned()))
276                    .collect();
277                Err(Error::Config(format!(
278                    "Multiple node identities found at {}: [{}]. Specify --root-dir to select one.",
279                    nodes_dir.display(),
280                    dirs.join(", ")
281                )))
282            }
283        }
284    }
285
286    /// Load an existing identity from `dir/node_identity.key`, or generate and save a new one.
287    async fn load_or_generate_identity(dir: &std::path::Path) -> Result<NodeIdentity> {
288        let key_path = dir.join(NODE_IDENTITY_FILENAME);
289        if key_path.exists() {
290            NodeIdentity::load_from_file(&key_path)
291                .await
292                .map_err(|e| Error::Startup(format!("Failed to load node identity: {e}")))
293        } else {
294            let identity = NodeIdentity::generate()
295                .map_err(|e| Error::Startup(format!("Failed to generate node identity: {e}")))?;
296            std::fs::create_dir_all(dir)?;
297            identity
298                .save_to_file(&key_path)
299                .await
300                .map_err(|e| Error::Startup(format!("Failed to save node identity: {e}")))?;
301            Ok(identity)
302        }
303    }
304
305    /// Scan `base_dir` for immediate subdirectories that contain `node_identity.key`.
306    fn scan_identity_dirs(base_dir: &std::path::Path) -> Result<Vec<PathBuf>> {
307        let mut dirs = Vec::new();
308        let read_dir = match std::fs::read_dir(base_dir) {
309            Ok(rd) => rd,
310            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(dirs),
311            Err(e) => return Err(e.into()),
312        };
313        for entry in read_dir {
314            let entry = entry?;
315            let path = entry.path();
316            if path.is_dir() && path.join(NODE_IDENTITY_FILENAME).exists() {
317                dirs.push(path);
318            }
319        }
320        Ok(dirs)
321    }
322
323    fn build_upgrade_monitor(config: &NodeConfig, node_id_seed: &[u8]) -> UpgradeMonitor {
324        let mut monitor = UpgradeMonitor::new(
325            config.upgrade.github_repo.clone(),
326            config.upgrade.channel,
327            config.upgrade.check_interval_hours,
328        );
329
330        if let Ok(cache_dir) = upgrade_cache_dir() {
331            monitor = monitor.with_release_cache(ReleaseCache::new(
332                cache_dir,
333                std::time::Duration::from_secs(3600),
334            ));
335        }
336
337        if config.upgrade.staged_rollout_hours > 0 {
338            monitor =
339                monitor.with_staged_rollout(node_id_seed, config.upgrade.staged_rollout_hours);
340        }
341
342        monitor
343    }
344
345    /// Build the ANT protocol handler from config.
346    ///
347    /// Initializes LMDB storage, payment verifier, and quote generator.
348    /// Wires ML-DSA-65 signing from the node's identity into the quote generator.
349    async fn build_ant_protocol(
350        config: &NodeConfig,
351        identity: &NodeIdentity,
352        close_group_size: usize,
353    ) -> Result<AntProtocol> {
354        // Create LMDB storage
355        let storage_config = LmdbStorageConfig {
356            root_dir: config.root_dir.clone(),
357            verify_on_read: config.storage.verify_on_read,
358            max_map_size: config.storage.db_size_gb.saturating_mul(1024 * 1024 * 1024),
359            disk_reserve: config.storage.disk_reserve_mb.saturating_mul(MIB),
360        };
361        let storage = LmdbStorage::new(storage_config)
362            .await
363            .map_err(|e| Error::Startup(format!("Failed to create LMDB storage: {e}")))?;
364
365        // Parse rewards address (required — node must know where to receive payments)
366        let rewards_address = match config.payment.rewards_address {
367            Some(ref addr) => parse_rewards_address(addr)?,
368            None => {
369                return Err(Error::Startup(
370                    "No rewards address configured. Set --rewards-address or payment.rewards_address in config.".to_string(),
371                ));
372            }
373        };
374
375        // Create payment verifier
376        let evm_network = config.payment.evm_network.clone().into_evm_network();
377        let payment_config = PaymentVerifierConfig {
378            evm: EvmVerifierConfig {
379                network: evm_network,
380            },
381            cache_capacity: config.payment.cache_capacity,
382            close_group_size,
383            local_rewards_address: rewards_address,
384        };
385        let payment_verifier = PaymentVerifier::new(payment_config);
386        let metrics_tracker = QuotingMetricsTracker::new(0);
387        let mut quote_generator = QuoteGenerator::new(rewards_address, metrics_tracker);
388
389        // Wire ML-DSA-65 signing from node identity.
390        // This same signer is used for both regular quotes and merkle candidate quotes.
391        crate::payment::wire_ml_dsa_signer(&mut quote_generator, identity)?;
392
393        let storage = Arc::new(storage);
394        let payment_verifier = Arc::new(payment_verifier);
395
396        let protocol = AntProtocol::new(storage, payment_verifier, Arc::new(quote_generator));
397
398        info!(
399            "ANT protocol handler initialized with ML-DSA-65 signing (protocol={CHUNK_PROTOCOL_ID})"
400        );
401
402        Ok(protocol)
403    }
404}
405
406/// A running Ant node.
407pub struct RunningNode {
408    config: NodeConfig,
409    p2p_node: Arc<P2PNode>,
410    shutdown: CancellationToken,
411    events_tx: NodeEventsSender,
412    events_rx: Option<NodeEventsChannel>,
413    upgrade_monitor: Option<UpgradeMonitor>,
414    /// ANT protocol handler for chunk storage.
415    ant_protocol: Option<Arc<AntProtocol>>,
416    /// Replication engine (manages neighbor sync, verification, audits).
417    replication_engine: Option<ReplicationEngine>,
418    /// Protocol message routing background task.
419    protocol_task: Option<JoinHandle<()>>,
420    /// Exit code requested by a successful upgrade (-1 = no upgrade exit pending).
421    upgrade_exit_code: Arc<AtomicI32>,
422}
423
424impl RunningNode {
425    /// Get the node's root directory.
426    #[must_use]
427    pub fn root_dir(&self) -> &PathBuf {
428        &self.config.root_dir
429    }
430
431    /// Get a receiver for node events.
432    ///
433    /// Note: Can only be called once. Subsequent calls return None.
434    pub fn events(&mut self) -> Option<NodeEventsChannel> {
435        self.events_rx.take()
436    }
437
438    /// Subscribe to node events.
439    #[must_use]
440    pub fn subscribe_events(&self) -> NodeEventsChannel {
441        self.events_tx.subscribe()
442    }
443
444    /// Run the node until shutdown is requested.
445    ///
446    /// # Errors
447    ///
448    /// Returns an error if the node encounters a fatal error.
449    #[allow(clippy::too_many_lines)]
450    pub async fn run(&mut self) -> Result<()> {
451        info!("Node runtime loop starting");
452
453        // Subscribe to DHT events BEFORE starting the P2P node so the
454        // bootstrap-sync task does not miss the BootstrapComplete event
455        // emitted during P2PNode::start().
456        let dht_events_for_bootstrap = self
457            .replication_engine
458            .as_ref()
459            .map(|_| self.p2p_node.dht_manager().subscribe_events());
460
461        // Start the P2P node
462        self.p2p_node
463            .start()
464            .await
465            .map_err(|e| Error::Startup(format!("Failed to start P2P node: {e}")))?;
466
467        let listen_addrs = self.p2p_node.listen_addrs().await;
468        info!(listen_addrs = ?listen_addrs, "P2P node started");
469
470        // Extract the actual bound port (config port may be 0 = auto-select)
471        let actual_port = listen_addrs
472            .first()
473            .and_then(MultiAddr::port)
474            .unwrap_or(self.config.port);
475        info!(
476            port = actual_port,
477            "Node is running on port: {}", actual_port
478        );
479
480        // Emit started event
481        if let Err(e) = self.events_tx.send(NodeEvent::Started) {
482            warn!("Failed to send Started event: {e}");
483        }
484
485        // Start protocol message routing (P2P → AntProtocol → P2P response)
486        self.start_protocol_routing();
487
488        // Start replication engine background tasks
489        if let Some(ref mut engine) = self.replication_engine {
490            // Safety: dht_events_for_bootstrap is Some when replication_engine
491            // is Some (both arms use the same condition).
492            if let Some(dht_events) = dht_events_for_bootstrap {
493                engine.start(dht_events);
494            }
495            info!("Replication engine started");
496        }
497
498        // Start upgrade monitor if enabled
499        if let Some(monitor) = self.upgrade_monitor.take() {
500            let events_tx = self.events_tx.clone();
501            let shutdown = self.shutdown.clone();
502            let stop_on_upgrade = self.config.upgrade.stop_on_upgrade;
503            let upgrade_exit_code = Arc::clone(&self.upgrade_exit_code);
504
505            tokio::spawn(async move {
506                let mut monitor = monitor;
507                let mut upgrader = AutoApplyUpgrader::new().with_stop_on_upgrade(stop_on_upgrade);
508                if let Ok(cache_dir) = upgrade_cache_dir() {
509                    upgrader = upgrader.with_binary_cache(BinaryCache::new(cache_dir));
510                }
511
512                // Add randomized jitter before the first upgrade check to prevent all nodes
513                // from hitting the GitHub API simultaneously when started together.
514                {
515                    let jitter_duration = jittered_interval(monitor.check_interval());
516                    let first_check_time = chrono::Utc::now()
517                        + chrono::Duration::from_std(jitter_duration).unwrap_or_else(|e| {
518                            warn!("chrono::Duration::from_std failed for jitter ({e}), defaulting to 1 minute");
519                            chrono::Duration::minutes(1)
520                        });
521                    info!(
522                        "First upgrade check scheduled for {} (jitter: {}s)",
523                        first_check_time.to_rfc3339(),
524                        jitter_duration.as_secs()
525                    );
526                    tokio::time::sleep(jitter_duration).await;
527                }
528
529                loop {
530                    tokio::select! {
531                        () = shutdown.cancelled() => {
532                            break;
533                        }
534                        result = monitor.check_for_ready_upgrade() => {
535                            match result {
536                                Ok(Some(upgrade_info)) => {
537                                    info!(
538                                        current_version = %upgrader.current_version(),
539                                        new_version = %upgrade_info.version,
540                                        "Upgrade available"
541                                    );
542
543                                    // Send notification event
544                                    if let Err(e) = events_tx.send(NodeEvent::UpgradeAvailable {
545                                        version: upgrade_info.version.to_string(),
546                                    }) {
547                                        warn!("Failed to send UpgradeAvailable event: {e}");
548                                    }
549
550                                    // Auto-apply the upgrade
551                                    info!("Starting auto-apply upgrade...");
552                                    match upgrader.apply_upgrade(&upgrade_info).await {
553                                        Ok(UpgradeResult::Success { version, exit_code }) => {
554                                            info!("Upgrade to {} successful, initiating graceful shutdown", version);
555                                            upgrade_exit_code.store(exit_code, Ordering::SeqCst);
556                                            shutdown.cancel();
557                                            break;
558                                        }
559                                        Ok(UpgradeResult::RolledBack { reason }) => {
560                                            warn!("Error during upgrade process: {}", reason);
561                                        }
562                                        Ok(UpgradeResult::NoUpgrade) => {
563                                            info!("Already running latest version");
564                                        }
565                                        Err(e) => {
566                                            error!("Error during upgrade process: {}", e);
567                                        }
568                                    }
569                                }
570                                Ok(None) => {
571                                    if let Some(remaining) = monitor.time_until_upgrade() {
572                                        info!(
573                                            "Upgrade pending, rollout delay remaining: {}m {}s",
574                                            remaining.as_secs() / 60,
575                                            remaining.as_secs() % 60
576                                        );
577                                    } else {
578                                        info!("No upgrade available");
579                                    }
580                                }
581                                Err(e) => {
582                                    warn!("Error during upgrade process: {}", e);
583                                }
584                            }
585                            // If an upgrade is pending, sleep for exactly the remaining
586                            // rollout delay so the node restarts at its scheduled time
587                            // rather than waiting for the next check interval tick.
588                            let sleep_duration = monitor.time_until_upgrade().map_or_else(
589                                || {
590                                    // No pending upgrade - schedule next check with jitter
591                                    let jittered_duration =
592                                        jittered_interval(monitor.check_interval());
593                                    let next_check = chrono::Utc::now()
594                                        + chrono::Duration::from_std(jittered_duration).unwrap_or_else(|e| {
595                                            warn!("chrono::Duration::from_std failed for interval ({e}), defaulting to 1 hour");
596                                            chrono::Duration::hours(1)
597                                        });
598                                    info!("Next upgrade check scheduled for {}", next_check.to_rfc3339());
599                                    jittered_duration
600                                },
601                                |remaining| {
602                                    // If the rollout delay has fully elapsed but the upgrade was
603                                    // not successfully applied, avoid a tight loop by backing off
604                                    // at least one check interval before retrying.
605                                    if remaining.is_zero() {
606                                        let backoff = jittered_interval(monitor.check_interval());
607                                        let next_check = chrono::Utc::now()
608                                            + chrono::Duration::from_std(backoff).unwrap_or_else(|e| {
609                                                warn!("chrono::Duration::from_std failed for backoff ({e}), defaulting to 1 hour");
610                                                chrono::Duration::hours(1)
611                                            });
612                                        info!(
613                                            "Upgrade rollout delay elapsed but previous apply did not succeed; \
614                                             backing off, next check scheduled for {}",
615                                            next_check.to_rfc3339()
616                                        );
617                                        backoff
618                                    } else {
619                                        let wake_time = chrono::Utc::now()
620                                            + chrono::Duration::from_std(remaining).unwrap_or_else(|e| {
621                                                warn!("chrono::Duration::from_std failed for rollout delay ({e}), defaulting to 1 minute");
622                                                chrono::Duration::minutes(1)
623                                            });
624                                        info!("Will apply upgrade at {}", wake_time.to_rfc3339());
625                                        remaining
626                                    }
627                                },
628                            );
629                            // Use select! so shutdown can interrupt long sleeps
630                            // (e.g. during a full rollout window delay).
631                            tokio::select! {
632                                () = shutdown.cancelled() => {
633                                    break;
634                                }
635                                () = tokio::time::sleep(sleep_duration) => {}
636                            }
637                        }
638                    }
639                }
640            });
641        }
642
643        info!("Node running, waiting for shutdown signal");
644
645        // Run the main event loop with signal handling
646        self.run_event_loop().await?;
647
648        // Shutdown replication engine before P2P so background tasks don't
649        // use a dead P2P layer, and Arc<LmdbStorage> references are released.
650        if let Some(ref mut engine) = self.replication_engine {
651            engine.shutdown().await;
652        }
653
654        // Stop protocol routing task
655        if let Some(handle) = self.protocol_task.take() {
656            handle.abort();
657        }
658
659        // Shutdown P2P node
660        info!("Shutting down P2P node...");
661        if let Err(e) = self.p2p_node.shutdown().await {
662            warn!("Error during P2P node shutdown: {e}");
663        }
664
665        if let Err(e) = self.events_tx.send(NodeEvent::ShuttingDown) {
666            warn!("Failed to send ShuttingDown event: {e}");
667        }
668        info!("Node shutdown complete");
669
670        // If an upgrade triggered the shutdown, exit with the requested code.
671        // This happens *after* all cleanup (P2P shutdown, log flush, etc.) so
672        // that destructors and async resources are properly torn down.
673        let exit_code = self.upgrade_exit_code.load(Ordering::SeqCst);
674        if exit_code >= 0 {
675            info!("Exiting with code {} for upgrade restart", exit_code);
676            std::process::exit(exit_code);
677        }
678
679        Ok(())
680    }
681
682    /// Run the main event loop, handling shutdown and signals.
683    #[cfg(unix)]
684    async fn run_event_loop(&self) -> Result<()> {
685        let mut sigterm = signal(SignalKind::terminate())?;
686        let mut sighup = signal(SignalKind::hangup())?;
687
688        loop {
689            tokio::select! {
690                () = self.shutdown.cancelled() => {
691                    info!("Shutdown signal received");
692                    break;
693                }
694                _ = tokio::signal::ctrl_c() => {
695                    info!("Received SIGINT (Ctrl-C), initiating shutdown");
696                    self.shutdown();
697                    break;
698                }
699                _ = sigterm.recv() => {
700                    info!("Received SIGTERM, initiating shutdown");
701                    self.shutdown();
702                    break;
703                }
704                _ = sighup.recv() => {
705                    info!("Received SIGHUP (config reload not yet supported)");
706                }
707            }
708        }
709        Ok(())
710    }
711
712    /// Run the main event loop, handling shutdown signals (non-Unix version).
713    #[cfg(not(unix))]
714    async fn run_event_loop(&self) -> Result<()> {
715        loop {
716            tokio::select! {
717                () = self.shutdown.cancelled() => {
718                    info!("Shutdown signal received");
719                    break;
720                }
721                _ = tokio::signal::ctrl_c() => {
722                    info!("Received Ctrl-C, initiating shutdown");
723                    self.shutdown();
724                    break;
725                }
726            }
727        }
728        Ok(())
729    }
730
731    /// Start the protocol message routing background task.
732    ///
733    /// Subscribes to P2P events and routes incoming chunk protocol messages
734    /// to the `AntProtocol` handler, sending responses back to the sender.
735    fn start_protocol_routing(&mut self) {
736        let protocol = match self.ant_protocol {
737            Some(ref p) => Arc::clone(p),
738            None => return,
739        };
740
741        let mut events = self.p2p_node.subscribe_events();
742        let p2p = Arc::clone(&self.p2p_node);
743        let semaphore = Arc::new(Semaphore::new(64));
744
745        self.protocol_task = Some(tokio::spawn(async move {
746            while let Ok(event) = events.recv().await {
747                if let P2PEvent::Message {
748                    topic,
749                    source: Some(source),
750                    data,
751                    ..
752                } = event
753                {
754                    let handler_info: Option<(&str, &str)> = if topic == CHUNK_PROTOCOL_ID {
755                        Some(("chunk", CHUNK_PROTOCOL_ID))
756                    } else {
757                        None
758                    };
759
760                    if let Some((data_type, response_topic)) = handler_info {
761                        debug!("Received {data_type} protocol message from {source}");
762                        let protocol = Arc::clone(&protocol);
763                        let p2p = Arc::clone(&p2p);
764                        let sem = semaphore.clone();
765                        tokio::spawn(async move {
766                            let Ok(_permit) = sem.acquire().await else {
767                                return;
768                            };
769                            let result = match data_type {
770                                "chunk" => protocol.try_handle_request(&data).await,
771                                _ => return,
772                            };
773                            match result {
774                                Ok(Some(response)) => {
775                                    if let Err(e) = p2p
776                                        .send_message(
777                                            &source,
778                                            response_topic,
779                                            response.to_vec(),
780                                            &[],
781                                        )
782                                        .await
783                                    {
784                                        warn!("Failed to send {data_type} protocol response to {source}: {e}");
785                                    }
786                                }
787                                Ok(None) => {}
788                                Err(e) => {
789                                    warn!("{data_type} protocol handler error: {e}");
790                                }
791                            }
792                        });
793                    }
794                }
795            }
796        }));
797        info!("Protocol message routing started");
798    }
799
800    /// Request the node to shut down.
801    pub fn shutdown(&self) {
802        self.shutdown.cancel();
803    }
804}
805
806/// Apply ±5% jitter to a base interval to prevent thundering-herd behaviour
807/// when multiple nodes check for upgrades on the same schedule.
808fn jittered_interval(base: std::time::Duration) -> std::time::Duration {
809    let secs = base.as_secs();
810    let variance = secs / 20; // 5%
811    if variance == 0 {
812        return base;
813    }
814    let jitter = rand::thread_rng().gen_range(0..=variance * 2);
815    std::time::Duration::from_secs(secs.saturating_sub(variance) + jitter)
816}
817
818#[cfg(test)]
819#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
820mod tests {
821    use super::*;
822    use crate::config::NODES_SUBDIR;
823
824    #[test]
825    fn test_build_upgrade_monitor_staged_rollout_enabled() {
826        let config = NodeConfig {
827            upgrade: crate::config::UpgradeConfig {
828                staged_rollout_hours: 24,
829                ..Default::default()
830            },
831            ..Default::default()
832        };
833        let seed = b"node-seed";
834
835        let monitor = NodeBuilder::build_upgrade_monitor(&config, seed);
836        assert!(monitor.has_staged_rollout());
837    }
838
839    #[test]
840    fn test_build_upgrade_monitor_staged_rollout_disabled() {
841        let config = NodeConfig {
842            upgrade: crate::config::UpgradeConfig {
843                staged_rollout_hours: 0,
844                ..Default::default()
845            },
846            ..Default::default()
847        };
848        let seed = b"node-seed";
849
850        let monitor = NodeBuilder::build_upgrade_monitor(&config, seed);
851        assert!(!monitor.has_staged_rollout());
852    }
853
854    #[test]
855    fn test_build_core_config_sets_production_mode() {
856        let config = NodeConfig {
857            network_mode: NetworkMode::Production,
858            ..Default::default()
859        };
860        let core = NodeBuilder::build_core_config(&config).expect("core config");
861        assert!(core.diversity_config.is_some());
862    }
863
864    #[test]
865    fn test_build_core_config_ipv4_only() {
866        let config = NodeConfig {
867            ipv4_only: true,
868            ..Default::default()
869        };
870        let core = NodeBuilder::build_core_config(&config).expect("core config");
871        assert!(!core.ipv6, "ipv4_only should disable IPv6");
872    }
873
874    #[test]
875    fn test_build_core_config_dual_stack_by_default() {
876        let config = NodeConfig::default();
877        let core = NodeBuilder::build_core_config(&config).expect("core config");
878        assert!(core.ipv6, "dual-stack should be the default");
879    }
880
881    #[test]
882    fn test_build_core_config_sets_development_mode_permissive() {
883        let config = NodeConfig {
884            network_mode: NetworkMode::Development,
885            ..Default::default()
886        };
887        let core = NodeBuilder::build_core_config(&config).expect("core config");
888        let diversity = core.diversity_config.expect("diversity");
889        assert_eq!(diversity.max_per_ip, Some(usize::MAX));
890        assert_eq!(diversity.max_per_subnet, Some(usize::MAX));
891    }
892
893    #[test]
894    fn test_scan_identity_dirs_empty_dir() {
895        let tmp = tempfile::tempdir().unwrap();
896        let dirs = NodeBuilder::scan_identity_dirs(tmp.path()).unwrap();
897        assert!(dirs.is_empty());
898    }
899
900    #[test]
901    fn test_scan_identity_dirs_nonexistent_dir() {
902        let tmp = tempfile::tempdir().unwrap();
903        let path = tmp.path().join("nonexistent_identity_dir");
904        let dirs = NodeBuilder::scan_identity_dirs(&path).unwrap();
905        assert!(dirs.is_empty());
906    }
907
908    #[test]
909    fn test_scan_identity_dirs_finds_one() {
910        let tmp = tempfile::tempdir().unwrap();
911        let node_dir = tmp.path().join("abc123");
912        std::fs::create_dir_all(&node_dir).unwrap();
913        std::fs::write(node_dir.join(NODE_IDENTITY_FILENAME), "{}").unwrap();
914
915        let dirs = NodeBuilder::scan_identity_dirs(tmp.path()).unwrap();
916        assert_eq!(dirs.len(), 1);
917        assert_eq!(dirs[0], node_dir);
918    }
919
920    #[test]
921    fn test_scan_identity_dirs_finds_multiple() {
922        let tmp = tempfile::tempdir().unwrap();
923        for name in &["node_a", "node_b"] {
924            let dir = tmp.path().join(name);
925            std::fs::create_dir_all(&dir).unwrap();
926            std::fs::write(dir.join(NODE_IDENTITY_FILENAME), "{}").unwrap();
927        }
928        // A directory without a key file should be ignored
929        std::fs::create_dir_all(tmp.path().join("no_key")).unwrap();
930
931        let dirs = NodeBuilder::scan_identity_dirs(tmp.path()).unwrap();
932        assert_eq!(dirs.len(), 2);
933    }
934
935    #[tokio::test]
936    async fn test_resolve_identity_first_run_creates_identity() {
937        let tmp = tempfile::tempdir().unwrap();
938        let mut config = NodeConfig {
939            root_dir: tmp.path().to_path_buf(),
940            ..Default::default()
941        };
942
943        let identity = NodeBuilder::resolve_identity(&mut config).await.unwrap();
944        // Key file should exist
945        assert!(tmp.path().join(NODE_IDENTITY_FILENAME).exists());
946        // peer_id should be derivable from the identity
947        let peer_id = identity.peer_id().to_hex();
948        assert_eq!(peer_id.len(), 64); // 32 bytes hex-encoded
949    }
950
951    #[tokio::test]
952    async fn test_resolve_identity_loads_existing() {
953        let tmp = tempfile::tempdir().unwrap();
954
955        // Generate and save an identity
956        let original = NodeIdentity::generate().unwrap();
957        original
958            .save_to_file(&tmp.path().join(NODE_IDENTITY_FILENAME))
959            .await
960            .unwrap();
961
962        let mut config = NodeConfig {
963            root_dir: tmp.path().to_path_buf(),
964            ..Default::default()
965        };
966
967        let loaded = NodeBuilder::resolve_identity(&mut config).await.unwrap();
968        assert_eq!(loaded.peer_id(), original.peer_id());
969    }
970
971    #[test]
972    fn test_peer_id_hex_length() {
973        let id = saorsa_core::identity::PeerId::from_bytes([0x42; 32]);
974        let hex = id.to_hex();
975        assert_eq!(hex.len(), 64); // 32 bytes = 64 hex chars
976    }
977
978    /// Simulates a node restart: first run creates identity in a scoped subdir
979    /// under `nodes/`, second run discovers and reloads it — `peer_id` must be
980    /// identical and the directory name is the full 64-char hex peer ID.
981    #[tokio::test]
982    async fn test_identity_persisted_across_restarts() {
983        let base_dir = tempfile::tempdir().unwrap();
984        let nodes_dir = base_dir.path().join(NODES_SUBDIR);
985
986        // First "boot": generate identity, save it in nodes/{peer_id}/
987        let identity1 = NodeIdentity::generate().unwrap();
988        let peer_id1 = identity1.peer_id().to_hex();
989        let peer_dir = nodes_dir.join(&peer_id1);
990        std::fs::create_dir_all(&peer_dir).unwrap();
991        identity1
992            .save_to_file(&peer_dir.join(NODE_IDENTITY_FILENAME))
993            .await
994            .unwrap();
995
996        // Verify directory name is the full 64-char hex peer ID
997        assert_eq!(peer_id1.len(), 64);
998        assert_eq!(peer_dir.file_name().unwrap().to_string_lossy(), peer_id1);
999
1000        // Second "boot": scan should find and reload the same identity
1001        let identity_dirs = NodeBuilder::scan_identity_dirs(&nodes_dir).unwrap();
1002        assert_eq!(identity_dirs.len(), 1);
1003        let loaded = NodeIdentity::load_from_file(&identity_dirs[0].join(NODE_IDENTITY_FILENAME))
1004            .await
1005            .unwrap();
1006        let peer_id2 = loaded.peer_id().to_hex();
1007
1008        assert_eq!(peer_id1, peer_id2, "peer_id must survive restart");
1009        assert_eq!(
1010            identity_dirs[0], peer_dir,
1011            "root_dir must be the same directory"
1012        );
1013    }
1014
1015    /// When two identity subdirs exist under `nodes/`, the scan finds multiple
1016    /// and the resolve path would error asking for `--root-dir`.
1017    #[tokio::test]
1018    async fn test_multiple_identities_errors() {
1019        let base_dir = tempfile::tempdir().unwrap();
1020        let nodes_dir = base_dir.path().join(NODES_SUBDIR);
1021
1022        // Create two identity subdirectories under nodes/
1023        for name in &["aaaa", "bbbb"] {
1024            let dir = nodes_dir.join(name);
1025            std::fs::create_dir_all(&dir).unwrap();
1026            let identity = NodeIdentity::generate().unwrap();
1027            identity
1028                .save_to_file(&dir.join(NODE_IDENTITY_FILENAME))
1029                .await
1030                .unwrap();
1031        }
1032
1033        let identity_dirs = NodeBuilder::scan_identity_dirs(&nodes_dir).unwrap();
1034        assert_eq!(identity_dirs.len(), 2, "should find both identity dirs");
1035    }
1036
1037    /// With a non-default `root_dir` (explicit path), the identity is created on
1038    /// first run and reloaded on subsequent runs from the same directory.
1039    #[tokio::test]
1040    async fn test_explicit_root_dir_persists_across_restarts() {
1041        let tmp = tempfile::tempdir().unwrap();
1042
1043        // First boot — non-default root_dir triggers explicit path
1044        let mut config1 = NodeConfig {
1045            root_dir: tmp.path().to_path_buf(),
1046            ..Default::default()
1047        };
1048        let identity1 = NodeBuilder::resolve_identity(&mut config1).await.unwrap();
1049
1050        // Second boot — same dir
1051        let mut config2 = NodeConfig {
1052            root_dir: tmp.path().to_path_buf(),
1053            ..Default::default()
1054        };
1055        let identity2 = NodeBuilder::resolve_identity(&mut config2).await.unwrap();
1056
1057        assert_eq!(
1058            identity1.peer_id(),
1059            identity2.peer_id(),
1060            "explicit --root-dir must yield stable identity"
1061        );
1062    }
1063}