1use std::{
10 collections::{BTreeMap, HashMap},
11 marker::PhantomData,
12};
13
14use alloy::primitives::{FixedBytes, U256};
15use bitvec::{bitvec, vec::BitVec};
16use committable::{Commitment, Committable};
17use hotshot_utils::anytrace::*;
18use tracing::error;
19use vbs::version::Version;
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 ) -> Result<()>;
88
89 fn signers(
91 &self,
92 stake_table: &[<TYPES::SignatureKey as SignatureKey>::StakeTableEntry],
93 threshold: U256,
94 ) -> Result<Vec<<TYPES::SignatureKey as SignatureKey>::VerificationKeyType>>;
95
96 fn threshold(membership: &EpochMembership<TYPES>) -> U256;
99
100 fn stake_table(membership: &EpochMembership<TYPES>) -> HSStakeTable<TYPES>;
102
103 fn total_nodes(membership: &EpochMembership<TYPES>) -> usize;
105
106 fn stake_table_entry(
108 membership: &EpochMembership<TYPES>,
109 pub_key: &TYPES::SignatureKey,
110 ) -> Option<PeerConfig<TYPES>>;
111
112 fn data(&self) -> &Self::Voteable;
114
115 fn data_commitment(
117 &self,
118 upgrade_lock: &UpgradeLock<TYPES>,
119 ) -> 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 fn new(upgrade_lock: UpgradeLock<TYPES>) -> Self {
162 Self {
163 vote_outcomes: HashMap::new(),
164 signers: HashMap::new(),
165 phantom: PhantomData,
166 upgrade_lock,
167 }
168 }
169 pub fn accumulate(&mut self, vote: &VOTE, membership: EpochMembership<TYPES>) -> Option<CERT> {
173 let key = vote.signing_key();
174
175 let vote_commitment = match VersionedVoteData::new(
176 vote.date().clone(),
177 vote.view_number(),
178 &self.upgrade_lock,
179 ) {
180 Ok(data) => data.commit(),
181 Err(e) => {
182 tracing::warn!("Failed to generate versioned vote data: {e}");
183 return None;
184 },
185 };
186
187 if self.upgrade_lock.version(vote.view_number()).ok()? < (Version { major: 0, minor: 6 })
188 && !key.validate(&vote.signature(), vote_commitment.as_ref())
189 {
190 error!("Invalid vote! Vote Data {:?}", vote.date());
191 return None;
192 }
193
194 let stake_table_entry = CERT::stake_table_entry(&membership, &key)?;
195 let stake_table = CERT::stake_table(&membership);
196 let total_nodes = CERT::total_nodes(&membership);
197 let threshold = CERT::threshold(&membership);
198
199 let vote_node_id = stake_table
200 .iter()
201 .position(|x| *x == stake_table_entry.clone())?;
202
203 let original_signature: <TYPES::SignatureKey as SignatureKey>::PureAssembledSignatureType =
204 vote.signature();
205
206 let (total_stake_casted, total_vote_map) = self
207 .vote_outcomes
208 .entry(vote_commitment)
209 .or_insert_with(|| (U256::from(0), BTreeMap::new()));
210
211 if total_vote_map.contains_key(&key) {
213 return None;
214 }
215 let (signers, sig_list) = self
216 .signers
217 .entry(vote_commitment)
218 .or_insert((bitvec![0; total_nodes], Vec::new()));
219 if signers.get(vote_node_id).as_deref() == Some(&true) {
220 error!("Node id is already in signers list");
221 return None;
222 }
223 signers.set(vote_node_id, true);
224 sig_list.push(original_signature);
225
226 *total_stake_casted += stake_table_entry.stake_table_entry.stake();
227 total_vote_map.insert(key, (vote.signature(), vote_commitment));
228
229 if *total_stake_casted >= threshold {
230 let stake_table_entries = StakeTableEntries::<TYPES>::from(stake_table).0;
232 let real_qc_pp: <<TYPES as NodeType>::SignatureKey as SignatureKey>::QcParams<'_> =
233 <TYPES::SignatureKey as SignatureKey>::public_parameter(
234 &stake_table_entries,
235 threshold,
236 );
237
238 let real_qc_sig = <TYPES::SignatureKey as SignatureKey>::assemble(
239 &real_qc_pp,
240 signers.as_bitslice(),
241 &sig_list[..],
242 );
243
244 let cert = CERT::create_signed_certificate(
245 vote_commitment,
246 vote.date().clone(),
247 real_qc_sig,
248 vote.view_number(),
249 );
250 return Some(cert);
251 }
252 None
253 }
254}
255
256type VoteMap2<COMMITMENT, PK, SIG> = HashMap<COMMITMENT, (U256, BTreeMap<PK, (SIG, COMMITMENT)>)>;
258
259#[allow(clippy::type_complexity)]
261pub struct LightClientStateUpdateVoteAccumulator<TYPES: NodeType> {
262 pub vote_outcomes: HashMap<
263 (LightClientState, StakeTableState, FixedBytes<32>),
264 (
265 U256,
266 HashMap<
267 TYPES::StateSignatureKey,
268 (
269 <TYPES::StateSignatureKey as StateSignatureKey>::StateSignature, <TYPES::StateSignatureKey as StateSignatureKey>::StateSignature, ),
272 >,
273 ),
274 >,
275
276 pub upgrade_lock: UpgradeLock<TYPES>,
277}
278
279impl<TYPES: NodeType> LightClientStateUpdateVoteAccumulator<TYPES> {
280 pub fn accumulate(
284 &mut self,
285 key: &TYPES::SignatureKey,
286 vote: &LightClientStateUpdateVote2<TYPES>,
287 membership: &EpochMembership<TYPES>,
288 ) -> Option<LightClientStateUpdateCertificateV2<TYPES>> {
289 let epoch = membership.epoch()?;
290 let threshold = membership.success_threshold();
291 let PeerConfig {
292 stake_table_entry,
293 state_ver_key,
294 ..
295 } = membership.stake(key)?;
296
297 if !<TYPES::StateSignatureKey as LCV2StateSignatureKey>::verify_state_sig(
298 &state_ver_key,
299 &vote.v2_signature,
300 &vote.light_client_state,
301 &vote.next_stake_table_state,
302 ) {
303 error!("Invalid light client state update vote {vote:?}");
304 return None;
305 }
306 if self
308 .upgrade_lock
309 .proposal2_version(ViewNumber::new(vote.light_client_state.view_number))
310 && !<TYPES::StateSignatureKey as LCV3StateSignatureKey>::verify_state_sig(
311 &state_ver_key,
312 &vote.signature,
313 vote.signed_state_digest,
314 )
315 {
316 error!("Invalid light client state update vote {vote:?}");
317 return None;
318 }
319 let (total_stake_casted, vote_map) = self
320 .vote_outcomes
321 .entry((
322 vote.light_client_state,
323 vote.next_stake_table_state,
324 vote.auth_root,
325 ))
326 .or_insert_with(|| (U256::from(0), HashMap::new()));
327
328 if vote_map.contains_key(&state_ver_key) {
330 tracing::warn!("Duplicate vote (key: {key:?}, vote: {vote:?})");
331 return None;
332 }
333
334 *total_stake_casted += stake_table_entry.stake();
335 vote_map.insert(
336 state_ver_key.clone(),
337 (vote.signature.clone(), vote.v2_signature.clone()),
338 );
339
340 if *total_stake_casted >= threshold {
341 return Some(LightClientStateUpdateCertificateV2 {
342 epoch,
343 light_client_state: vote.light_client_state,
344 next_stake_table_state: vote.next_stake_table_state,
345 signatures: Vec::from_iter(
346 vote_map
347 .iter()
348 .map(|(k, (v3, v2))| (k.clone(), v3.clone(), v2.clone())),
349 ),
350 auth_root: vote.auth_root,
351 });
352 }
353 None
354 }
355}