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