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

hotshot_testing/
helpers.rs
