1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! Convenience structs and functions for generating a random set of Diem ndoes without the
//! genesis.blob.

use crate::{
    config::{
        DiscoveryMethod, NetworkConfig, NodeConfig, Peer, PeerRole, PeerSet, TestConfig,
        HANDSHAKE_VERSION,
    },
    network_id::NetworkId,
};
use rand::{rngs::StdRng, SeedableRng};
use std::collections::{HashMap, HashSet};

pub struct ValidatorSwarm {
    pub nodes: Vec<NodeConfig>,
}

pub fn validator_swarm(
    template: &NodeConfig,
    count: usize,
    seed: [u8; 32],
    randomize_ports: bool,
) -> ValidatorSwarm {
    let mut rng = StdRng::from_seed(seed);
    let mut nodes = Vec::new();

    for index in 0..count {
        let mut node = NodeConfig::random_with_template(index as u32, template, &mut rng);
        if randomize_ports {
            node.randomize_ports();
        }

        // For a validator node, any of its validator peers are considered an upstream peer
        let network = node.validator_network.as_mut().unwrap();
        network.discovery_method = DiscoveryMethod::Onchain;
        network.mutual_authentication = true;
        network.network_id = NetworkId::Validator;

        nodes.push(node);
    }

    // set the first validator as every validators' initial configured seed peer.
    let seed_config = &nodes[0].validator_network.as_ref().unwrap();
    let seeds = build_seed_for_network(seed_config, PeerRole::Validator);
    for node in &mut nodes {
        let network = node.validator_network.as_mut().unwrap();
        network.seeds = seeds.clone();
    }

    ValidatorSwarm { nodes }
}

pub fn validator_swarm_for_testing(nodes: usize) -> ValidatorSwarm {
    let config = NodeConfig {
        test: Some(TestConfig::open_module()),
        ..Default::default()
    };
    validator_swarm(&config, nodes, [1u8; 32], true)
}

/// Convenience function that builds a `PeerSet` containing a single peer for testing
/// with a fully formatted `NetworkAddress` containing its network identity pubkey
/// and handshake protocol version.
pub fn build_seed_for_network(seed_config: &NetworkConfig, seed_role: PeerRole) -> PeerSet {
    let seed_pubkey = diem_crypto::PrivateKey::public_key(&seed_config.identity_key());
    let seed_addr = seed_config
        .listen_address
        .clone()
        .append_prod_protos(seed_pubkey, HANDSHAKE_VERSION);

    let mut keys = HashSet::new();
    keys.insert(seed_pubkey);
    let mut seeds = HashMap::default();
    seeds.insert(
        seed_config.peer_id(),
        Peer::new(vec![seed_addr], keys, seed_role),
    );
    seeds
}