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
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181

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

use crate::testutils::{
    test_framework::TestFramework,
    test_node::{NodeId, NodeType, TestNode},
};
use diem_config::{
    config::NodeConfig,
    network_id::{NetworkId, PeerNetworkId},
};
use rand::{rngs::StdRng, SeedableRng};
use std::{collections::HashMap, marker::PhantomData};

/// A builder for a [`TestFramework`] implementation.
///
/// This handles making sure that nodes are unique, and that they have the channels attached
/// to each other to send messages between each other.
///
pub struct TestFrameworkBuilder<Framework: TestFramework<Node>, Node: TestNode> {
    /// Owners are a simple stand in for actual [`AccountAddress`] uniqueness in the network.
    /// An `owner` can have 1 of each [`NodeType`] of `Node`. This simplifies linking
    /// [`NodeType::Validator`] and [`NodeType::ValidatorFullNode`] later, as well as keeps
    /// a unique identifier for each node as [`NodeId`].
    owners: u32,
    /// The unique mapping of `Node` to ensure no duplicates.
    nodes: HashMap<NodeId, Node>,
    /// Random generator for randomness in keys and [`PeerId`]s. Hardcoded to remove non-determinism.
    rng: StdRng,
    _framework_marker: PhantomData<Framework>,
}

impl<Framework: TestFramework<Node>, Node: TestNode> TestFrameworkBuilder<Framework, Node> {
    /// Create a new [`TestFrameworkBuilder`], ensuring that there is a fixed number of `owners`.
    pub fn new(owners: u32) -> Self {
        Self {
            owners,
            nodes: HashMap::new(),
            rng: StdRng::from_seed([0u8; 32]),
            _framework_marker: PhantomData::default(),
        }
    }

    /// Builds the [`TestFramework`]
    pub fn build(self) -> Framework {
        TestFramework::new(self.nodes)
    }

    /// Adds a [`TestNode`] of [`NodeType::Validator`]
    pub fn add_validator(mut self, owner: u32) -> Self {
        let config = NodeConfig::random_with_template(
            owner,
            &NodeConfig::default_for_validator(),
            &mut self.rng,
        );
        let peer_id = config
            .validator_network
            .as_ref()
            .expect("Validator must have a validator network")
            .peer_id();

        self.add_node(
            owner,
            NodeType::Validator,
            config,
            &[
                PeerNetworkId::new(NetworkId::Validator, peer_id),
                PeerNetworkId::new(NetworkId::Vfn, peer_id),
            ],
        )
    }

    /// Adds a [`TestNode`] of [`NodeType::ValidatorFullNode`]
    pub fn add_vfn(mut self, owner: u32) -> Self {
        let config = NodeConfig::random_with_template(
            owner,
            &NodeConfig::default_for_validator_full_node(),
            &mut self.rng,
        );
        let peer_id = config
            .full_node_networks
            .iter()
            .find(|network| network.network_id == NetworkId::Public)
            .expect("Vfn must have a public network")
            .peer_id();

        self.add_node(
            owner,
            NodeType::ValidatorFullNode,
            config,
            &[
                PeerNetworkId::new(NetworkId::Vfn, peer_id),
                PeerNetworkId::new(NetworkId::Public, peer_id),
            ],
        )
    }

    /// Adds a [`TestNode`] of [`NodeType::PublicFullNode`]
    pub fn add_pfn(mut self, owner: u32) -> Self {
        let config = NodeConfig::random_with_template(
            owner,
            &NodeConfig::default_for_public_full_node(),
            &mut self.rng,
        );
        let peer_id = config
            .full_node_networks
            .iter()
            .find(|network| network.network_id == NetworkId::Public)
            .expect("Pfn must have a public network")
            .peer_id();

        self.add_node(
            owner,
            NodeType::PublicFullNode,
            config,
            &[PeerNetworkId::new(NetworkId::Public, peer_id)],
        )
    }

    /// Add a node to the network, ensuring that it doesn't already exist
    fn add_node(
        mut self,
        owner: u32,
        node_type: NodeType,
        config: NodeConfig,
        peer_network_ids: &[PeerNetworkId],
    ) -> Self {
        assert!(owner < self.owners);

        let node_id = NodeId { owner, node_type };
        assert!(!self.nodes.contains_key(&node_id));

        let mut node = Framework::build_node(node_id, config, peer_network_ids);

        // Add node's sender to every possible node that it could be connected to.
        for (other_node_id, other_node) in self.nodes.iter_mut() {
            if let Some(network_id) = other_node.find_common_network(&node) {
                // The VFN network only goes between the same owner, skip it if it doesn't match
                if network_id == NetworkId::Vfn && owner != other_node_id.owner {
                    continue;
                }

                // Add the inbound handle for the new node to the other node
                add_inbound_peer_handle_to_node(network_id, other_node, &node);

                // Add the inbound handle for the other node to the new node
                add_inbound_peer_handle_to_node(network_id, &mut node, other_node);
            }
        }

        self.nodes.insert(node_id, node);
        self
    }

    pub fn single_validator() -> Node {
        let mut test_framework: Framework = TestFrameworkBuilder::new(1).add_validator(0).build();
        test_framework.take_node(NodeId::validator(0))
    }

    pub fn single_vfn() -> Node {
        let mut test_framework: Framework = TestFrameworkBuilder::new(1).add_vfn(0).build();
        test_framework.take_node(NodeId::vfn(0))
    }

    pub fn single_pfn() -> Node {
        let mut test_framework: Framework = TestFrameworkBuilder::new(1).add_pfn(0).build();
        test_framework.take_node(NodeId::pfn(0))
    }
}

/// Adds `receiving_node`'s peer_handle to the `sending_node`
fn add_inbound_peer_handle_to_node<Node: TestNode>(
    network_id: NetworkId,
    sending_node: &mut Node,
    receiving_node: &Node,
) {
    sending_node.add_inbound_handle_for_peer(
        receiving_node.peer_network_id(network_id),
        receiving_node.get_inbound_handle(network_id),
    );
}