Clean up the Substrate block processing code

coordinator/src/main.rs

@@ -2,6 +2,8 @@
 #![allow(unreachable_code)]
 #![allow(clippy::diverging_sub_expression)]
 
+use std::time::Duration;
+
 use zeroize::Zeroizing;
 
 use ciphersuite::{group::ff::Field, Ciphersuite, Ristretto};
@@ -9,6 +11,8 @@ use ciphersuite::{group::ff::Field, Ciphersuite, Ristretto};
 use serai_db::{Db, MemDb};
 use serai_client::Serai;
 
+use tokio::time::sleep;
+
 mod db;
 pub use db::*;
 
@@ -33,16 +37,26 @@ async fn run<D: Db, P: P2p>(
 
   let mut last_substrate_block = db.last_substrate_block();
 
+  tokio::spawn(async move {
     loop {
-    match substrate::handle_new_blocks(&mut db, &key, &p2p, &serai, &mut last_substrate_block).await
+      match substrate::handle_new_blocks(&mut db, &key, &p2p, &serai, &mut last_substrate_block)
+        .await
       {
         Ok(()) => {}
-      Err(e) => log::error!("couldn't communicate with serai node: {e}"),
+        Err(e) => {
+          log::error!("couldn't communicate with serai node: {e}");
+          sleep(Duration::from_secs(5)).await;
         }
+      }
+    }
+  });
 
+  loop {
     // Handle all messages from tributaries
 
     // Handle all messages from processors
+
+    todo!()
   }
 }
 

coordinator/src/substrate.rs

@@ -1,5 +1,5 @@
-use core::{time::Duration, ops::Deref};
+use core::ops::Deref;
-use std::collections::HashMap;
+use std::collections::{HashSet, HashMap};
 
 use zeroize::Zeroizing;
 
@@ -7,13 +7,11 @@ use transcript::{Transcript, RecommendedTranscript};
 use ciphersuite::{group::GroupEncoding, Ciphersuite, Ristretto};
 use frost::{Participant, ThresholdParams};
 
