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
83
84
85
86

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

// Copyright 2021 Conflux Foundation. All rights reserved.
// Conflux is free software and distributed under GNU General Public License.
// See http://www.gnu.org/licenses/

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

use crate::config::{
    NetworkConfig, NodeConfig, Peer, PeerRole, PeerSet, TestConfig,
    HANDSHAKE_VERSION,
};
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 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
}