[fortuna] Add subsampling of hash chains (#1624)

* add hash chain test

* add sampling

* works

* works

* ok this works

* cleanup

* parallelize

* gr

* spawn

apps/fortuna/Cargo.lock

@@ -1488,7 +1488,7 @@ dependencies = [
 
 [[package]]
 name = "fortuna"
-version = "6.0.1"
+version = "6.1.0"
 dependencies = [
  "anyhow",
  "axum",

apps/fortuna/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name    = "fortuna"
-version = "6.0.1"
+version = "6.1.0"
 edition = "2021"
 
 [dependencies]

apps/fortuna/config.sample.yaml

@@ -18,6 +18,7 @@ chains:
 provider:
   uri: http://localhost:8080/
   chain_length: 100000
+  chain_sample_interval: 10
 
   # An ethereum wallet address and private key. Generate with `cast wallet new`
   address: 0xADDRESS

apps/fortuna/src/api.rs

@@ -233,11 +233,11 @@ mod test {
         // initialize the BlockchainStates below, but they aren't cloneable (nor do they need to be cloned).
         static ref ETH_CHAIN: Arc<HashChainState> = Arc::new(HashChainState::from_chain_at_offset(
             0,
-            PebbleHashChain::new([0u8; 32], 1000),
+            PebbleHashChain::new([0u8; 32], 1000, 1),
         ));
         static ref AVAX_CHAIN: Arc<HashChainState> = Arc::new(HashChainState::from_chain_at_offset(
             100,
-            PebbleHashChain::new([1u8; 32], 1000),
+            PebbleHashChain::new([1u8; 32], 1000, 1),
         ));
     }
 

apps/fortuna/src/command/register_provider.rs

@@ -65,18 +65,21 @@ pub async fn register_provider_from_config(
         .ok_or(anyhow!("Please specify a provider secret in the config"))?;
 
     let commitment_length = provider_config.chain_length;
-    let mut chain = PebbleHashChain::from_config(
+    tracing::info!("Generating hash chain");
+    let chain = PebbleHashChain::from_config(
         &secret,
         &chain_id,
         &private_key_string.parse::<LocalWallet>()?.address(),
         &chain_config.contract_addr,
         &random,
         commitment_length,
+        provider_config.chain_sample_interval,
     )?;
+    tracing::info!("Done generating hash chain");
 
     // Arguments to the contract to register our new provider.
     let fee_in_wei = chain_config.fee;
-    let commitment = chain.reveal()?;
+    let commitment = chain.reveal_ith(0)?;
     // Store the random seed and chain length in the metadata field so that we can regenerate the hash
     // chain at-will. (This is secure because you can't generate the chain unless you also have the secret)
     let commitment_metadata = CommitmentMetadata {

apps/fortuna/src/command/run.rs

@@ -36,6 +36,7 @@ use {
             BlockNumber,
         },
     },
+    futures::future::join_all,
     prometheus_client::{
         encoding::EncodeLabelSet,
         metrics::{
@@ -157,10 +158,29 @@ pub async fn run(opts: &RunOptions) -> Result<()> {
     ))?;
     let (tx_exit, rx_exit) = watch::channel(false);
 
+    let mut tasks = Vec::new();
+    for (chain_id, chain_config) in config.chains.clone() {
+        let secret_copy = secret.clone();
+
+        tasks.push(spawn(async move {
+            let state = setup_chain_state(
+                &config.provider.address,
+                &secret_copy,
+                config.provider.chain_sample_interval,
+                &chain_id,
+                &chain_config,
+            )
+            .await;
+
+            (chain_id, state)
+        }));
+    }
+    let states = join_all(tasks).await;
+
     let mut chains: HashMap<ChainId, BlockchainState> = HashMap::new();
-    for (chain_id, chain_config) in &config.chains {
-        let state =
-            setup_chain_state(&config.provider.address, &secret, chain_id, chain_config).await;
+    for result in states {
+        let (chain_id, state) = result?;
+
         match state {
             Ok(state) => {
                 chains.insert(chain_id.clone(), state);
@@ -211,6 +231,7 @@ pub async fn run(opts: &RunOptions) -> Result<()> {
 async fn setup_chain_state(
     provider: &Address,
     secret: &String,
+    chain_sample_interval: u64,
     chain_id: &ChainId,
     chain_config: &EthereumConfig,
 ) -> Result<BlockchainState> {
@@ -267,6 +288,7 @@ async fn setup_chain_state(
             &chain_config.contract_addr,
             &commitment.seed,
             commitment.chain_length,
+            chain_sample_interval,
         )?;
         hash_chains.push(pebble_hash_chain);
     }

apps/fortuna/src/command/setup_provider.rs

@@ -106,6 +106,7 @@ async fn setup_chain_provider(
             &chain_config.contract_addr,
             &metadata.seed,
             provider_config.chain_length,
+            provider_config.chain_sample_interval,
         )?;
         let chain_state = HashChainState {
             offsets:     vec![provider_info

apps/fortuna/src/config.rs

@@ -180,6 +180,15 @@ pub struct ProviderConfig {
 
     /// The length of the hash chain to generate.
     pub chain_length: u64,
+
+    /// How frequently the hash chain is sampled -- increase this value to tradeoff more
+    /// compute per request for less RAM use.
+    #[serde(default = "default_chain_sample_interval")]
+    pub chain_sample_interval: u64,
+}
+
+fn default_chain_sample_interval() -> u64 {
+    1
 }
 
 /// Configuration values for the keeper service that are shared across chains.

apps/fortuna/src/state.rs

@@ -11,28 +11,44 @@ use {
     },
 };
 
-/// A HashChain.
+/// A hash chain of a specific length. The hash chain has the property that
+/// hash(chain.reveal_ith(i)) == chain.reveal_ith(i - 1)
+///
+/// The implementation subsamples the elements of the chain such that it uses less memory
+/// to keep the chain around.
 #[derive(Clone)]
 pub struct PebbleHashChain {
-    hash: Vec<[u8; 32]>,
-    next: usize,
+    hash:            Vec<[u8; 32]>,
+    sample_interval: usize,
+    length:          usize,
 }
 
 impl PebbleHashChain {
     // Given a secret, we hash it with Keccak256 len times to get the final hash, this is an S/KEY
     // like protocol in which revealing the hashes in reverse proves knowledge.
-    pub fn new(secret: [u8; 32], length: usize) -> Self {
+    pub fn new(secret: [u8; 32], length: usize, sample_interval: usize) -> Self {
+        assert!(sample_interval > 0, "Sample interval must be positive");
         let mut hash = Vec::<[u8; 32]>::with_capacity(length);
-        hash.push(Keccak256::digest(secret).into());
-        for _ in 1..length {
-            hash.push(Keccak256::digest(&hash[hash.len() - 1]).into());
+        let mut current: [u8; 32] = Keccak256::digest(secret).into();
+
+        hash.push(current.clone());
+        for i in 1..length {
+            current = Keccak256::digest(&current).into();
+            if i % sample_interval == 0 {
+                hash.push(current);
+            }
         }
 
         hash.reverse();
 
-        Self { hash, next: 0 }
+        Self {
+            hash,
+            sample_interval,
+            length,
+        }
     }
 
+
     pub fn from_config(
         secret: &str,
         chain_id: &ChainId,
@@ -40,6 +56,7 @@ impl PebbleHashChain {
         contract_address: &Address,
         random: &[u8; 32],
         chain_length: u64,
+        sample_interval: u64,
     ) -> Result<Self> {
         let mut input: Vec<u8> = vec![];
         input.extend_from_slice(&hex::decode(secret.trim())?);
@@ -49,24 +66,32 @@ impl PebbleHashChain {
         input.extend_from_slice(random);
 
         let secret: [u8; 32] = Keccak256::digest(input).into();
-        Ok(Self::new(secret, chain_length.try_into()?))
-    }
-
-    /// Reveal the next hash in the chain using the previous proof.
-    pub fn reveal(&mut self) -> Result<[u8; 32]> {
-        ensure!(self.next < self.len(), "no more hashes in the chain");
-        let next = self.hash[self.next].clone();
-        self.next += 1;
-        Ok(next)
+        Ok(Self::new(
+            secret,
+            chain_length.try_into()?,
+            sample_interval.try_into()?,
+        ))
     }
 
     pub fn reveal_ith(&self, i: usize) -> Result<[u8; 32]> {
         ensure!(i < self.len(), "index not in range");
-        Ok(self.hash[i].clone())
+
+        // Note that subsample_interval may not perfectly divide length, in which case the uneven segment is
+        // actually at the *front* of the list. Thus, it's easier to compute indexes from the end of the list.
+        let index_from_end_of_subsampled_list = ((self.len() - 1) - i) / self.sample_interval;
+        let mut i_index = self.len() - 1 - index_from_end_of_subsampled_list * self.sample_interval;
+        let mut val = self.hash[self.hash.len() - 1 - index_from_end_of_subsampled_list].clone();
+
+        while i_index > i {
+            val = Keccak256::digest(val).into();
+            i_index -= 1;
+        }
+
+        Ok(val)
     }
 
     pub fn len(&self) -> usize {
-        self.hash.len()
+        self.length
     }
 }
 
@@ -99,3 +124,56 @@ impl HashChainState {
         self.hash_chains[chain_index].reveal_ith(sequence_number - self.offsets[chain_index])
     }
 }
+
+#[cfg(test)]
+mod test {
+    use {
+        crate::state::PebbleHashChain,
+        sha3::{
+            Digest,
+            Keccak256,
+        },
+    };
+
+    fn run_hash_chain_test(secret: [u8; 32], length: usize, sample_interval: usize) {
+        // Calculate the hash chain the naive way as a comparison point to the subsampled implementation.
+        let mut basic_chain = Vec::<[u8; 32]>::with_capacity(length);
+        let mut current: [u8; 32] = Keccak256::digest(secret).into();
+        basic_chain.push(current.clone());
+        for _ in 1..length {
+            current = Keccak256::digest(&current).into();
+            basic_chain.push(current);
+        }
+
+        basic_chain.reverse();
+
+        let chain = PebbleHashChain::new(secret, length, sample_interval);
+
+        let mut last_val = chain.reveal_ith(0).unwrap();
+        for i in 1..length {
+            let cur_val = chain.reveal_ith(i).unwrap();
+            println!("{}", i);
+            assert_eq!(basic_chain[i], cur_val);
+
+            let expected_last_val: [u8; 32] = Keccak256::digest(cur_val).into();
+            assert_eq!(expected_last_val, last_val);
+            last_val = cur_val;
+        }
+    }
+
+    #[test]
+    fn test_hash_chain() {
+        run_hash_chain_test([0u8; 32], 10, 1);
+        run_hash_chain_test([0u8; 32], 10, 2);
+        run_hash_chain_test([0u8; 32], 10, 3);
+        run_hash_chain_test([1u8; 32], 10, 1);
+        run_hash_chain_test([1u8; 32], 10, 2);
+        run_hash_chain_test([1u8; 32], 10, 3);
+        run_hash_chain_test([0u8; 32], 100, 1);
+        run_hash_chain_test([0u8; 32], 100, 2);
+        run_hash_chain_test([0u8; 32], 100, 3);
+        run_hash_chain_test([0u8; 32], 100, 7);
+        run_hash_chain_test([0u8; 32], 100, 50);
+        run_hash_chain_test([0u8; 32], 100, 55);
+    }
+}