-use tokio::time::sleep;
-
 use serai_client::{
   SeraiError, Block, Serai,
   primitives::BlockHash,
   validator_sets::{
-    primitives::{Session, ValidatorSet, ValidatorSetData},
+    primitives::{Session, ValidatorSet, KeyPair},
     ValidatorSetsEvent,
   },
   in_instructions::InInstructionsEvent,
@@ -26,77 +24,55 @@ use processor_messages::{SubstrateContext, key_gen::KeyGenId};
 
 use crate::{Db, MainDb, TributaryTransaction, P2p};
 
-async fn get_set(serai: &Serai, set: ValidatorSet) -> ValidatorSetData {
+async fn get_coin_key(serai: &Serai, set: ValidatorSet) -> Result<Option<Vec<u8>>, SeraiError> {
-  loop {
+  Ok(serai.get_keys(set).await?.map(|keys| keys.1.into_inner()))
-    match serai.get_validator_set(set).await {
-      Ok(data) => return data.unwrap(),
-      Err(e) => {
-        log::error!("couldn't get validator set data: {e}");
-        sleep(Duration::from_secs(5)).await;
-      }
-    }
-  }
-}
-
-async fn get_coin_keys(serai: &Serai, set: ValidatorSet) -> Vec<u8> {
-  loop {
-    match serai.get_keys(set).await {
-      Ok(data) => return data.unwrap().1.into_inner(),
-      Err(e) => {
-        log::error!("couldn't get validator set's keys: {e}");
-        sleep(Duration::from_secs(5)).await;
-      }
-    }
-  }
 }
 
 async fn in_set(
   key: &Zeroizing<<Ristretto as Ciphersuite>::F>,
   serai: &Serai,
   set: ValidatorSet,
-) -> bool {
+) -> Result<Option<Option<Participant>>, SeraiError> {
-  let data = get_set(serai, set).await;
+  let Some(data) = serai.get_validator_set(set).await? else {
-  let key = Ristretto::generator() * key.deref();
+    return Ok(None);
-  data.participants.iter().any(|(participant, _)| participant.0 == key.to_bytes())
+  };
+  let key = (Ristretto::generator() * key.deref()).to_bytes();
+  Ok(Some(
+    data
+      .participants
+      .iter()
+      .position(|(participant, _)| participant.0 == key)
+      .map(|index| Participant::new((index + 1).try_into().unwrap()).unwrap()),
+  ))
 }
 
-async fn handle_block<D: Db, P: P2p>(
+async fn handle_new_set<D: Db, P: P2p>(
   db: &mut MainDb<D>,
   key: &Zeroizing<<Ristretto as Ciphersuite>::F>,
   p2p: &P,
   serai: &Serai,
-  block: Block,
+  block: &Block,
+  set: ValidatorSet,
 ) -> Result<(), SeraiError> {
-  let hash = block.hash();
+  if let Some(i) = in_set(key, serai, set).await?.expect("NewSet for set which doesn't exist") {
+    let set_data = serai.get_validator_set(set).await?.expect("NewSet for set which doesn't exist");
 
-  let mut event_id = 0;
+    let n = u16::try_from(set_data.participants.len()).unwrap();
+    let t = (2 * (n / 3)) + 1;
 
-  // If a new validator set was activated, create tributary/inform processor to do a DKG
-  for new_set in serai.get_new_set_events(hash).await? {
-    if !db.handled_event(hash, event_id) {
-      if let ValidatorSetsEvent::NewSet { set } = new_set {
-        let set_data = serai.get_validator_set(set).await?.unwrap();
-
-        let mut i = None;
     let mut validators = HashMap::new();
     for (l, (participant, amount)) in set_data.participants.iter().enumerate() {
-          // TODO2: Ensure an invalid public key can't be a validator
+      // TODO: Ban invalid keys from being validators on the Serai side
-          let participant =
+      let participant = <Ristretto as Ciphersuite>::read_G::<&[u8]>(&mut participant.0.as_ref())
-            <Ristretto as Ciphersuite>::read_G::<&[u8]>(&mut participant.0.as_ref()).unwrap();
+        .expect("invalid key registered as participant");
-          if participant == (Ristretto::generator() * key.deref()) {
-            i = Some(Participant::new((l + 1).try_into().unwrap()).unwrap());
-          }
-
       // Give one weight on Tributary per bond instance
       validators.insert(participant, amount.0 / set_data.bond.0);
     }
 
-        if let Some(i) = i {
+    // Calculate the genesis for this Tributary
-          let n = u16::try_from(set_data.participants.len()).unwrap();
-          let t = (2 * (n / 3)) + 1;
-
     let mut genesis = RecommendedTranscript::new(b"Serai Tributary Genesis");
-          genesis.append_message(b"serai_block", hash);
+    // This locks it to a specific Serai chain
+    genesis.append_message(b"serai_block", block.hash());
     genesis.append_message(b"session", set.session.0.to_le_bytes());
     genesis.append_message(b"network", set.network.0.to_le_bytes());
     let genesis = genesis.challenge(b"genesis");
@@ -108,7 +84,7 @@ async fn handle_block<D: Db, P: P2p>(
       // TODO2: Use a DB on a dedicated volume
       db.0.clone(),
       genesis,
-            block.time().unwrap(),
+      block.time().expect("Serai block didn't have a timestamp set"),
       key.clone(),
       validators,
       p2p.clone(),
@@ -118,11 +94,176 @@ async fn handle_block<D: Db, P: P2p>(
 
     // Trigger a DKG
     // TODO: Send this to processor. Check how it handles it being fired multiple times
+    // We already have a unique event ID based on block, event index (where event index is
+    // the one generated in this handle_block function)
+    // We could use that on this end and the processor end?
     let msg = processor_messages::key_gen::CoordinatorMessage::GenerateKey {
       id: KeyGenId { set, attempt: 0 },
       params: ThresholdParams::new(t, n, i).unwrap(),
     };
   }
+
+  Ok(())
+}
+
+async fn handle_key_gen<D: Db>(
+  db: &mut MainDb<D>,
+  key: &Zeroizing<<Ristretto as Ciphersuite>::F>,
+  serai: &Serai,
+  block: &Block,
+  set: ValidatorSet,
+  key_pair: KeyPair,
+) -> Result<(), SeraiError> {
+  if in_set(key, serai, set)
+    .await?
+    .expect("KeyGen occurred for a set which doesn't exist")
+    .is_some()
+  {
+    // TODO: Send this to processor. Check how it handles it being fired multiple times
+    let msg = processor_messages::key_gen::CoordinatorMessage::ConfirmKeyPair {
+      context: SubstrateContext {
+        coin_latest_finalized_block: serai
+          .get_latest_block_for_network(block.hash(), set.network)
+          .await?
+          .unwrap_or(BlockHash([0; 32])), // TODO: Have the processor override this
+      },
+      // TODO: Check the DB for which attempt used this key pair
+      id: KeyGenId { set, attempt: todo!() },
+    };
+  }
+
+  Ok(())
+}
+
+async fn handle_batch_and_burns<D: Db>(
+  db: &mut MainDb<D>,
+  key: &Zeroizing<<Ristretto as Ciphersuite>::F>,
+  serai: &Serai,
+  block: &Block,
+) -> Result<(), SeraiError> {
+  let hash = block.hash();
+
+  // Track which networks had events with a Vec in ordr to preserve the insertion order
+  // While that shouldn't be needed, ensuring order never hurts, and may enable design choices
+  // with regards to Processor <-> Coordinator message passing
+  let mut networks_with_event = vec![];
+  let mut network_had_event = |burns: &mut HashMap<_, _>, network| {
+    // Don't insert this network multiple times
+    // A Vec is still used in order to maintain the insertion order
+    if !networks_with_event.contains(&network) {
+      networks_with_event.push(network);
+      burns.insert(network, vec![]);
+    }
+  };
+
+  let mut batch_block = HashMap::new();
+  let mut burns = HashMap::new();
+
+  for batch in serai.get_batch_events(hash).await? {
+    if let InInstructionsEvent::Batch { network, id: _, block: network_block } = batch {
+      network_had_event(&mut burns, network);
+
+      // Track what Serai acknowledges as the latest block for this network
+      // If this Substrate block has multiple batches, the last batch's block will overwrite the
+      // prior batches
+      // Since batches within a block are guaranteed to be ordered, thanks to their incremental ID,
+      // the last batch will be the latest batch, so its block will be the latest block
+      batch_block.insert(network, network_block);
+
+      // TODO: Send this to processor. Check how it handles it being fired multiple times
+      let msg = processor_messages::coordinator::CoordinatorMessage::BatchSigned {
+        key: get_coin_key(
+          serai,
+          // TODO2
+          ValidatorSet { network, session: Session(0) },
+        )
+        .await?
+        .expect("ValidatorSet without keys signed a batch"),
+        block: network_block,
+      };
+    } else {
+      panic!("Batch event wasn't Batch: {batch:?}");
+    }
+  }
+
+  for burn in serai.get_burn_events(hash).await? {
+    if let TokensEvent::Burn { address: _, balance, instruction } = burn {
+      // TODO: Move Network/Coin to an enum and provide this mapping
+      let network = {
+        use serai_client::primitives::*;
+        match balance.coin {
+          BITCOIN => BITCOIN_NET_ID,
+          ETHER => ETHEREUM_NET_ID,
+          DAI => ETHEREUM_NET_ID,
+          MONERO => MONERO_NET_ID,
+          invalid => panic!("burn from unrecognized coin: {invalid:?}"),
+        }
+      };
+
+      network_had_event(&mut burns, network);
+
+      // network_had_event should register an entry in burns
+      let mut burns_so_far = burns.remove(&network).unwrap();
+      burns_so_far.push(OutInstructionWithBalance { balance, instruction });
+      burns.insert(network, burns_so_far);
+    } else {
+      panic!("Burn event wasn't Burn: {burn:?}");
+    }
+  }
+
+  assert_eq!(HashSet::<&_>::from_iter(networks_with_event.iter()).len(), networks_with_event.len());
+
+  for network in networks_with_event {
+    let coin_latest_finalized_block = if let Some(block) = batch_block.remove(&network) {
+      block
+    } else {
+      // If it's had a batch or a burn, it must have had a block acknowledged
+      serai
+        .get_latest_block_for_network(hash, network)
+        .await?
+        .expect("network had a batch/burn yet never set a latest block")
+    };
+
+    // TODO: Send this to processor. Check how it handles it being fired multiple times
+    let msg = processor_messages::substrate::CoordinatorMessage::SubstrateBlock {
+      context: SubstrateContext { coin_latest_finalized_block },
+      key: get_coin_key(
+        serai,
+        // TODO2
+        ValidatorSet { network, session: Session(0) },
+      )
+      .await?
+      .expect("batch/burn for network which never set keys"),
+      burns: burns.remove(&network).unwrap(),
+    };
+  }
+
+  Ok(())
+}
+
+// Handle a specific Substrate block, returning an error when it fails to get data
+// (not blocking / holding)
+async fn handle_block<D: Db, P: P2p>(
+  db: &mut MainDb<D>,
+  key: &Zeroizing<<Ristretto as Ciphersuite>::F>,
+  p2p: &P,
+  serai: &Serai,
+  block: Block,
+) -> Result<(), SeraiError> {
+  let hash = block.hash();
+
+  // Define an indexed event ID.
+  let mut event_id = 0;
+
+  // If a new validator set was activated, create tributary/inform processor to do a DKG
+  for new_set in serai.get_new_set_events(hash).await? {
+    // Individually mark each event as handled so on reboot, we minimize duplicates
+    // Additionally, if the Serai connection also fails 1/100 times, this means a block with 1000
+    // events will successfully be incrementally handled (though the Serai connection should be
+    // stable)
+    if !db.handled_event(hash, event_id) {
+      if let ValidatorSetsEvent::NewSet { set } = new_set {
+        handle_new_set(db, key, p2p, serai, &block, set).await?;
       } else {
         panic!("NewSet event wasn't NewSet: {new_set:?}");
       }
@@ -135,18 +276,7 @@ async fn handle_block<D: Db, P: P2p>(
   for key_gen in serai.get_key_gen_events(hash).await? {
     if !db.handled_event(hash, event_id) {
       if let ValidatorSetsEvent::KeyGen { set, key_pair } = key_gen {
-        if in_set(key, serai, set).await {
+        handle_key_gen(db, key, serai, &block, set, key_pair).await?;
-          // TODO: Send this to processor. Check how it handles it being fired multiple times
-          let msg = processor_messages::key_gen::CoordinatorMessage::ConfirmKeyPair {
-            context: SubstrateContext {
-              coin_latest_finalized_block: serai
-                .get_latest_block_for_network(hash, set.network)
-                .await?
-                .unwrap_or(BlockHash([0; 32])), // TODO: Have the processor override this
-            },
-            id: KeyGenId { set, attempt: todo!() },
-          };
-        }
       } else {
         panic!("KeyGen event wasn't KeyGen: {key_gen:?}");
       }
@@ -155,76 +285,13 @@ async fn handle_block<D: Db, P: P2p>(
     event_id += 1;
   }
 
-  if !db.handled_event(hash, event_id) {
   // Finally, tell the processor of acknowledged blocks/burns
-    let mut coins_with_event = vec![];
+  // This uses a single event as. unlike prior events which individually executed code, all
-    let mut batch_block = HashMap::new();
+  // following events share data collection
-    let mut burns = HashMap::new();
+  // This does break the uniqueness of (hash, event_id) -> one event, yet
-
+  // (network, (hash, event_id)) remains valid as a unique ID for an event
-    for batch in serai.get_batch_events(hash).await? {
+  if !db.handled_event(hash, event_id) {
-      if let InInstructionsEvent::Batch { network, id: _, block: coin_block } = batch {
+    handle_batch_and_burns(db, key, serai, &block).await?;
-        // Don't insert this multiple times, yet use a Vec to maintain the insertion order
-        if !coins_with_event.contains(&network) {
-          coins_with_event.push(network);
-          burns.insert(network, vec![]);
-        }
-
-        // Use the last specified block
-        batch_block.insert(network, coin_block);
-
-        // TODO: Send this to processor. Check how it handles it being fired multiple times
-        let msg = processor_messages::coordinator::CoordinatorMessage::BatchSigned {
-          key: get_coin_keys(serai, ValidatorSet { network, session: Session(0) }).await, // TODO2
-          block: coin_block,
-        };
-      } else {
-        panic!("Batch event wasn't Batch: {batch:?}");
-      }
-    }
-
-    for burn in serai.get_burn_events(hash).await? {
-      if let TokensEvent::Burn { address: _, balance, instruction } = burn {
-        let network = {
-          use serai_client::primitives::*;
-          match balance.coin {
-            BITCOIN => BITCOIN_NET_ID,
-            ETHER => ETHEREUM_NET_ID,
-            DAI => ETHEREUM_NET_ID,
-            MONERO => MONERO_NET_ID,
-            invalid => panic!("burn from unrecognized coin: {invalid:?}"),
-          }
-        };
-
-        if !coins_with_event.contains(&network) {
-          coins_with_event.push(network);
-          burns.insert(network, vec![]);
-        }
-
-        let mut burns_so_far = burns.remove(&network).unwrap_or(vec![]);
-        burns_so_far.push(OutInstructionWithBalance { balance, instruction });
-        burns.insert(network, burns_so_far);
-      } else {
-        panic!("Burn event wasn't Burn: {burn:?}");
-      }
-    }
-
-    for network in coins_with_event {
-      let coin_latest_finalized_block = if let Some(block) = batch_block.remove(&network) {
-        block
-      } else {
-        // If it's had a batch or a burn, it must have had a block acknowledged
-        serai.get_latest_block_for_network(hash, network).await?.unwrap()
-      };
-
-      // TODO: Send this to processor. Check how it handles it being fired multiple times
-      let msg = processor_messages::substrate::CoordinatorMessage::SubstrateBlock {
-        context: SubstrateContext { coin_latest_finalized_block },
-        key: get_coin_keys(serai, ValidatorSet { network, session: Session(0) }).await, // TODO2
-        // Use remove not only to avoid a clone, yet so if network is present twice somehow, this
-        // isn't fired multiple times
-        burns: burns.remove(&network).unwrap(),
-      };
-    }
   }
   db.handle_event(hash, event_id);