1use std::{
10 collections::{BTreeMap, HashMap},
11 future::Future,
12 marker::PhantomData,
13};
14
15use alloy::primitives::U256;
16use bitvec::{bitvec, vec::BitVec};
17use committable::{Commitment, Committable};
18use hotshot_utils::anytrace::*;
19use tracing::error;
20
21use crate::{
22 epoch_membership::EpochMembership,
23 light_client::{LightClientState, StakeTableState},
24 message::UpgradeLock,
25 simple_certificate::{LightClientStateUpdateCertificateV2, Threshold},
26 simple_vote::{LightClientStateUpdateVote, VersionedVoteData, Voteable},
27 stake_table::{HSStakeTable, StakeTableEntries},
28 traits::{
29 node_implementation::{NodeType, Versions},
30 signature_key::{
31 LCV2StateSignatureKey, LCV3StateSignatureKey, SignatureKey, StakeTableEntryType,
32 StateSignatureKey,
33 },
34 },
35 PeerConfig,
36};
37
38pub trait Vote<TYPES: NodeType>: HasViewNumber<TYPES> {
40 type Commitment: Voteable<TYPES>;
42
43 fn signature(&self) -> <TYPES::SignatureKey as SignatureKey>::PureAssembledSignatureType;
45 fn date(&self) -> &Self::Commitment;
47 fn data_commitment(&self) -> Commitment<Self::Commitment>;
49
50 fn signing_key(&self) -> TYPES::SignatureKey;
52}
53
54pub trait HasViewNumber<TYPES: NodeType> {
56 fn view_number(&self) -> TYPES::View;
58}
59
60pub trait Certificate<TYPES: NodeType, T>: HasViewNumber<TYPES> {
66 type Voteable: Voteable<TYPES>;
68
69 type Threshold: Threshold<TYPES>;
71
72 fn create_signed_certificate<V: Versions>(
74 vote_commitment: Commitment<VersionedVoteData<TYPES, Self::Voteable, V>>,
75 data: Self::Voteable,
76 sig: <TYPES::SignatureKey as SignatureKey>::QcType,
77 view: TYPES::View,
78 ) -> Self;
79
80 fn is_valid_cert<V: Versions>(
82 &self,
83 stake_table: &[<TYPES::SignatureKey as SignatureKey>::StakeTableEntry],
84 threshold: U256,
85 upgrade_lock: &UpgradeLock<TYPES, V>,
86 ) -> impl std::future::Future<Output = Result<()>>;
87 fn threshold(membership: &EpochMembership<TYPES>) -> impl Future<Output = U256> + Send;
90
91 fn stake_table(
93 membership: &EpochMembership<TYPES>,
94 ) -> impl Future<Output = HSStakeTable<TYPES>> + Send;
95
96 fn total_nodes(membership: &EpochMembership<TYPES>) -> impl Future<Output = usize> + Send;
98
99 fn stake_table_entry(
101 membership: &EpochMembership<TYPES>,
102 pub_key: &TYPES::SignatureKey,
103 ) -> impl Future<Output = Option<PeerConfig<TYPES>>> + Send;
104
105 fn data(&self) -> &Self::Voteable;
107
108 fn data_commitment<V: Versions>(
110 &self,
111 upgrade_lock: &UpgradeLock<TYPES, V>,
112 ) -> impl std::future::Future<Output = Result<Commitment<VersionedVoteData<TYPES, Self::Voteable, V>>>>;
113}
114type SignersMap<COMMITMENT, KEY> = HashMap<
116 COMMITMENT,
117 (
118 BitVec,
119 Vec<<KEY as SignatureKey>::PureAssembledSignatureType>,
120 ),
121>;
122
123#[allow(clippy::type_complexity)]
124pub struct VoteAccumulator<
126 TYPES: NodeType,
127 VOTE: Vote<TYPES>,
128 CERT: Certificate<TYPES, VOTE::Commitment, Voteable = VOTE::Commitment>,
129 V: Versions,
130> {
131 pub vote_outcomes: VoteMap2<
133 Commitment<VersionedVoteData<TYPES, <VOTE as Vote<TYPES>>::Commitment, V>>,
134 TYPES::SignatureKey,
135 <TYPES::SignatureKey as SignatureKey>::PureAssembledSignatureType,
136 >,
137 pub signers: SignersMap<
140 Commitment<VersionedVoteData<TYPES, <VOTE as Vote<TYPES>>::Commitment, V>>,
141 TYPES::SignatureKey,
142 >,
143 pub phantom: PhantomData<(TYPES, VOTE, CERT)>,
145 pub upgrade_lock: UpgradeLock<TYPES, V>,
147}
148
149impl<
150 TYPES: NodeType,
151 VOTE: Vote<TYPES>,
152 CERT: Certificate<TYPES, VOTE::Commitment, Voteable = VOTE::Commitment>,
153 V: Versions,
154 > VoteAccumulator<TYPES, VOTE, CERT, V>
155{
156 pub async fn accumulate(
160 &mut self,
161 vote: &VOTE,
162 membership: EpochMembership<TYPES>,
163 ) -> Option<CERT> {
164 let key = vote.signing_key();
165
166 let vote_commitment = match VersionedVoteData::new(
167 vote.date().clone(),
168 vote.view_number(),
169 &self.upgrade_lock,
170 )
171 .await
172 {
173 Ok(data) => data.commit(),
174 Err(e) => {
175 tracing::warn!("Failed to generate versioned vote data: {e}");
176 return None;
177 },
178 };
179
180 if !key.validate(&vote.signature(), vote_commitment.as_ref()) {
181 error!("Invalid vote! Vote Data {:?}", vote.date());
182 return None;
183 }
184
185 let stake_table_entry = CERT::stake_table_entry(&membership, &key).await?;
186 let stake_table = CERT::stake_table(&membership).await;
187 let total_nodes = CERT::total_nodes(&membership).await;
188 let threshold = CERT::threshold(&membership).await;
189
190 let vote_node_id = stake_table
191 .iter()
192 .position(|x| *x == stake_table_entry.clone())?;
193
194 let original_signature: <TYPES::SignatureKey as SignatureKey>::PureAssembledSignatureType =
195 vote.signature();
196
197 let (total_stake_casted, total_vote_map) = self
198 .vote_outcomes
199 .entry(vote_commitment)
200 .or_insert_with(|| (U256::from(0), BTreeMap::new()));
201
202 if total_vote_map.contains_key(&key) {
204 return None;
205 }
206 let (signers, sig_list) = self
207 .signers
208 .entry(vote_commitment)
209 .or_insert((bitvec![0; total_nodes], Vec::new()));
210 if signers.get(vote_node_id).as_deref() == Some(&true) {
211 error!("Node id is already in signers list");
212 return None;
213 }
214 signers.set(vote_node_id, true);
215 sig_list.push(original_signature);
216
217 *total_stake_casted += stake_table_entry.stake_table_entry.stake();
218 total_vote_map.insert(key, (vote.signature(), vote_commitment));
219
220 if *total_stake_casted >= threshold {
221 let stake_table_entries = StakeTableEntries::<TYPES>::from(stake_table).0;
223 let real_qc_pp: <<TYPES as NodeType>::SignatureKey as SignatureKey>::QcParams<'_> =
224 <TYPES::SignatureKey as SignatureKey>::public_parameter(
225 &stake_table_entries,
226 threshold,
227 );
228
229 let real_qc_sig = <TYPES::SignatureKey as SignatureKey>::assemble(
230 &real_qc_pp,
231 signers.as_bitslice(),
232 &sig_list[..],
233 );
234
235 let cert = CERT::create_signed_certificate::<V>(
236 vote_commitment,
237 vote.date().clone(),
238 real_qc_sig,
239 vote.view_number(),
240 );
241 return Some(cert);
242 }
243 None
244 }
245}
246
247type VoteMap2<COMMITMENT, PK, SIG> = HashMap<COMMITMENT, (U256, BTreeMap<PK, (SIG, COMMITMENT)>)>;
249
250#[allow(clippy::type_complexity)]
252pub struct LightClientStateUpdateVoteAccumulator<TYPES: NodeType> {
253 pub vote_outcomes: HashMap<
254 (LightClientState, StakeTableState),
255 (
256 U256,
257 HashMap<
258 TYPES::StateSignatureKey,
259 (
260 <TYPES::StateSignatureKey as StateSignatureKey>::StateSignature, <TYPES::StateSignatureKey as StateSignatureKey>::StateSignature, ),
263 >,
264 ),
265 >,
266}
267
268impl<TYPES: NodeType> LightClientStateUpdateVoteAccumulator<TYPES> {
269 pub async fn accumulate(
273 &mut self,
274 key: &TYPES::SignatureKey,
275 vote: &LightClientStateUpdateVote<TYPES>,
276 membership: &EpochMembership<TYPES>,
277 ) -> Option<LightClientStateUpdateCertificateV2<TYPES>> {
278 let epoch = membership.epoch()?;
279 let threshold = membership.success_threshold().await;
280 let PeerConfig {
281 stake_table_entry,
282 state_ver_key,
283 } = membership.stake(key).await?;
284
285 if !<TYPES::StateSignatureKey as LCV2StateSignatureKey>::verify_state_sig(
286 &state_ver_key,
287 &vote.v2_signature,
288 &vote.light_client_state,
289 &vote.next_stake_table_state,
290 ) {
291 error!("Invalid light client state update vote {vote:?}");
292 return None;
293 }
294 if !<TYPES::StateSignatureKey as LCV3StateSignatureKey>::verify_state_sig(
295 &state_ver_key,
296 &vote.signature,
297 vote.signed_state_digest,
298 ) {
299 error!("Invalid light client state update vote {vote:?}");
300 return None;
301 }
302 let (total_stake_casted, vote_map) = self
303 .vote_outcomes
304 .entry((vote.light_client_state, vote.next_stake_table_state))
305 .or_insert_with(|| (U256::from(0), HashMap::new()));
306
307 if vote_map.contains_key(&state_ver_key) {
309 tracing::warn!("Duplicate vote (key: {key:?}, vote: {vote:?})");
310 return None;
311 }
312
313 *total_stake_casted += stake_table_entry.stake();
314 vote_map.insert(
315 state_ver_key.clone(),
316 (vote.signature.clone(), vote.v2_signature.clone()),
317 );
318
319 if *total_stake_casted >= threshold {
320 return Some(LightClientStateUpdateCertificateV2 {
321 epoch,
322 light_client_state: vote.light_client_state,
323 next_stake_table_state: vote.next_stake_table_state,
324 signatures: Vec::from_iter(
325 vote_map
326 .iter()
327 .map(|(k, (v3, v2))| (k.clone(), v3.clone(), v2.clone())),
328 ),
329 auth_root: vote.auth_root,
330 });
331 }
332 None
333 }
334}