Variable stake and randomized leader selection (#2638)

ss-es · web-flow · commit 36d87023def5 · 2025-02-24T13:15:59.000-05:00
diff --git a/hotshot-example-types/src/node_types.rs b/hotshot-example-types/src/node_types.rs
@@ -11,7 +11,7 @@ pub use hotshot::traits::election::helpers::{
 };
 use hotshot::traits::{
     election::{
-        helpers::QuorumFilterConfig, randomized_committee::RandomizedCommittee,
+        helpers::QuorumFilterConfig, randomized_committee::Committee,
         randomized_committee_members::RandomizedCommitteeMembers,
         static_committee::StaticCommittee,
         static_committee_leader_two_views::StaticCommitteeLeaderForTwoViews,
@@ -97,7 +97,7 @@ impl NodeType for TestTypesRandomizedLeader {
     type Transaction = TestTransaction;
     type ValidatedState = TestValidatedState;
     type InstanceState = TestInstanceState;
-    type Membership = RandomizedCommittee<TestTypesRandomizedLeader>;
+    type Membership = Committee<TestTypesRandomizedLeader>;
     type BuilderSignatureKey = BuilderKey;
 }
 
diff --git a/hotshot-types/src/drb.rs b/hotshot-types/src/drb.rs
@@ -126,3 +126,122 @@ impl<TYPES: NodeType> Default for DrbSeedsAndResults<TYPES> {
         Self::new()
     }
 }
+
+/// Functions for leader selection based on the DRB.
+///
+/// The algorithm we use is:
+///
+/// Initialization:
+/// - obtain `drb: [u8; 32]` from the DRB calculation
+/// - sort the stake table for a given epoch by `xor(drb, public_key)`
+/// - generate a cdf of the cumulative stake using this newly-sorted table,
+///   along with a hash of the stake table entries
+///
+/// Selecting a leader:
+/// - calculate the SHA512 hash of the `drb_result`, `view_number` and `stake_table_hash`
+/// - find the first index in the cdf for which the remainder of this hash modulo the `total_stake`
+///   is strictly smaller than the cdf entry
+/// - return the corresponding node as the leader for that view
+pub mod election {
+    use primitive_types::{U256, U512};
+    use sha2::{Digest, Sha256, Sha512};
+
+    use crate::traits::signature_key::{SignatureKey, StakeTableEntryType};
+
+    /// Calculate `xor(drb.cycle(), public_key)`, returning the result as a vector of bytes
+    fn cyclic_xor(drb: [u8; 32], public_key: Vec<u8>) -> Vec<u8> {
+        let drb: Vec<u8> = drb.to_vec();
+
+        let mut result: Vec<u8> = vec![];
+
+        for (drb_byte, public_key_byte) in public_key.iter().zip(drb.iter().cycle()) {
+            result.push(drb_byte ^ public_key_byte);
+        }
+
+        result
+    }
+
+    /// Generate the stake table CDF, as well as a hash of the resulting stake table
+    pub fn generate_stake_cdf<Key: SignatureKey, Entry: StakeTableEntryType<Key>>(
+        mut stake_table: Vec<Entry>,
+        drb: [u8; 32],
+    ) -> RandomizedCommittee<Entry> {
+        // sort by xor(public_key, drb_result)
+        stake_table.sort_by(|a, b| {
+            cyclic_xor(drb, a.public_key().to_bytes())
+                .cmp(&cyclic_xor(drb, b.public_key().to_bytes()))
+        });
+
+        let mut hasher = Sha256::new();
+
+        let mut cumulative_stake = U256::from(0);
+        let mut cdf = vec![];
+
+        for entry in stake_table {
+            cumulative_stake += entry.stake();
+            hasher.update(entry.public_key().to_bytes());
+
+            cdf.push((entry, cumulative_stake));
+        }
+
+        RandomizedCommittee {
+            cdf,
+            stake_table_hash: hasher.finalize().into(),
+            drb,
+        }
+    }
+
+    /// select the leader for a view
+    ///
+    /// # Panics
+    /// Panics if `cdf` is empty. Results in undefined behaviour if `cdf` is not ordered.
+    ///
+    /// Note that we try to downcast a U512 to a U256,
+    /// but this should never panic because the U512 should be strictly smaller than U256::MAX by construction.
+    pub fn select_randomized_leader<
+        SignatureKey,
+        Entry: StakeTableEntryType<SignatureKey> + Clone,
+    >(
+        randomized_committee: &RandomizedCommittee<Entry>,
+        view: u64,
+    ) -> Entry {
+        let RandomizedCommittee {
+            cdf,
+            stake_table_hash,
+            drb,
+        } = randomized_committee;
+        // We hash the concatenated drb, view and stake table hash.
+        let mut hasher = Sha512::new();
+        hasher.update(drb);
+        hasher.update(view.to_le_bytes());
+        hasher.update(stake_table_hash);
+        let raw_breakpoint: [u8; 64] = hasher.finalize().into();
+
+        // then calculate the remainder modulo the total stake as a U512
+        let remainder: U512 =
+            U512::from_little_endian(&raw_breakpoint) % U512::from(cdf.last().unwrap().1);
+
+        // and drop the top 32 bytes, downcasting to a U256
+        let breakpoint: U256 = U256::try_from(remainder).unwrap();
+
+        // now find the first index where the breakpoint is strictly smaller than the cdf
+        //
+        // in principle, this may result in an index larger than `cdf.len()`.
+        // however, we have ensured by construction that `breakpoint < total_stake`
+        // and so the largest index we can actually return is `cdf.len() - 1`
+        let index = cdf.partition_point(|(_, cumulative_stake)| breakpoint >= *cumulative_stake);
+
+        // and return the corresponding entry
+        cdf[index].0.clone()
+    }
+
+    #[derive(Clone, Debug)]
+    pub struct RandomizedCommittee<Entry> {
+        /// cdf of nodes by cumulative stake
+        cdf: Vec<(Entry, U256)>,
+        /// Hash of the stake table
+        stake_table_hash: [u8; 32],
+        /// DRB result
+        drb: [u8; 32],
+    }
+}
diff --git a/hotshot/src/traits/election/randomized_committee.rs b/hotshot/src/traits/election/randomized_committee.rs
@@ -4,24 +4,27 @@
 // You should have received a copy of the MIT License
 // along with the HotShot repository. If not, see <https://mit-license.org/>.
 
+use std::{cmp::max, collections::BTreeMap, num::NonZeroU64};
+
 use hotshot_types::{
-    drb::DrbResult,
+    drb::{
+        election::{generate_stake_cdf, select_randomized_leader, RandomizedCommittee},
+        DrbResult,
+    },
     traits::{
         election::Membership,
         node_implementation::NodeType,
         signature_key::{SignatureKey, StakeTableEntryType},
     },
     PeerConfig,
 };
-use hotshot_utils::anytrace::Result;
+use hotshot_utils::anytrace::*;
 use primitive_types::U256;
-use rand::{rngs::StdRng, Rng};
-use std::{cmp::max, collections::BTreeMap, num::NonZeroU64};
 
-#[derive(Clone, Debug, Eq, PartialEq, Hash)]
+#[derive(Clone, Debug)]
 
 /// The static committee election
-pub struct RandomizedCommittee<T: NodeType> {
+pub struct Committee<T: NodeType> {
     /// The nodes eligible for leadership.
     /// NOTE: This is currently a hack because the DA leader needs to be the quorum
     /// leader but without voting rights.
@@ -33,6 +36,9 @@ pub struct RandomizedCommittee<T: NodeType> {
     /// The nodes on the committee and their stake
     da_stake_table: Vec<<T::SignatureKey as SignatureKey>::StakeTableEntry>,
 
+    /// Stake tables randomized with the DRB, used (only) for leader election
+    randomized_committee: RandomizedCommittee<<T::SignatureKey as SignatureKey>::StakeTableEntry>,
+
     /// The nodes on the committee and their stake, indexed by public key
     indexed_stake_table:
         BTreeMap<T::SignatureKey, <T::SignatureKey as SignatureKey>::StakeTableEntry>,
@@ -42,7 +48,7 @@ pub struct RandomizedCommittee<T: NodeType> {
         BTreeMap<T::SignatureKey, <T::SignatureKey as SignatureKey>::StakeTableEntry>,
 }
 
-impl<TYPES: NodeType> Membership<TYPES> for RandomizedCommittee<TYPES> {
+impl<TYPES: NodeType> Membership<TYPES> for Committee<TYPES> {
     type Error = hotshot_utils::anytrace::Error;
 
     /// Create a new election
@@ -91,10 +97,14 @@ impl<TYPES: NodeType> Membership<TYPES> for RandomizedCommittee<TYPES> {
             .map(|entry| (TYPES::SignatureKey::public_key(entry), entry.clone()))
             .collect();
 
+        // We use a constant value of `[0u8; 32]` for the drb, since this is just meant to be used in tests
+        let randomized_committee = generate_stake_cdf(eligible_leaders.clone(), [0u8; 32]);
+
         Self {
             eligible_leaders,
             stake_table: members,
             da_stake_table: da_members,
+            randomized_committee,
             indexed_stake_table,
             indexed_da_stake_table,
         }
@@ -205,13 +215,7 @@ impl<TYPES: NodeType> Membership<TYPES> for RandomizedCommittee<TYPES> {
         view_number: <TYPES as NodeType>::View,
         _epoch: Option<<TYPES as NodeType>::Epoch>,
     ) -> Result<TYPES::SignatureKey> {
-        let mut rng: StdRng = rand::SeedableRng::seed_from_u64(*view_number);
-
-        let randomized_view_number: u64 = rng.gen_range(0..=u64::MAX);
-        #[allow(clippy::cast_possible_truncation)]
-        let index = randomized_view_number as usize % self.eligible_leaders.len();
-
-        let res = self.eligible_leaders[index].clone();
+        let res = select_randomized_leader(&self.randomized_committee, *view_number);
 
         Ok(TYPES::SignatureKey::public_key(&res))
     }
@@ -248,5 +252,5 @@ impl<TYPES: NodeType> Membership<TYPES> for RandomizedCommittee<TYPES> {
         .unwrap()
     }
 
-    fn add_drb_result(&mut self, _epoch: <TYPES as NodeType>::Epoch, _drb_result: DrbResult) {}
+    fn add_drb_result(&mut self, _epoch: <TYPES as NodeType>::Epoch, _drb: DrbResult) {}
 }
diff --git a/types/src/v0/impls/stake_table.rs b/types/src/v0/impls/stake_table.rs
@@ -1,6 +1,8 @@
-use super::{
-    v0_3::{DAMembers, StakeTable, StakeTables},
-    Header, L1Client, NodeState, PubKey, SeqTypes,
+use std::{
+    cmp::max,
+    collections::{BTreeMap, BTreeSet, HashMap},
+    num::NonZeroU64,
+    str::FromStr,
 };
 
 use async_trait::async_trait;
@@ -11,7 +13,10 @@ use hotshot::types::{BLSPubKey, SignatureKey as _};
 use hotshot_contract_adapter::stake_table::{bls_alloy_to_jf, NodeInfoJf};
 use hotshot_types::{
     data::EpochNumber,
-    drb::DrbResult,
+    drb::{
+        election::{generate_stake_cdf, select_randomized_leader, RandomizedCommittee},
+        DrbResult,
+    },
     stake_table::StakeTableEntry,
     traits::{
         election::Membership,
@@ -21,16 +26,14 @@ use hotshot_types::{
     PeerConfig,
 };
 use itertools::Itertools;
-use std::{
-    cmp::max,
-    collections::{BTreeMap, BTreeSet, HashMap},
-    num::NonZeroU64,
-    str::FromStr,
-};
 use thiserror::Error;
-
 use url::Url;
 
+use super::{
+    v0_3::{DAMembers, StakeTable, StakeTables},
+    Header, L1Client, NodeState, PubKey, SeqTypes,
+};
+
 type Epoch = <SeqTypes as NodeType>::Epoch;
 
 impl StakeTables {
@@ -107,8 +110,8 @@ pub struct EpochCommittees {
     /// Address of Stake Table Contract
     contract_address: Option<Address>,
 
-    /// The results of DRB calculations
-    drb_result_table: BTreeMap<Epoch, DrbResult>,
+    /// Randomized committees, filled when we receive the DrbResult
+    randomized_committees: BTreeMap<Epoch, RandomizedCommittee<StakeTableEntry<PubKey>>>,
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -256,7 +259,7 @@ impl EpochCommittees {
             _epoch_size: epoch_size,
             l1_client: instance_state.l1_client.clone(),
             contract_address: instance_state.chain_config.stake_table_contract,
-            drb_result_table: BTreeMap::new(),
+            randomized_committees: BTreeMap::new(),
         }
     }
 
@@ -337,7 +340,7 @@ impl Membership<SeqTypes> for EpochCommittees {
             l1_client: L1Client::new(vec![Url::from_str("http:://ab.b").unwrap()])
                 .expect("Failed to create L1 client"),
             contract_address: None,
-            drb_result_table: BTreeMap::new(),
+            randomized_committees: BTreeMap::new(),
         }
     }
 
@@ -432,15 +435,26 @@ impl Membership<SeqTypes> for EpochCommittees {
         view_number: <SeqTypes as NodeType>::View,
         epoch: Option<Epoch>,
     ) -> Result<PubKey, Self::Error> {
-        let leaders = self
-            .state(&epoch)
-            .ok_or(LeaderLookupError)?
-            .eligible_leaders
-            .clone();
+        if let Some(epoch) = epoch {
+            let Some(randomized_committee) = self.randomized_committees.get(&epoch) else {
+                tracing::error!(
+                    "We are missing the randomized committee for epoch {}",
+                    epoch
+                );
+                return Err(LeaderLookupError);
+            };
+
+            Ok(PubKey::public_key(&select_randomized_leader(
+                randomized_committee,
+                *view_number,
+            )))
+        } else {
+            let leaders = &self.non_epoch_committee.eligible_leaders;
 
-        let index = *view_number as usize % leaders.len();
-        let res = leaders[index].clone();
-        Ok(PubKey::public_key(&res))
+            let index = *view_number as usize % leaders.len();
+            let res = leaders[index].clone();
+            Ok(PubKey::public_key(&res))
+        }
     }
 
     /// Get the total number of nodes in the committee
@@ -504,8 +518,17 @@ impl Membership<SeqTypes> for EpochCommittees {
             })
     }
 
-    fn add_drb_result(&mut self, epoch: Epoch, drb_result: DrbResult) {
-        self.drb_result_table.insert(epoch, drb_result);
+    fn add_drb_result(&mut self, epoch: Epoch, drb: DrbResult) {
+        let Some(raw_stake_table) = self.state.get(&epoch) else {
+            tracing::error!("add_drb_result({}, {:?}) was called, but we do not yet have the stake table for epoch {}", epoch, drb, epoch);
+            return;
+        };
+
+        let randomized_committee =
+            generate_stake_cdf(raw_stake_table.eligible_leaders.clone(), drb);
+
+        self.randomized_committees
+            .insert(epoch, randomized_committee);
     }
 }
 
