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    .await
196    .expect("Failed to sign data!");
197
198    let vote_commitment =
199        VersionedVoteData::new(vote.date().clone(), vote.view_number(), upgrade_lock)
200            .await
201            .expect("Failed to create VersionedVoteData!")
202            .commit();
203
204    CERT::create_signed_certificate(
205        vote_commitment,
206        vote.date().clone(),
207        real_qc_sig,
208        vote.view_number(),
209    )
210}
211
212pub fn vid_share<TYPES: NodeType>(
213    shares: &[Proposal<TYPES, VidDisperseShare<TYPES>>],
214    pub_key: TYPES::SignatureKey,
215) -> Proposal<TYPES, VidDisperseShare<TYPES>> {
216    shares
217        .iter()
218        .filter(|s| *s.data.recipient_key() == pub_key)
219        .cloned()
220        .collect::<Vec<_>>()
221        .first()
222        .expect("No VID for key")
223        .clone()
224}
225
226/// create signature
227/// # Panics
228/// if fails to convert node id into keypair
229pub async fn build_assembled_sig<
230    TYPES: NodeType,
231    VOTE: Vote<TYPES>,
232    CERT: Certificate<TYPES, VOTE::Commitment, Voteable = VOTE::Commitment>,
233    DATAType: Committable + Clone + Eq + Hash + Serialize + Debug + 'static,
234>(
235    data: &DATAType,
236    epoch_membership: &EpochMembership<TYPES>,
237    view: ViewNumber,
238    upgrade_lock: &UpgradeLock<TYPES>,
239) -> <TYPES::SignatureKey as SignatureKey>::QcType {
240    let stake_table = CERT::stake_table(epoch_membership).await;
241    let stake_table_entries = StakeTableEntries::<TYPES>::from(stake_table.clone()).0;
242    let real_qc_pp: <TYPES::SignatureKey as SignatureKey>::QcParams<'_> =
243        <TYPES::SignatureKey as SignatureKey>::public_parameter(
244            &stake_table_entries,
245            CERT::threshold(epoch_membership).await,
246        );
247
248    let total_nodes = stake_table.len();
249    let signers = bitvec![1; total_nodes];
250    let mut sig_lists = Vec::new();
251
252    // assemble the vote
253    for node_id in 0..total_nodes {
254        let (private_key_i, public_key_i) = key_pair_for_id::<TYPES>(node_id.try_into().unwrap());
255        let vote: SimpleVote<TYPES, DATAType> = SimpleVote::<TYPES, DATAType>::create_signed_vote(
256            data.clone(),
257            view,
258            &public_key_i,
259            &private_key_i,
260            upgrade_lock,
261        )
262        .await
263        .expect("Failed to sign data!");
264        let original_signature: <TYPES::SignatureKey as SignatureKey>::PureAssembledSignatureType =
265            vote.signature();
266        sig_lists.push(original_signature);
267    }
268
269    <TYPES::SignatureKey as SignatureKey>::assemble(
270        &real_qc_pp,
271        signers.as_bitslice(),
272        &sig_lists[..],
273    )
274}
275
276/// get the keypair for a node id
277#[must_use]
278pub fn key_pair_for_id<TYPES: NodeType>(
279    node_id: u64,
280) -> (
281    <TYPES::SignatureKey as SignatureKey>::PrivateKey,
282    TYPES::SignatureKey,
283) {
284    let private_key = TYPES::SignatureKey::generated_from_seed_indexed([0u8; 32], node_id).1;
285    let public_key = <TYPES as NodeType>::SignatureKey::from_private(&private_key);
286    (private_key, public_key)
287}
288
289pub async fn da_payload_commitment<TYPES: NodeType>(
290    membership: &EpochMembership<TYPES>,
291    transactions: Vec<TestTransaction>,
292    metadata: &<TYPES::BlockPayload as BlockPayload<TYPES>>::Metadata,
293    version: Version,
294) -> VidCommitment {
295    let encoded_transactions = TestTransaction::encode(&transactions);
296
297    vid_commitment(
298        &encoded_transactions,
299        &metadata.encode(),
300        membership.total_nodes().await,
301        version,
302    )
303}
304
305pub async fn build_payload_commitment<TYPES: NodeType>(
306    membership: &EpochMembership<TYPES>,
307    view: ViewNumber,
308    version: Version,
309) -> VidCommitment {
310    // Make some empty encoded transactions, we just care about having a commitment handy for the
311    // later calls. We need the VID commitment to be able to propose later.
312    let encoded_transactions = Vec::new();
313    let num_storage_nodes = membership.committee_members(view).await.len();
314    vid_commitment(&encoded_transactions, &[], num_storage_nodes, version)
315}
316
317pub async fn build_vid_proposal<TYPES: NodeType>(
318    membership: &EpochMembership<TYPES>,
319    view_number: ViewNumber,
320    epoch_number: Option<EpochNumber>,
321    payload: &TYPES::BlockPayload,
322    metadata: &<TYPES::BlockPayload as BlockPayload<TYPES>>::Metadata,
323    private_key: &<TYPES::SignatureKey as SignatureKey>::PrivateKey,
324    upgrade_lock: &UpgradeLock<TYPES>,
325) -> (
326    Proposal<TYPES, VidDisperse<TYPES>>,
327    Vec<Proposal<TYPES, VidDisperseShare<TYPES>>>,
328) {
329    let VidDisperseAndDuration {
330        disperse: vid_disperse,
331        duration: _,
332    } = VidDisperse::calculate_vid_disperse(
333        payload,
334        &membership.coordinator,
335        view_number,
336        epoch_number,
337        epoch_number,
338        metadata,
339        upgrade_lock,
340    )
341    .await
342    .unwrap();
343
344    let signature =
345        TYPES::SignatureKey::sign(private_key, vid_disperse.payload_commitment().as_ref())
346            .expect("Failed to sign VID commitment");
347    let vid_disperse_proposal = Proposal {
348        data: vid_disperse.clone(),
349        signature,
350        _pd: PhantomData,
351    };
352
353    (
354        vid_disperse_proposal,
355        vid_disperse
356            .to_shares()
357            .into_iter()
358            .map(|share| {
359                share
360                    .to_proposal(private_key)
361                    .expect("Failed to sign payload commitment")
362            })
363            .collect(),
364    )
365}
366
367#[allow(clippy::too_many_arguments)]
368pub async fn build_da_certificate<TYPES: NodeType>(
369    membership: &EpochMembership<TYPES>,
370    view_number: ViewNumber,
371    epoch_number: Option<EpochNumber>,
372    transactions: Vec<TestTransaction>,
373    metadata: &<TYPES::BlockPayload as BlockPayload<TYPES>>::Metadata,
374    public_key: &TYPES::SignatureKey,
375    private_key: &<TYPES::SignatureKey as SignatureKey>::PrivateKey,
376    upgrade_lock: &UpgradeLock<TYPES>,
377) -> anyhow::Result<DaCertificate2<TYPES>> {
378    let encoded_transactions = TestTransaction::encode(&transactions);
379
380    let da_payload_commitment = vid_commitment(
381        &encoded_transactions,
382        &metadata.encode(),
383        membership.total_nodes().await,
384        upgrade_lock.version_infallible(view_number),
385    );
386
387    let next_epoch_da_payload_commitment =
388        if upgrade_lock.epochs_enabled(view_number) && membership.epoch().is_some() {
389            Some(vid_commitment(
390                &encoded_transactions,
391                &metadata.encode(),
392                membership
393                    .next_epoch_stake_table()
394                    .await?
395                    .total_nodes()
396                    .await,
397                upgrade_lock.version_infallible(view_number),
398            ))
399        } else {
400            None
401        };
402
403    let da_data = DaData2 {
404        payload_commit: da_payload_commitment,
405        next_epoch_payload_commit: next_epoch_da_payload_commitment,
406        epoch: epoch_number,
407    };
408
409    anyhow::Ok(
410        build_cert::<TYPES, DaData2, DaVote2<TYPES>, DaCertificate2<TYPES>>(
411            da_data,
412            membership,
413            view_number,
414            public_key,
415            private_key,
416            upgrade_lock,
417        )
418        .await,
419    )
420}
421
422/// This function permutes the provided input vector `inputs`, given some order provided within the
423/// `order` vector.
424///
425/// # Examples
426/// let output = permute_input_with_index_order(vec![1, 2, 3], vec![2, 1, 0]);
427/// // Output is [3, 2, 1] now
428pub fn permute_input_with_index_order<T>(inputs: Vec<T>, order: Vec<usize>) -> Vec<T>
429where
430    T: Clone,
431{
432    let mut ordered_inputs = Vec::with_capacity(inputs.len());
433    for &index in &order {
434        ordered_inputs.push(inputs[index].clone());
435    }
436    ordered_inputs
437}
438
439/// This function will create a fake [`View`] from a provided [`Leaf`].
440pub async fn build_fake_view_with_leaf(
441    leaf: Leaf2<TestTypes>,
442    upgrade_lock: &UpgradeLock<TestTypes>,
443    epoch_height: u64,
444) -> View<TestTypes> {
445    build_fake_view_with_leaf_and_state(
446        leaf,
447        TestValidatedState::default(),
448        upgrade_lock,
449        epoch_height,
450    )
451    .await
452}
453
454/// This function will create a fake [`View`] from a provided [`Leaf`] and `state`.
455pub async fn build_fake_view_with_leaf_and_state(
456    leaf: Leaf2<TestTypes>,
457    state: TestValidatedState,
458    _upgrade_lock: &UpgradeLock<TestTypes>,
459    epoch_height: u64,
460) -> View<TestTypes> {
461    let epoch = option_epoch_from_block_number(leaf.with_epoch, leaf.height(), epoch_height);
462    View {
463        view_inner: ViewInner::Leaf {
464            leaf: leaf.commit(),
465            state: state.into(),
466            delta: None,
467            epoch,
468        },
469    }
470}
hotshot_testing/helpers.rs

hotshot_testing/
helpers.rs
