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