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