Skip to main content

iota_swarm/memory/
swarm.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    collections::HashMap,
7    net::SocketAddr,
8    num::NonZeroUsize,
9    ops,
10    path::{Path, PathBuf},
11};
12
13use anyhow::{Context, Result};
14use futures::future::try_join_all;
15use iota_config::{
16    ExecutionCacheConfig, IOTA_GENESIS_FILENAME, NodeConfig,
17    node::{AuthorityOverloadConfig, GrpcApiConfig, RunWithRange},
18    p2p::DiscoveryConfig,
19    transaction_deny_config::TransactionDenyConfig,
20};
21use iota_macros::nondeterministic;
22use iota_names::config::IotaNamesConfig;
23use iota_node::IotaNodeHandle;
24use iota_protocol_config::{Chain, ProtocolVersion};
25use iota_swarm_config::{
26    genesis_config::{AccountConfig, GenesisConfig, ValidatorGenesisConfig},
27    network_config::NetworkConfig,
28    network_config_builder::{
29        CommitteeConfig, ConfigBuilder, GlobalStateHashV1EnabledConfig, ProtocolVersionsConfig,
30        SupportedProtocolVersionsCallback,
31    },
32    node_config_builder::FullnodeConfigBuilder,
33    node_config_override::{
34        NodeConfigOverride, OverrideScope, apply_node_config_overrides,
35        check_validator_override_scopes, overrides_for_fullnode, overrides_for_validator,
36    },
37};
38use iota_types::{
39    base_types::AuthorityName,
40    object::Object,
41    supported_protocol_versions::SupportedProtocolVersions,
42    traffic_control::{PolicyConfig, RemoteFirewallConfig},
43};
44use rand::rngs::OsRng;
45use tempfile::TempDir;
46use tracing::info;
47
48use super::Node;
49
50pub struct SwarmBuilder<R = OsRng> {
51    rng: R,
52    // template: NodeConfig,
53    dir: Option<PathBuf>,
54    committee: CommitteeConfig,
55    genesis_config: Option<GenesisConfig>,
56    network_config: Option<NetworkConfig>,
57    chain_override: Option<Chain>,
58    additional_objects: Vec<Object>,
59    fullnode_count: usize,
60    fullnode_db_path: Option<PathBuf>,
61    fullnode_rpc_port: Option<u16>,
62    fullnode_rpc_addr: Option<SocketAddr>,
63    supported_protocol_versions_config: ProtocolVersionsConfig,
64    // Default to supported_protocol_versions_config, but can be overridden.
65    fullnode_supported_protocol_versions_config: Option<ProtocolVersionsConfig>,
66    num_unpruned_validators: Option<usize>,
67    authority_overload_config: Option<AuthorityOverloadConfig>,
68    transaction_deny_config: Option<TransactionDenyConfig>,
69    execution_cache_config: Option<ExecutionCacheConfig>,
70    data_ingestion_dir: Option<PathBuf>,
71    fullnode_run_with_range: Option<RunWithRange>,
72    validator_policy_config: Option<PolicyConfig>,
73    fullnode_policy_config: Option<PolicyConfig>,
74    fullnode_fw_config: Option<RemoteFirewallConfig>,
75    max_submit_position: Option<usize>,
76    submit_delay_step_override_millis: Option<u64>,
77    global_state_hash_v1_enabled_config: GlobalStateHashV1EnabledConfig,
78    disable_fullnode_pruning: bool,
79    iota_names_config: Option<IotaNamesConfig>,
80    fullnode_enable_grpc_api: bool,
81    fullnode_grpc_api_config: Option<GrpcApiConfig>,
82    disable_address_verification_cooldown: bool,
83    deterministic_validator_port_base: Option<u16>,
84    fullnode_genesis_config: Option<ValidatorGenesisConfig>,
85    node_config_overrides: Vec<NodeConfigOverride>,
86}
87
88impl SwarmBuilder {
89    #[expect(clippy::new_without_default)]
90    pub fn new() -> Self {
91        Self {
92            rng: OsRng,
93            dir: None,
94            committee: CommitteeConfig::Size(NonZeroUsize::new(1).unwrap()),
95            genesis_config: None,
96            network_config: None,
97            chain_override: None,
98            additional_objects: vec![],
99            fullnode_count: 0,
100            fullnode_db_path: None,
101            fullnode_rpc_port: None,
102            fullnode_rpc_addr: None,
103            supported_protocol_versions_config: ProtocolVersionsConfig::Default,
104            fullnode_supported_protocol_versions_config: None,
105            num_unpruned_validators: None,
106            authority_overload_config: None,
107            transaction_deny_config: None,
108            execution_cache_config: None,
109            data_ingestion_dir: None,
110            fullnode_run_with_range: None,
111            validator_policy_config: None,
112            fullnode_policy_config: None,
113            fullnode_fw_config: None,
114            max_submit_position: None,
115            submit_delay_step_override_millis: None,
116            global_state_hash_v1_enabled_config: GlobalStateHashV1EnabledConfig::Global(true),
117            disable_fullnode_pruning: false,
118            iota_names_config: None,
119            fullnode_enable_grpc_api: false,
120            fullnode_grpc_api_config: None,
121            disable_address_verification_cooldown: false,
122            deterministic_validator_port_base: None,
123            fullnode_genesis_config: None,
124            node_config_overrides: vec![],
125        }
126    }
127}
128
129impl<R> SwarmBuilder<R> {
130    pub fn rng<N: rand::RngCore + rand::CryptoRng>(self, rng: N) -> SwarmBuilder<N> {
131        SwarmBuilder {
132            rng,
133            dir: self.dir,
134            committee: self.committee,
135            genesis_config: self.genesis_config,
136            network_config: self.network_config,
137            chain_override: self.chain_override,
138            additional_objects: self.additional_objects,
139            fullnode_count: self.fullnode_count,
140            fullnode_db_path: self.fullnode_db_path,
141            fullnode_rpc_port: self.fullnode_rpc_port,
142            fullnode_rpc_addr: self.fullnode_rpc_addr,
143            supported_protocol_versions_config: self.supported_protocol_versions_config,
144            fullnode_supported_protocol_versions_config: self
145                .fullnode_supported_protocol_versions_config,
146            num_unpruned_validators: self.num_unpruned_validators,
147            authority_overload_config: self.authority_overload_config,
148            transaction_deny_config: self.transaction_deny_config,
149            execution_cache_config: self.execution_cache_config,
150            data_ingestion_dir: self.data_ingestion_dir,
151            fullnode_run_with_range: self.fullnode_run_with_range,
152            validator_policy_config: self.validator_policy_config,
153            fullnode_policy_config: self.fullnode_policy_config,
154            fullnode_fw_config: self.fullnode_fw_config,
155            max_submit_position: self.max_submit_position,
156            submit_delay_step_override_millis: self.submit_delay_step_override_millis,
157            global_state_hash_v1_enabled_config: self.global_state_hash_v1_enabled_config,
158            disable_fullnode_pruning: self.disable_fullnode_pruning,
159            iota_names_config: self.iota_names_config,
160            fullnode_enable_grpc_api: self.fullnode_enable_grpc_api,
161            fullnode_grpc_api_config: self.fullnode_grpc_api_config,
162            disable_address_verification_cooldown: self.disable_address_verification_cooldown,
163            deterministic_validator_port_base: self.deterministic_validator_port_base,
164            fullnode_genesis_config: self.fullnode_genesis_config,
165            node_config_overrides: self.node_config_overrides,
166        }
167    }
168
169    /// Set the directory that should be used by the Swarm for any on-disk data.
170    ///
171    /// If a directory is provided, it will not be cleaned up when the Swarm is
172    /// dropped.
173    ///
174    /// Defaults to using a temporary directory that will be cleaned up when the
175    /// Swarm is dropped.
176    pub fn dir<P: Into<PathBuf>>(mut self, dir: P) -> Self {
177        self.dir = Some(dir.into());
178        self
179    }
180
181    /// Set the committee size (the number of validators in the validator set).
182    ///
183    /// Defaults to 1.
184    pub fn committee_size(mut self, committee_size: NonZeroUsize) -> Self {
185        self.committee = CommitteeConfig::Size(committee_size);
186        self
187    }
188
189    pub fn with_validators(mut self, validators: Vec<ValidatorGenesisConfig>) -> Self {
190        self.committee = CommitteeConfig::Validators(validators);
191        self
192    }
193
194    /// Lay the generated validators out as
195    /// [`ConfigBuilder::with_deterministic_ports`] describes.
196    ///
197    /// Has no effect when the validators come from a network config or from
198    /// `with_validators`.
199    pub fn with_deterministic_validator_ports(mut self, port_base: u16) -> Self {
200        self.deterministic_validator_port_base = Some(port_base);
201        self
202    }
203
204    /// Take the first fullnode's key pairs and addresses from
205    /// `fullnode_genesis_config` instead of generating them. This gives it the
206    /// same config, and the same db path, on every build.
207    ///
208    /// Further fullnodes keep generated key pairs and addresses, since an
209    /// address can only be used once.
210    pub fn with_fullnode_genesis_config(
211        mut self,
212        fullnode_genesis_config: ValidatorGenesisConfig,
213    ) -> Self {
214        self.fullnode_genesis_config = Some(fullnode_genesis_config);
215        self
216    }
217
218    pub fn with_genesis_config(mut self, genesis_config: GenesisConfig) -> Self {
219        assert!(self.network_config.is_none() && self.genesis_config.is_none());
220        self.genesis_config = Some(genesis_config);
221        self
222    }
223
224    pub fn with_chain_override(mut self, chain: Chain) -> Self {
225        assert!(self.chain_override.is_none());
226        self.chain_override = Some(chain);
227        self
228    }
229
230    pub fn with_num_unpruned_validators(mut self, n: usize) -> Self {
231        assert!(self.network_config.is_none());
232        self.num_unpruned_validators = Some(n);
233        self
234    }
235
236    pub fn with_network_config(mut self, network_config: NetworkConfig) -> Self {
237        assert!(self.network_config.is_none() && self.genesis_config.is_none());
238        self.network_config = Some(network_config);
239        self
240    }
241
242    pub fn with_accounts(mut self, accounts: Vec<AccountConfig>) -> Self {
243        self.get_or_init_genesis_config().accounts = accounts;
244        self
245    }
246
247    pub fn with_objects<I: IntoIterator<Item = Object>>(mut self, objects: I) -> Self {
248        self.additional_objects.extend(objects);
249        self
250    }
251
252    pub fn with_fullnode_count(mut self, fullnode_count: usize) -> Self {
253        self.fullnode_count = fullnode_count;
254        self
255    }
256
257    pub fn with_fullnode_db_path(mut self, fullnode_db_path: PathBuf) -> Self {
258        self.fullnode_db_path = Some(fullnode_db_path);
259        self
260    }
261
262    pub fn with_fullnode_rpc_port(mut self, fullnode_rpc_port: u16) -> Self {
263        assert!(self.fullnode_rpc_addr.is_none());
264        self.fullnode_rpc_port = Some(fullnode_rpc_port);
265        self
266    }
267
268    pub fn with_fullnode_rpc_addr(mut self, fullnode_rpc_addr: SocketAddr) -> Self {
269        assert!(self.fullnode_rpc_port.is_none());
270        self.fullnode_rpc_addr = Some(fullnode_rpc_addr);
271        self
272    }
273
274    pub fn with_epoch_duration_ms(mut self, epoch_duration_ms: u64) -> Self {
275        self.get_or_init_genesis_config()
276            .parameters
277            .epoch_duration_ms = epoch_duration_ms;
278        self
279    }
280
281    pub fn with_protocol_version(mut self, v: ProtocolVersion) -> Self {
282        self.get_or_init_genesis_config()
283            .parameters
284            .protocol_version = v;
285        self
286    }
287
288    pub fn with_supported_protocol_versions(mut self, c: SupportedProtocolVersions) -> Self {
289        self.supported_protocol_versions_config = ProtocolVersionsConfig::Global(c);
290        self
291    }
292
293    pub fn with_supported_protocol_version_callback(
294        mut self,
295        func: SupportedProtocolVersionsCallback,
296    ) -> Self {
297        self.supported_protocol_versions_config = ProtocolVersionsConfig::PerValidator(func);
298        self
299    }
300
301    pub fn with_supported_protocol_versions_config(mut self, c: ProtocolVersionsConfig) -> Self {
302        self.supported_protocol_versions_config = c;
303        self
304    }
305
306    pub fn with_global_state_hash_v1_enabled_config(
307        mut self,
308        c: GlobalStateHashV1EnabledConfig,
309    ) -> Self {
310        self.global_state_hash_v1_enabled_config = c;
311        self
312    }
313
314    pub fn with_fullnode_supported_protocol_versions_config(
315        mut self,
316        c: ProtocolVersionsConfig,
317    ) -> Self {
318        self.fullnode_supported_protocol_versions_config = Some(c);
319        self
320    }
321
322    pub fn with_authority_overload_config(
323        mut self,
324        authority_overload_config: AuthorityOverloadConfig,
325    ) -> Self {
326        assert!(self.network_config.is_none());
327        self.authority_overload_config = Some(authority_overload_config);
328        self
329    }
330
331    pub fn with_transaction_deny_config(
332        mut self,
333        transaction_deny_config: TransactionDenyConfig,
334    ) -> Self {
335        assert!(self.network_config.is_none());
336        self.transaction_deny_config = Some(transaction_deny_config);
337        self
338    }
339
340    pub fn with_execution_cache_config(
341        mut self,
342        execution_cache_config: ExecutionCacheConfig,
343    ) -> Self {
344        self.execution_cache_config = Some(execution_cache_config);
345        self
346    }
347
348    pub fn with_data_ingestion_dir(mut self, path: PathBuf) -> Self {
349        self.data_ingestion_dir = Some(path);
350        self
351    }
352
353    pub fn with_fullnode_run_with_range(mut self, run_with_range: Option<RunWithRange>) -> Self {
354        if let Some(run_with_range) = run_with_range {
355            self.fullnode_run_with_range = Some(run_with_range);
356        }
357        self
358    }
359
360    /// Set the traffic control policy of every validator, whether the
361    /// committee is generated here or taken from a network config.
362    pub fn with_validator_policy_config(mut self, config: Option<PolicyConfig>) -> Self {
363        self.validator_policy_config = config;
364        self
365    }
366
367    pub fn with_fullnode_policy_config(mut self, config: Option<PolicyConfig>) -> Self {
368        self.fullnode_policy_config = config;
369        self
370    }
371
372    pub fn with_fullnode_fw_config(mut self, config: Option<RemoteFirewallConfig>) -> Self {
373        self.fullnode_fw_config = config;
374        self
375    }
376
377    pub fn with_fullnode_enable_grpc_api(mut self, enable: bool) -> Self {
378        self.fullnode_enable_grpc_api = enable;
379        self
380    }
381
382    pub fn with_fullnode_grpc_api_config(mut self, config: GrpcApiConfig) -> Self {
383        self.fullnode_grpc_api_config = Some(config);
384        self
385    }
386
387    fn get_or_init_genesis_config(&mut self) -> &mut GenesisConfig {
388        if self.genesis_config.is_none() {
389            assert!(self.network_config.is_none());
390            self.genesis_config = Some(GenesisConfig::for_local_testing());
391        }
392        self.genesis_config.as_mut().unwrap()
393    }
394
395    pub fn with_max_submit_position(mut self, max_submit_position: usize) -> Self {
396        self.max_submit_position = Some(max_submit_position);
397        self
398    }
399
400    pub fn with_disable_fullnode_pruning(mut self) -> Self {
401        self.disable_fullnode_pruning = true;
402        self
403    }
404
405    pub fn with_submit_delay_step_override_millis(
406        mut self,
407        submit_delay_step_override_millis: u64,
408    ) -> Self {
409        self.submit_delay_step_override_millis = Some(submit_delay_step_override_millis);
410        self
411    }
412
413    pub fn with_iota_names_config(mut self, iota_names_config: IotaNamesConfig) -> Self {
414        self.iota_names_config = Some(iota_names_config);
415        self
416    }
417
418    /// Disable address verification cooldown for test environments where nodes
419    /// frequently restart. This prevents nodes from being blocked from
420    /// reconnecting after crashes/restarts.
421    pub fn with_disabled_address_verification_cooldown(mut self) -> Self {
422        self.disable_address_verification_cooldown = true;
423        self
424    }
425
426    /// Set overrides applied to every node config this builder produces, in
427    /// the given order, after all other configuration. Nodes spawned on the
428    /// built [`Swarm`] later get them too, except `validator-<N>` scoped
429    /// overrides, which refer to positions in the initial network config.
430    pub fn with_node_config_overrides(
431        mut self,
432        node_config_overrides: Vec<NodeConfigOverride>,
433    ) -> Self {
434        self.node_config_overrides = node_config_overrides;
435        self
436    }
437}
438
439impl<R: rand::RngCore + rand::CryptoRng> SwarmBuilder<R> {
440    /// Create the configured Swarm.
441    ///
442    /// # Panics
443    ///
444    /// Panics if [`SwarmBuilder::try_build`] returns an error.
445    pub fn build(self) -> Swarm {
446        self.try_build().unwrap_or_else(|err| panic!("{err:#}"))
447    }
448
449    /// Create the configured Swarm.
450    ///
451    /// # Errors
452    ///
453    /// - A `validator-<N>` override names a validator the network does not
454    ///   have.
455    /// - An override fails to apply to a built config.
456    /// - The network has a fullnode and a validator config has no
457    ///   `p2p-config.external-address`.
458    ///
459    /// # Panics
460    ///
461    /// Panics on failures the swarm cannot run without: creating its temporary
462    /// directory, saving the genesis blob, parsing a generated network address,
463    /// and building the genesis (e.g. on invalid genesis parameters or a
464    /// validator below the minimum stake).
465    pub fn try_build(mut self) -> Result<Swarm> {
466        let mut fullnode_genesis_config = self.fullnode_genesis_config.take();
467        let dir = if let Some(dir) = self.dir {
468            SwarmDirectory::Persistent(dir)
469        } else {
470            SwarmDirectory::new_temporary()
471        };
472
473        let ingest_data = self.data_ingestion_dir.clone();
474
475        let mut network_config = self.network_config.unwrap_or_else(|| {
476            let mut config_builder = ConfigBuilder::new(dir.as_ref());
477
478            if let Some(genesis_config) = self.genesis_config {
479                config_builder = config_builder.with_genesis_config(genesis_config);
480            }
481
482            if let Some(chain_override) = self.chain_override {
483                config_builder = config_builder.with_chain_override(chain_override);
484            }
485
486            if let Some(num_unpruned_validators) = self.num_unpruned_validators {
487                config_builder =
488                    config_builder.with_num_unpruned_validators(num_unpruned_validators);
489            }
490
491            if let Some(authority_overload_config) = self.authority_overload_config {
492                config_builder =
493                    config_builder.with_authority_overload_config(authority_overload_config);
494            }
495
496            if let Some(transaction_deny_config) = self.transaction_deny_config {
497                config_builder =
498                    config_builder.with_transaction_deny_config(transaction_deny_config);
499            }
500
501            if let Some(execution_cache_config) = self.execution_cache_config {
502                config_builder = config_builder.with_execution_cache_config(execution_cache_config);
503            }
504
505            if let Some(path) = self.data_ingestion_dir {
506                config_builder = config_builder.with_data_ingestion_dir(path);
507            }
508
509            if let Some(port_base) = self.deterministic_validator_port_base {
510                config_builder = config_builder.with_deterministic_ports(port_base);
511            }
512
513            if let Some(max_submit_position) = self.max_submit_position {
514                config_builder = config_builder.with_max_submit_position(max_submit_position);
515            }
516
517            if let Some(submit_delay_step_override_millis) = self.submit_delay_step_override_millis
518            {
519                config_builder = config_builder
520                    .with_submit_delay_step_override_millis(submit_delay_step_override_millis);
521            }
522
523            let mut network_config = config_builder
524                .committee(self.committee)
525                .rng(self.rng)
526                .with_objects(self.additional_objects)
527                .with_empty_validator_genesis()
528                .with_supported_protocol_versions_config(
529                    self.supported_protocol_versions_config.clone(),
530                )
531                .with_global_state_hash_v1_enabled_config(
532                    self.global_state_hash_v1_enabled_config.clone(),
533                )
534                .build();
535            // Populate validator genesis by pointing to the blob
536            let genesis_path = dir.join(IOTA_GENESIS_FILENAME);
537            network_config
538                .genesis
539                .save(&genesis_path)
540                .expect("genesis should be saved successfully");
541            for validator in &mut network_config.validator_configs {
542                validator.genesis = iota_config::node::Genesis::new_from_file(&genesis_path);
543            }
544            network_config
545        });
546
547        if let Some(policy_config) = self.validator_policy_config {
548            for validator in &mut network_config.validator_configs {
549                validator.policy_config = Some(policy_config.clone());
550            }
551        }
552
553        if self.disable_address_verification_cooldown {
554            for validator in &mut network_config.validator_configs {
555                if let Some(ref mut discovery_config) = validator.p2p_config.discovery {
556                    discovery_config.address_verification_failure_cooldown_sec = Some(0);
557                } else {
558                    validator.p2p_config.discovery = Some(DiscoveryConfig {
559                        address_verification_failure_cooldown_sec: Some(0),
560                        ..Default::default()
561                    });
562                }
563            }
564        }
565
566        check_validator_override_scopes(
567            &self.node_config_overrides,
568            network_config.validator_configs.len(),
569        )?;
570        for (index, validator) in network_config.validator_configs.iter_mut().enumerate() {
571            apply_node_config_overrides(
572                overrides_for_validator(&self.node_config_overrides, index),
573                validator,
574            )
575            .with_context(|| {
576                format!("failed to apply node config overrides to validator {index}")
577            })?;
578        }
579
580        let mut nodes: HashMap<_, _> = network_config
581            .validator_configs()
582            .iter()
583            .map(|config| {
584                info!(
585                    "SwarmBuilder configuring validator with name {}",
586                    config.authority_public_key()
587                );
588                (config.authority_public_key(), Node::new(config.to_owned()))
589            })
590            .collect();
591
592        let mut fullnode_config_builder = FullnodeConfigBuilder::new()
593            .with_config_directory(dir.as_ref().into())
594            .with_run_with_range(self.fullnode_run_with_range)
595            .with_policy_config(self.fullnode_policy_config)
596            .with_data_ingestion_dir(ingest_data)
597            .with_fw_config(self.fullnode_fw_config)
598            .with_disable_pruning(self.disable_fullnode_pruning)
599            .with_iota_names_config(self.iota_names_config);
600        if let Some(fullnode_db_path) = self.fullnode_db_path {
601            fullnode_config_builder = fullnode_config_builder.with_db_path(fullnode_db_path);
602        }
603
604        if self.disable_address_verification_cooldown {
605            let discovery_config = DiscoveryConfig {
606                address_verification_failure_cooldown_sec: Some(0),
607                ..Default::default()
608            };
609
610            fullnode_config_builder =
611                fullnode_config_builder.with_discovery_config(discovery_config);
612        }
613
614        if let Some(chain) = self.chain_override {
615            fullnode_config_builder = fullnode_config_builder.with_chain_override(chain);
616        }
617
618        if let Some(spvc) = &self.fullnode_supported_protocol_versions_config {
619            let supported_versions = match spvc {
620                ProtocolVersionsConfig::Default => SupportedProtocolVersions::SYSTEM_DEFAULT,
621                ProtocolVersionsConfig::Global(v) => *v,
622                ProtocolVersionsConfig::PerValidator(func) => func(0, None),
623            };
624            fullnode_config_builder =
625                fullnode_config_builder.with_supported_protocol_versions(supported_versions);
626        }
627
628        // Add gRPC config wiring
629        fullnode_config_builder =
630            fullnode_config_builder.with_enable_grpc_api(self.fullnode_enable_grpc_api);
631        if let Some(grpc_config) = &self.fullnode_grpc_api_config {
632            fullnode_config_builder =
633                fullnode_config_builder.with_grpc_api_config(grpc_config.clone());
634        }
635
636        for idx in 0..self.fullnode_count {
637            let mut builder = fullnode_config_builder.clone();
638            // Only the first fullnode is used as the rpc fullnode, and only it
639            // takes the given genesis config: an address can only be used once.
640            let genesis_config = if idx == 0 {
641                if let Some(rpc_addr) = self.fullnode_rpc_addr {
642                    builder = builder.with_rpc_addr(rpc_addr);
643                }
644                if let Some(rpc_port) = self.fullnode_rpc_port {
645                    builder = builder.with_rpc_port(rpc_port);
646                }
647                fullnode_genesis_config.take()
648            } else {
649                None
650            };
651            let mut config = match genesis_config {
652                Some(genesis_config) => {
653                    builder.try_build_with_genesis_config(genesis_config, &network_config)
654                }
655                None => builder.try_build(&mut OsRng, &network_config),
656            }
657            .context("failed to build the fullnode config")?;
658            apply_node_config_overrides(
659                overrides_for_fullnode(&self.node_config_overrides),
660                &mut config,
661            )
662            .with_context(|| format!("failed to apply node config overrides to fullnode {idx}"))?;
663            info!(
664                "SwarmBuilder configuring full node with name {}",
665                config.authority_public_key()
666            );
667            nodes.insert(config.authority_public_key(), Node::new(config));
668        }
669        Ok(Swarm {
670            dir,
671            network_config,
672            nodes,
673            fullnode_config_builder,
674            node_config_overrides: self.node_config_overrides,
675        })
676    }
677}
678
679/// A handle to an in-memory IOTA Network.
680#[derive(Debug)]
681pub struct Swarm {
682    dir: SwarmDirectory,
683    network_config: NetworkConfig,
684    nodes: HashMap<AuthorityName, Node>,
685    // Save a copy of the fullnode config builder to build future fullnodes.
686    fullnode_config_builder: FullnodeConfigBuilder,
687    // Applied to the configs of nodes spawned after the initial build too.
688    node_config_overrides: Vec<NodeConfigOverride>,
689}
690
691impl Drop for Swarm {
692    fn drop(&mut self) {
693        self.nodes_iter_mut().for_each(|node| node.stop());
694    }
695}
696
697impl Swarm {
698    fn nodes_iter_mut(&mut self) -> impl Iterator<Item = &mut Node> {
699        self.nodes.values_mut()
700    }
701
702    /// Return a new Builder
703    pub fn builder() -> SwarmBuilder {
704        SwarmBuilder::new()
705    }
706
707    /// Start all nodes associated with this Swarm
708    pub async fn launch(&mut self) -> Result<()> {
709        try_join_all(self.nodes_iter_mut().map(|node| node.start())).await?;
710        tracing::info!("Successfully launched Swarm");
711        Ok(())
712    }
713
714    /// Return the path to the directory where this Swarm's on-disk data is
715    /// kept.
716    pub fn dir(&self) -> &Path {
717        self.dir.as_ref()
718    }
719
720    /// Return a reference to this Swarm's `NetworkConfig`.
721    pub fn config(&self) -> &NetworkConfig {
722        &self.network_config
723    }
724
725    /// Return a mutable reference to this Swarm's `NetworkConfig`.
726    // TODO: It's not ideal to mutate network config. We should consider removing
727    // this.
728    pub fn config_mut(&mut self) -> &mut NetworkConfig {
729        &mut self.network_config
730    }
731
732    pub fn all_nodes(&self) -> impl Iterator<Item = &Node> {
733        self.nodes.values()
734    }
735
736    pub fn node(&self, name: &AuthorityName) -> Option<&Node> {
737        self.nodes.get(name)
738    }
739
740    pub fn node_mut(&mut self, name: &AuthorityName) -> Option<&mut Node> {
741        self.nodes.get_mut(name)
742    }
743
744    /// Return an iterator over shared references of all nodes that are set up
745    /// as validators. This means that they have a consensus config. This
746    /// however doesn't mean this validator is currently active (i.e. it's
747    /// not necessarily in the validator set at the moment).
748    pub fn validator_nodes(&self) -> impl Iterator<Item = &Node> {
749        self.nodes
750            .values()
751            .filter(|node| node.config().is_validator())
752    }
753
754    pub fn validator_node_handles(&self) -> Vec<IotaNodeHandle> {
755        self.validator_nodes()
756            .map(|node| node.get_node_handle().unwrap())
757            .collect()
758    }
759
760    /// Returns an iterator over all current active validators.
761    pub fn active_validators(&self) -> impl Iterator<Item = &Node> {
762        self.validator_nodes().filter(|node| {
763            node.get_node_handle().is_some_and(|handle| {
764                let state = handle.state();
765                state.is_active_validator(&state.epoch_store_for_testing())
766            })
767        })
768    }
769
770    /// Returns an iterator over all current active validators.
771    pub fn committee_validators(&self) -> impl Iterator<Item = &Node> {
772        self.validator_nodes().filter(|node| {
773            node.get_node_handle().is_some_and(|handle| {
774                let state = handle.state();
775                state.is_committee_validator(&state.epoch_store_for_testing())
776            })
777        })
778    }
779
780    /// Return an iterator over shared references of all Fullnodes.
781    pub fn fullnodes(&self) -> impl Iterator<Item = &Node> {
782        self.nodes
783            .values()
784            .filter(|node| !node.config().is_validator())
785    }
786
787    /// Start a node from `config` and add it to the swarm.
788    ///
789    /// The swarm's node config overrides are applied to the config first.
790    ///
791    /// # Panics
792    ///
793    /// Panics on an override that fails to apply and on a node that fails
794    /// to start.
795    pub async fn spawn_new_node(&mut self, mut config: NodeConfig) -> IotaNodeHandle {
796        self.apply_node_config_overrides_for_spawn(&mut config);
797        let name = config.authority_public_key();
798        let node = Node::new(config);
799        node.start().await.unwrap();
800        let handle = node.get_node_handle().unwrap();
801        self.nodes.insert(name, node);
802        handle
803    }
804
805    /// Apply the swarm's overrides to the config of a node spawned after the
806    /// initial build. `validator-<N>` scoped overrides refer to positions in
807    /// the initial network config, so they are skipped here.
808    ///
809    /// # Panics
810    ///
811    /// Panics on an override that fails to apply.
812    fn apply_node_config_overrides_for_spawn(&self, config: &mut NodeConfig) {
813        let overrides: Vec<&NodeConfigOverride> = if config.is_validator() {
814            self.node_config_overrides
815                .iter()
816                .filter(|config_override| {
817                    matches!(
818                        config_override.scope,
819                        OverrideScope::All | OverrideScope::AllValidators
820                    )
821                })
822                .collect()
823        } else {
824            overrides_for_fullnode(&self.node_config_overrides).collect()
825        };
826        apply_node_config_overrides(overrides, config).unwrap_or_else(|err| panic!("{err:#}"));
827    }
828
829    pub fn get_fullnode_config_builder(&self) -> FullnodeConfigBuilder {
830        self.fullnode_config_builder.clone()
831    }
832
833    /// The node config overrides the swarm was built with.
834    pub fn node_config_overrides(&self) -> &[NodeConfigOverride] {
835        &self.node_config_overrides
836    }
837}
838
839#[derive(Debug)]
840enum SwarmDirectory {
841    Persistent(PathBuf),
842    Temporary(TempDir),
843}
844
845impl SwarmDirectory {
846    fn new_temporary() -> Self {
847        SwarmDirectory::Temporary(nondeterministic!(TempDir::new().unwrap()))
848    }
849}
850
851impl ops::Deref for SwarmDirectory {
852    type Target = Path;
853
854    fn deref(&self) -> &Self::Target {
855        match self {
856            SwarmDirectory::Persistent(dir) => dir.deref(),
857            SwarmDirectory::Temporary(dir) => dir.path(),
858        }
859    }
860}
861
862impl AsRef<Path> for SwarmDirectory {
863    fn as_ref(&self) -> &Path {
864        match self {
865            SwarmDirectory::Persistent(dir) => dir.as_ref(),
866            SwarmDirectory::Temporary(dir) => dir.as_ref(),
867        }
868    }
869}
870
871#[cfg(test)]
872mod test {
873    use std::{collections::BTreeSet, num::NonZeroUsize};
874
875    use iota_swarm_config::{
876        genesis_config::ValidatorGenesisConfigBuilder,
877        network_config::NetworkConfig,
878        network_config_builder::ConfigBuilder,
879        node_config_override::{NodeConfigOverride, apply_node_config_overrides},
880    };
881    use iota_types::traffic_control::PolicyConfig;
882
883    use super::Swarm;
884
885    #[test]
886    fn the_validator_policy_config_applies_before_the_overrides() {
887        let policy_config = PolicyConfig {
888            connection_blocklist_ttl_sec: 4242,
889            ..PolicyConfig::default()
890        };
891        let swarm = Swarm::builder()
892            .committee_size(NonZeroUsize::new(2).unwrap())
893            .with_validator_policy_config(Some(policy_config))
894            .with_node_config_overrides(vec!["validator-0:policy-config=".parse().unwrap()])
895            .build();
896
897        let validators = swarm.config().validator_configs();
898        assert!(validators[0].policy_config.is_none());
899        assert_eq!(
900            validators[1]
901                .policy_config
902                .as_ref()
903                .unwrap()
904                .connection_blocklist_ttl_sec,
905            4242
906        );
907    }
908
909    #[test]
910    fn node_config_overrides() {
911        let swarm = Swarm::builder()
912            .committee_size(NonZeroUsize::new(2).unwrap())
913            .with_fullnode_count(1)
914            .with_node_config_overrides(vec![
915                "fullnode:authority-store-pruning-config.num-epochs-to-retain=18446744073709551615"
916                    .parse()
917                    .unwrap(),
918                "validator-0:authority-store-pruning-config.num-epochs-to-retain=5"
919                    .parse()
920                    .unwrap(),
921                "validator:enable-soft-locking=false".parse().unwrap(),
922            ])
923            .build();
924
925        let validators = swarm.config().validator_configs();
926        assert_eq!(
927            validators[0]
928                .authority_store_pruning_config
929                .num_epochs_to_retain,
930            5
931        );
932        assert_eq!(
933            validators[1]
934                .authority_store_pruning_config
935                .num_epochs_to_retain,
936            0
937        );
938        assert!(validators.iter().all(|config| !config.enable_soft_locking));
939
940        let fullnode = swarm.fullnodes().next().unwrap();
941        assert_eq!(
942            fullnode
943                .config()
944                .authority_store_pruning_config
945                .num_epochs_to_retain,
946            u64::MAX
947        );
948        assert!(fullnode.config().enable_soft_locking);
949    }
950
951    #[test]
952    fn node_config_overrides_apply_to_late_spawned_nodes() {
953        let swarm = Swarm::builder()
954            .committee_size(NonZeroUsize::new(2).unwrap())
955            .with_fullnode_count(1)
956            .with_node_config_overrides(vec![
957                "fullnode:authority-store-pruning-config.num-epochs-to-retain=18446744073709551615"
958                    .parse()
959                    .unwrap(),
960                "validator:enable-soft-locking=false".parse().unwrap(),
961                "validator-0:enable-index-processing=false".parse().unwrap(),
962            ])
963            .build();
964
965        let mut config = swarm
966            .get_fullnode_config_builder()
967            .build(&mut rand::rngs::OsRng, swarm.config());
968        assert_eq!(
969            config.authority_store_pruning_config.num_epochs_to_retain,
970            0
971        );
972        swarm.apply_node_config_overrides_for_spawn(&mut config);
973        assert_eq!(
974            config.authority_store_pruning_config.num_epochs_to_retain,
975            u64::MAX
976        );
977        // Validator-scoped overrides do not apply to a fullnode.
978        assert!(config.enable_soft_locking);
979
980        // A validator respawned from its own config: the batch it was built
981        // with applies again unchanged.
982        let mut config = swarm.config().validator_configs()[1].clone();
983        let num_epochs_to_retain = config.authority_store_pruning_config.num_epochs_to_retain;
984        assert!(!config.enable_soft_locking);
985        swarm.apply_node_config_overrides_for_spawn(&mut config);
986        assert!(!config.enable_soft_locking);
987        // The `validator-0` and fullnode scopes leave this validator alone.
988        assert!(config.enable_index_processing);
989        assert_eq!(
990            config.authority_store_pruning_config.num_epochs_to_retain,
991            num_epochs_to_retain
992        );
993    }
994
995    #[test]
996    fn try_build_rejects_a_consensus_override_that_reaches_the_fullnode() {
997        // `all:` applies cleanly to the validators and must still fail on
998        // the fullnode, which has no consensus config.
999        let err = Swarm::builder()
1000            .with_fullnode_count(1)
1001            .with_node_config_overrides(vec![
1002                "all:consensus-config.db-retention-epochs=2"
1003                    .parse()
1004                    .unwrap(),
1005            ])
1006            .try_build()
1007            .unwrap_err();
1008        let err = format!("{err:#}");
1009        assert!(
1010            err.contains("all:consensus-config.db-retention-epochs"),
1011            "{err}"
1012        );
1013        assert!(err.contains("on a fullnode"), "{err}");
1014    }
1015
1016    #[test]
1017    fn the_fullnodes_own_addresses_are_overridable() {
1018        // A fullnode is not a committee member, so the addresses that are
1019        // genesis data on a validator are ordinary config on it. The seed
1020        // peers it derives from the validators are unaffected.
1021        let swarm = Swarm::builder()
1022            .committee_size(NonZeroUsize::new(1).unwrap())
1023            .with_fullnode_count(1)
1024            .with_node_config_overrides(vec![
1025                "fullnode:p2p-config.external-address='/ip4/127.0.0.1/udp/19186'"
1026                    .parse()
1027                    .unwrap(),
1028            ])
1029            .try_build()
1030            .unwrap();
1031        let fullnode = swarm.fullnodes().next().unwrap();
1032        let config = fullnode.config();
1033        assert_eq!(
1034            config
1035                .p2p_config
1036                .external_address
1037                .as_ref()
1038                .unwrap()
1039                .to_string(),
1040            "/ip4/127.0.0.1/udp/19186"
1041        );
1042        assert_eq!(
1043            config.p2p_config.seed_peers[0].address,
1044            swarm.config().validator_configs()[0]
1045                .p2p_config
1046                .external_address
1047                .clone()
1048                .unwrap()
1049        );
1050    }
1051
1052    /// A network config whose validator 0 carries a firewall section, which
1053    /// its peers do not. Returns the config and the temporary directory it
1054    /// must outlive.
1055    fn network_config_with_a_firewall_on_validator_0(
1056        committee_size: usize,
1057    ) -> (NetworkConfig, tempfile::TempDir) {
1058        let dir = tempfile::TempDir::new().unwrap();
1059        let mut network_config = ConfigBuilder::new(dir.path())
1060            .committee_size(NonZeroUsize::new(committee_size).unwrap())
1061            .build();
1062        let overrides: Vec<NodeConfigOverride> = [
1063            "policy-config={}",
1064            "firewall-config={remote-fw-url: 'http://127.0.0.1:65000', destination-port: 65000}",
1065        ]
1066        .iter()
1067        .map(|input| input.parse().unwrap())
1068        .collect();
1069        apply_node_config_overrides(&overrides, &mut network_config.validator_configs[0]).unwrap();
1070        (network_config, dir)
1071    }
1072
1073    #[test]
1074    fn validator_override_failures_name_the_validator() {
1075        // A `validator:` scope carries no index, so only the error context
1076        // can say which validator rejected the override.
1077        let (network_config, _dir) = network_config_with_a_firewall_on_validator_0(2);
1078
1079        let err = Swarm::builder()
1080            .with_network_config(network_config)
1081            .with_node_config_overrides(vec![
1082                "validator:firewall-config.destination-port=65001"
1083                    .parse()
1084                    .unwrap(),
1085            ])
1086            .try_build()
1087            .unwrap_err();
1088        // Validator 0 has the section, validator 1 does not. The dotted
1089        // edit therefore leaves its required fields unset on validator 1.
1090        let err = format!("{err:#}");
1091        assert!(err.contains("validator 1"), "{err}");
1092        assert!(err.contains("remote-fw-url"), "{err}");
1093    }
1094
1095    #[test]
1096    fn try_build_rejects_an_out_of_range_validator_scope_for_a_supplied_network_config() {
1097        let dir = tempfile::TempDir::new().unwrap();
1098        let network_config = ConfigBuilder::new(dir.path())
1099            .committee_size(NonZeroUsize::new(1).unwrap())
1100            .build();
1101        let err = Swarm::builder()
1102            .with_network_config(network_config)
1103            .with_node_config_overrides(vec![
1104                "validator-1:enable-soft-locking=false".parse().unwrap(),
1105            ])
1106            .try_build()
1107            .unwrap_err();
1108        let err = format!("{err:#}");
1109        assert!(err.contains("validator-1:enable-soft-locking"), "{err}");
1110        assert!(err.contains("only 1 validator"), "{err}");
1111    }
1112
1113    #[test]
1114    fn validator_scopes_may_set_the_consensus_config() {
1115        let swarm = Swarm::builder()
1116            .with_node_config_overrides(vec![
1117                "validator:consensus-config.db-retention-epochs=2"
1118                    .parse()
1119                    .unwrap(),
1120                "validator-0:consensus-config.db-pruner-period-secs=60"
1121                    .parse()
1122                    .unwrap(),
1123            ])
1124            .try_build()
1125            .unwrap();
1126        let consensus_config = swarm.config().validator_configs()[0]
1127            .consensus_config
1128            .as_ref()
1129            .unwrap();
1130        assert_eq!(consensus_config.db_retention_epochs, Some(2));
1131        assert_eq!(consensus_config.db_pruner_period_secs, Some(60));
1132    }
1133
1134    #[test]
1135    fn overrides_apply_to_a_supplied_network_config() {
1136        // The localnet feeds a network config loaded from disk. Overrides
1137        // apply to those configs, not to freshly generated ones.
1138        let (network_config, _dir) = network_config_with_a_firewall_on_validator_0(1);
1139
1140        let swarm = Swarm::builder()
1141            .with_network_config(network_config)
1142            .with_fullnode_count(1)
1143            .with_node_config_overrides(vec![
1144                "validator:firewall-config.destination-port=65001"
1145                    .parse()
1146                    .unwrap(),
1147                "fullnode:enable-index-processing=false".parse().unwrap(),
1148            ])
1149            .try_build()
1150            .unwrap();
1151        assert_eq!(
1152            swarm.config().validator_configs()[0]
1153                .firewall_config
1154                .as_ref()
1155                .unwrap()
1156                .destination_port,
1157            65001
1158        );
1159        let fullnode = swarm.fullnodes().next().unwrap();
1160        assert!(!fullnode.config().enable_index_processing);
1161    }
1162
1163    #[test]
1164    fn try_build_rejects_a_fullnode_override_no_node_could_start_with() {
1165        let err = Swarm::builder()
1166            .committee_size(NonZeroUsize::new(1).unwrap())
1167            .with_fullnode_count(1)
1168            .with_node_config_overrides(vec![
1169                // A snapshot store without a backend: the node refuses to
1170                // start with it.
1171                "fullnode:state-snapshot-write-config.object-store-config.directory=/tmp/snapshots"
1172                    .parse()
1173                    .unwrap(),
1174            ])
1175            .try_build()
1176            .unwrap_err();
1177        let err = format!("{err:#}");
1178        assert!(err.contains("storage backend"), "{err}");
1179    }
1180
1181    #[test]
1182    fn try_build_fails_when_a_validator_has_no_p2p_external_address() {
1183        // The fullnode derives its seed peers from the validators' external
1184        // addresses.
1185        let dir = tempfile::TempDir::new().unwrap();
1186        let mut network_config = ConfigBuilder::new(dir.path())
1187            .committee_size(NonZeroUsize::new(1).unwrap())
1188            .build();
1189        network_config.validator_configs[0]
1190            .p2p_config
1191            .external_address = None;
1192
1193        let err = Swarm::builder()
1194            .with_network_config(network_config)
1195            .with_fullnode_count(1)
1196            .try_build()
1197            .unwrap_err();
1198        let err = format!("{err:#}");
1199        assert!(err.contains("validator 0"), "{err}");
1200        assert!(err.contains("seed peers"), "{err}");
1201    }
1202
1203    #[tokio::test]
1204    async fn launch() {
1205        telemetry_subscribers::init_for_testing();
1206        let mut swarm = Swarm::builder()
1207            .committee_size(NonZeroUsize::new(4).unwrap())
1208            .with_fullnode_count(1)
1209            .build();
1210
1211        swarm.launch().await.unwrap();
1212
1213        for validator in swarm.validator_nodes() {
1214            validator.health_check(true).await.unwrap();
1215        }
1216
1217        for fullnode in swarm.fullnodes() {
1218            fullnode.health_check(false).await.unwrap();
1219        }
1220
1221        println!("hello");
1222    }
1223
1224    #[test]
1225    fn deterministic_ports_reach_the_node_configs() {
1226        let swarm = Swarm::builder()
1227            .committee_size(NonZeroUsize::new(2).unwrap())
1228            .with_deterministic_validator_ports(9200)
1229            .build();
1230
1231        let validator_ports = swarm
1232            .validator_nodes()
1233            .map(|validator| {
1234                validator
1235                    .config()
1236                    .network_address
1237                    .to_socket_addr()
1238                    .unwrap()
1239                    .port()
1240            })
1241            .collect::<BTreeSet<_>>();
1242        assert_eq!(validator_ports, BTreeSet::from([9200, 9210]));
1243    }
1244
1245    #[test]
1246    fn the_first_fullnode_takes_the_given_genesis_config() {
1247        let mut fullnode_genesis_config = ValidatorGenesisConfigBuilder::new()
1248            .with_ip("127.0.0.1".to_owned())
1249            .build(&mut rand::rngs::OsRng);
1250        fullnode_genesis_config.metrics_address = ([127, 0, 0, 1], 19184).into();
1251        fullnode_genesis_config.admin_interface_address = ([127, 0, 0, 1], 19185).into();
1252        fullnode_genesis_config.p2p_address = "/ip4/127.0.0.1/udp/19186/http".parse().unwrap();
1253        let db_path_of = |swarm: &Swarm| swarm.fullnodes().next().unwrap().config().db_path.clone();
1254
1255        let swarm = Swarm::builder()
1256            .with_fullnode_count(1)
1257            .with_fullnode_genesis_config(fullnode_genesis_config.copy_with_private_keys())
1258            .build();
1259
1260        {
1261            let fullnode = swarm.fullnodes().next().unwrap().config();
1262            assert_eq!(fullnode.metrics_address.to_string(), "127.0.0.1:19184");
1263            assert_eq!(
1264                fullnode.admin_interface_address.to_string(),
1265                "127.0.0.1:19185"
1266            );
1267            assert_eq!(
1268                fullnode.p2p_config.listen_address.to_string(),
1269                "127.0.0.1:19186"
1270            );
1271        }
1272
1273        // The same entry gives the fullnode the same db path in a second
1274        // network, which is what lets a persisted network reuse its database.
1275        let same_swarm = Swarm::builder()
1276            .with_fullnode_count(1)
1277            .with_fullnode_genesis_config(fullnode_genesis_config)
1278            .build();
1279        assert_eq!(
1280            db_path_of(&swarm).file_name(),
1281            db_path_of(&same_swarm).file_name()
1282        );
1283    }
1284}