hotshot_testing/
helpers.rs

1// Copyright (c) 2021-2024 Espresso Systems (espressosys.com)
2// This file is part of the HotShot repository.
3
4// You should have received a copy of the MIT License
5// along with the HotShot repository. If not, see <https://mit-license.org/>.
6
7#![allow(clippy::panic)]
8use std::{collections::BTreeMap, fmt::Debug, hash::Hash, marker::PhantomData, sync::Arc};
9
10use async_broadcast::{Receiver, Sender};
11use async_lock::RwLock;
12use bitvec::bitvec;
13use committable::Committable;
14use hotshot::{
15    HotShotInitializer, SystemContext,
16    traits::{BlockPayload, NodeImplementation, TestableNodeImplementation},
17    types::{SignatureKey, SystemContextHandle},
18};
19use hotshot_example_types::{
20    block_types::TestTransaction,
21    node_types::TestTypes,
22    state_types::{TestInstanceState, TestValidatedState},
23    storage_types::TestStorage,
24};
25use hotshot_task_impls::events::HotShotEvent;
26use hotshot_types::{
27    ValidatorConfig,
28    consensus::ConsensusMetricsValue,
29    data::{
30        EpochNumber, Leaf2, VidCommitment, VidDisperse, VidDisperseAndDuration, VidDisperseShare,
31        ViewNumber, vid_commitment,
32    },
33    epoch_membership::{EpochMembership, EpochMembershipCoordinator},
34    message::{Proposal, UpgradeLock},
35    simple_certificate::DaCertificate2,
36    simple_vote::{DaData2, DaVote2, SimpleVote, VersionedVoteData},
37    stake_table::StakeTableEntries,
38    storage_metrics::StorageMetricsValue,
39    traits::{EncodeBytes, election::Membership, node_implementation::NodeType},
40    utils::{View, ViewInner, option_epoch_from_block_number},
41    vote::{Certificate, HasViewNumber, Vote},
42};
43use serde::Serialize;
44use vbs::version::Version;
45
46use crate::{test_builder::TestDescription, test_launcher::TestLauncher};
47
48pub type TestNodeKeyMap = BTreeMap<
49    <TestTypes as NodeType>::SignatureKey,
50    <<TestTypes as NodeType>::SignatureKey as SignatureKey>::PrivateKey,
51>;
52
53/// create the [`SystemContextHandle`] from a node id, with no epochs
54/// # Panics
55/// if cannot create a [`HotShotInitializer`]
56pub async fn build_system_handle<
57    TYPES: NodeType<InstanceState = TestInstanceState>,
58    I: NodeImplementation<TYPES, Storage = TestStorage<TYPES>> + TestableNodeImplementation<TYPES>,
59>(
60    node_id: u64,
61) -> (
62    SystemContextHandle<TYPES, I>,
63    Sender<Arc<HotShotEvent<TYPES>>>,
64    Receiver<Arc<HotShotEvent<TYPES>>>,
65    Arc<TestNodeKeyMap>,
66)
67where
68    <TYPES as NodeType>::Membership: Membership<TYPES, Storage = TestStorage<TYPES>>,
69{
70    let builder: TestDescription<TYPES, I> = TestDescription::default_multiple_rounds();
71
72    let launcher = builder.gen_launcher().map_hotshot_config(|hotshot_config| {
73        hotshot_config.epoch_height = 0;
74    });
75    build_system_handle_from_launcher(node_id, &launcher).await
76}
77
78/// create the [`SystemContextHandle`] from a node id and `TestLauncher`
79/// # Panics
80/// if cannot create a [`HotShotInitializer`]
81pub async fn build_system_handle_from_launcher<
82    TYPES: NodeType<InstanceState = TestInstanceState>,
83    I: NodeImplementation<TYPES, Storage = TestStorage<TYPES>> + TestableNodeImplementation<TYPES>,
84>(
85    node_id: u64,
86    launcher: &TestLauncher<TYPES, I>,
87) -> (
88    SystemContextHandle<TYPES, I>,
89    Sender<Arc<HotShotEvent<TYPES>>>,
90    Receiver<Arc<HotShotEvent<TYPES>>>,
91    Arc<TestNodeKeyMap>,
92)
93where
94    <TYPES as NodeType>::Membership: Membership<TYPES, Storage = TestStorage<TYPES>>,
95{
96    let network = (launcher.resource_generators.channel_generator)(node_id).await;
97    let storage = (launcher.resource_generators.storage)(node_id);
98    let hotshot_config = (launcher.resource_generators.hotshot_config)(node_id);
99
100    let initializer = HotShotInitializer::<TYPES>::from_genesis(
101        TestInstanceState::new(
102            launcher
103                .metadata
104                .async_delay_config
105                .get(&node_id)
106                .cloned()
107                .unwrap_or_default(),
108        ),
109        launcher.metadata.test_config.epoch_height,
110        launcher.metadata.test_config.epoch_start_block,
111        vec![],
112        launcher.metadata.upgrade,
113    )
114    .await
115    .unwrap();
116
117    // See whether or not we should be DA
118    let is_da = node_id < hotshot_config.da_staked_committee_size as u64;
119
120    // We assign node's public key and stake value rather than read from config file since it's a test
121    let validator_config: ValidatorConfig<TYPES> = ValidatorConfig::generated_from_seed_indexed(
122        [0u8; 32],
123        node_id,
124        launcher.metadata.node_stakes.get(node_id),
125        is_da,
126    );
127    let private_key = validator_config.private_key.clone();
128    let public_key = validator_config.public_key.clone();
129    let state_private_key = validator_config.state_private_key.clone();
130
131    let memberships = Arc::new(RwLock::new(TYPES::Membership::new::<I>(
132        hotshot_config.known_nodes_with_stake.clone(),
133        hotshot_config.known_da_nodes.clone(),
134        storage.clone(),
135        network.clone(),
136        public_key.clone(),
137        launcher.metadata.test_config.epoch_height,
138    )));
139
140    let coordinator =
141        EpochMembershipCoordinator::new(memberships, hotshot_config.epoch_height, &storage);
142    let node_key_map = launcher.metadata.build_node_key_map();
143
144    let (c, s, r) = SystemContext::init(
145        public_key,
146        private_key,
147        state_private_key,
148        node_id,
149        hotshot_config,
150        launcher.metadata.upgrade,
151        coordinator,
152        network,
153        initializer,
154        ConsensusMetricsValue::default(),
155        storage,
156        StorageMetricsValue::default(),
157    )
158    .await
159    .expect("Could not init hotshot");
160
161    (c, s, r, node_key_map)
162}
163
164/// create certificate
165/// # Panics
166/// if we fail to sign the data
167pub async fn build_cert<
168    TYPES: NodeType,
169    DATAType: Committable + Clone + Eq + Hash + Serialize + Debug + 'static,
170    VOTE: Vote<TYPES, Commitment = DATAType>,
171    CERT: Certificate<TYPES, VOTE::Commitment, Voteable = VOTE::Commitment>,
172>(
173    data: DATAType,
174    epoch_membership: &EpochMembership<TYPES>,
175    view: ViewNumber,
176    public_key: &TYPES::SignatureKey,
177    private_key: &<TYPES::SignatureKey as SignatureKey>::PrivateKey,
178    upgrade_lock: &UpgradeLock<TYPES>,
179) -> CERT {
180    let real_qc_sig = build_assembled_sig::<TYPES, VOTE, CERT, DATAType>(
181        &data,
182        epoch_membership,
183        view,
184        upgrade_lock,
185    )
186    .await;
187
188    let vote = SimpleVote::<TYPES, DATAType>::create_signed_vote(
189        data,
190        view,
191        public_key,
192        private_key,
193        upgrade_lock,
194    )
195    .expect("Failed to sign data!");
196
197    let vote_commitment =
198        VersionedVoteData::new(vote.date().clone(), vote.view_number(), upgrade_lock)
199            .expect("Failed to create VersionedVoteData!")
200            .commit();
201
202    CERT::create_signed_certificate(
203        vote_commitment,
204        vote.date().clone(),
205        real_qc_sig,
206        vote.view_number(),
207    )
208}
209
210pub fn vid_share<TYPES: NodeType>(
211    shares: &[Proposal<TYPES, VidDisperseShare<TYPES>>],
212    pub_key: TYPES::SignatureKey,
213) -> Proposal<TYPES, VidDisperseShare<TYPES>> {
214    shares
215        .iter()
216        .filter(|s| *s.data.recipient_key() == pub_key)
217        .cloned()
218        .collect::<Vec<_>>()
219        .first()
220        .expect("No VID for key")
221        .clone()
222}
223
224/// create signature
225/// # Panics
226/// if fails to convert node id into keypair
227pub async fn build_assembled_sig<
228    TYPES: NodeType,
229    VOTE: Vote<TYPES>,
230    CERT: Certificate<TYPES, VOTE::Commitment, Voteable = VOTE::Commitment>,
231    DATAType: Committable + Clone + Eq + Hash + Serialize + Debug + 'static,
232>(
233    data: &DATAType,
234    epoch_membership: &EpochMembership<TYPES>,
235    view: ViewNumber,
236    upgrade_lock: &UpgradeLock<TYPES>,
237) -> <TYPES::SignatureKey as SignatureKey>::QcType {
238    let stake_table = CERT::stake_table(epoch_membership).await;
239    let stake_table_entries = StakeTableEntries::<TYPES>::from(stake_table.clone()).0;
240    let real_qc_pp: <TYPES::SignatureKey as SignatureKey>::QcParams<'_> =
241        <TYPES::SignatureKey as SignatureKey>::public_parameter(
242            &stake_table_entries,
243            CERT::threshold(epoch_membership).await,
244        );
245
246    let total_nodes = stake_table.len();
247    let signers = bitvec![1; total_nodes];
248    let mut sig_lists = Vec::new();
249
250    // assemble the vote
251    for node_id in 0..total_nodes {
252        let (private_key_i, public_key_i) = key_pair_for_id::<TYPES>(node_id.try_into().unwrap());
253        let vote: SimpleVote<TYPES, DATAType> = SimpleVote::<TYPES, DATAType>::create_signed_vote(
254            data.clone(),
255            view,
256            &public_key_i,
257            &private_key_i,
258            upgrade_lock,
259        )
260        .expect("Failed to sign data!");
261        let original_signature: <TYPES::SignatureKey as SignatureKey>::PureAssembledSignatureType =
262            vote.signature();
263        sig_lists.push(original_signature);
264    }
265
266    <TYPES::SignatureKey as SignatureKey>::assemble(
267        &real_qc_pp,
268        signers.as_bitslice(),
269        &sig_lists[..],
270    )
271}
272
273/// get the keypair for a node id
274#[must_use]
275pub fn key_pair_for_id<TYPES: NodeType>(
276    node_id: u64,
277) -> (
278    <TYPES::SignatureKey as SignatureKey>::PrivateKey,
279    TYPES::SignatureKey,
280) {
281    let private_key = TYPES::SignatureKey::generated_from_seed_indexed([0u8; 32], node_id).1;
282    let public_key = <TYPES as NodeType>::SignatureKey::from_private(&private_key);
283    (private_key, public_key)
284}
285
286pub async fn da_payload_commitment<TYPES: NodeType>(
287    membership: &EpochMembership<TYPES>,
288    transactions: Vec<TestTransaction>,
289    metadata: &<TYPES::BlockPayload as BlockPayload<TYPES>>::Metadata,
290    version: Version,
291) -> VidCommitment {
292    let encoded_transactions = TestTransaction::encode(&transactions);
293
294    vid_commitment(
295        &encoded_transactions,
296        &metadata.encode(),
297        membership.total_nodes().await,
298        version,
299    )
300}
301
302pub async fn build_payload_commitment<TYPES: NodeType>(
303    membership: &EpochMembership<TYPES>,
304    view: ViewNumber,
305    version: Version,
306) -> VidCommitment {
307    // Make some empty encoded transactions, we just care about having a commitment handy for the
308    // later calls. We need the VID commitment to be able to propose later.
309    let encoded_transactions = Vec::new();
310    let num_storage_nodes = membership.committee_members(view).await.len();
311    vid_commitment(&encoded_transactions, &[], num_storage_nodes, version)
312}
313
314pub async fn build_vid_proposal<TYPES: NodeType>(
315    membership: &EpochMembership<TYPES>,
316    view_number: ViewNumber,
317    epoch_number: Option<EpochNumber>,
318    payload: &TYPES::BlockPayload,
319    metadata: &<TYPES::BlockPayload as BlockPayload<TYPES>>::Metadata,
320    private_key: &<TYPES::SignatureKey as SignatureKey>::PrivateKey,
321    upgrade_lock: &UpgradeLock<TYPES>,
322) -> (
323    Proposal<TYPES, VidDisperse<TYPES>>,
324    Vec<Proposal<TYPES, VidDisperseShare<TYPES>>>,
325) {
326    let VidDisperseAndDuration {
327        disperse: vid_disperse,
328        duration: _,
329    } = VidDisperse::calculate_vid_disperse(
330        payload,
331        &membership.coordinator,
332        view_number,
333        epoch_number,
334        epoch_number,
335        metadata,
336        upgrade_lock,
337    )
338    .await
339    .unwrap();
340
341    let signature =
342        TYPES::SignatureKey::sign(private_key, vid_disperse.payload_commitment().as_ref())
343            .expect("Failed to sign VID commitment");
344    let vid_disperse_proposal = Proposal {
345        data: vid_disperse.clone(),
346        signature,
347        _pd: PhantomData,
348    };
349
350    (
351        vid_disperse_proposal,
352        vid_disperse
353            .to_shares()
354            .into_iter()
355            .map(|share| {
356                share
357                    .to_proposal(private_key)
358                    .expect("Failed to sign payload commitment")
359            })
360            .collect(),
361    )
362}
363
364#[allow(clippy::too_many_arguments)]
365pub async fn build_da_certificate<TYPES: NodeType>(
366    membership: &EpochMembership<TYPES>,
367    view_number: ViewNumber,
368    epoch_number: Option<EpochNumber>,
369    transactions: Vec<TestTransaction>,
370    metadata: &<TYPES::BlockPayload as BlockPayload<TYPES>>::Metadata,
371    public_key: &TYPES::SignatureKey,
372    private_key: &<TYPES::SignatureKey as SignatureKey>::PrivateKey,
373    upgrade_lock: &UpgradeLock<TYPES>,
374) -> anyhow::Result<DaCertificate2<TYPES>> {
375    let encoded_transactions = TestTransaction::encode(&transactions);
376
377    let da_payload_commitment = vid_commitment(
378        &encoded_transactions,
379        &metadata.encode(),
380        membership.total_nodes().await,
381        upgrade_lock.version_infallible(view_number),
382    );
383
384    let next_epoch_da_payload_commitment =
385        if upgrade_lock.epochs_enabled(view_number) && membership.epoch().is_some() {
386            Some(vid_commitment(
387                &encoded_transactions,
388                &metadata.encode(),
389                membership
390                    .next_epoch_stake_table()
391                    .await?
392                    .total_nodes()
393                    .await,
394                upgrade_lock.version_infallible(view_number),
395            ))
396        } else {
397            None
398        };
399
400    let da_data = DaData2 {
401        payload_commit: da_payload_commitment,
402        next_epoch_payload_commit: next_epoch_da_payload_commitment,
403        epoch: epoch_number,
404    };
405
406    anyhow::Ok(
407        build_cert::<TYPES, DaData2, DaVote2<TYPES>, DaCertificate2<TYPES>>(
408            da_data,
409            membership,
410            view_number,
411            public_key,
412            private_key,
413            upgrade_lock,
414        )
415        .await,
416    )
417}
418
419/// This function permutes the provided input vector `inputs`, given some order provided within the
420/// `order` vector.
421///
422/// # Examples
423/// let output = permute_input_with_index_order(vec![1, 2, 3], vec![2, 1, 0]);
424/// // Output is [3, 2, 1] now
425pub fn permute_input_with_index_order<T>(inputs: Vec<T>, order: Vec<usize>) -> Vec<T>
426where
427    T: Clone,
428{
429    let mut ordered_inputs = Vec::with_capacity(inputs.len());
430    for &index in &order {
431        ordered_inputs.push(inputs[index].clone());
432    }
433    ordered_inputs
434}
435
436/// This function will create a fake [`View`] from a provided [`Leaf`].
437pub async fn build_fake_view_with_leaf(
438    leaf: Leaf2<TestTypes>,
439    upgrade_lock: &UpgradeLock<TestTypes>,
440    epoch_height: u64,
441) -> View<TestTypes> {
442    build_fake_view_with_leaf_and_state(
443        leaf,
444        TestValidatedState::default(),
445        upgrade_lock,
446        epoch_height,
447    )
448    .await
449}
450
451/// This function will create a fake [`View`] from a provided [`Leaf`] and `state`.
452pub async fn build_fake_view_with_leaf_and_state(
453    leaf: Leaf2<TestTypes>,
454    state: TestValidatedState,
455    _upgrade_lock: &UpgradeLock<TestTypes>,
456    epoch_height: u64,
457) -> View<TestTypes> {
458    let epoch = option_epoch_from_block_number(leaf.with_epoch, leaf.height(), epoch_height);
459    View {
460        view_inner: ViewInner::Leaf {
461            leaf: leaf.commit(),
462            state: state.into(),
463            delta: None,
464            epoch,
465        },
466    }
467}
hotshot_testing/helpers.rs

hotshot_testing/
helpers.rs
