Finish replacing key with session

coordinator/src/db.rs

@@ -34,18 +34,12 @@ impl<D: Db> MainDb<D> {
     getter.get(Self::handled_message_key(network, id)).is_some()
   }
 
-  fn in_tributary_key(set: ValidatorSet) -> Vec<u8> {
-    Self::main_key(b"in_tributary", set.encode())
-  }
   fn active_tributaries_key() -> Vec<u8> {
     Self::main_key(b"active_tributaries", [])
   }
   fn retired_tributary_key(set: ValidatorSet) -> Vec<u8> {
     Self::main_key(b"retired_tributary", set.encode())
   }
-  pub fn in_tributary<G: Get>(getter: &G, set: ValidatorSet) -> bool {
-    getter.get(Self::in_tributary_key(set)).is_some()
-  }
   pub fn active_tributaries<G: Get>(getter: &G) -> (Vec<u8>, Vec<TributarySpec>) {
     let bytes = getter.get(Self::active_tributaries_key()).unwrap_or(vec![]);
     let mut bytes_ref: &[u8] = bytes.as_ref();
@@ -58,8 +52,6 @@ impl<D: Db> MainDb<D> {
     (bytes, tributaries)
   }
   pub fn add_participating_in_tributary(txn: &mut D::Transaction<'_>, spec: &TributarySpec) {
-    txn.put(Self::in_tributary_key(spec.set()), []);
-
     let key = Self::active_tributaries_key();
     let (mut existing_bytes, existing) = Self::active_tributaries(txn);
     for tributary in &existing {

coordinator/src/main.rs

@@ -57,7 +57,7 @@ pub mod processors;
 use processors::Processors;
 
 mod substrate;
-use substrate::{CosignTransactions, SubstrateDb};
+use substrate::CosignTransactions;
 
 mod cosign_evaluator;
 use cosign_evaluator::CosignEvaluator;
@@ -205,49 +205,40 @@ async fn handle_processor_message<D: Db, P: P2p>(
     // TODO: Review replacing key with Session in messages?
     ProcessorMessage::Sign(inner_msg) => match inner_msg {
       // We'll only receive InvalidParticipant/Preprocess/Share if we're actively signing
-      sign::ProcessorMessage::InvalidParticipant { id, .. } => {
-        Some(SubstrateDb::<D>::session_for_key(&txn, &id.key).unwrap())
-      }
-      sign::ProcessorMessage::Preprocess { id, .. } => {
-        Some(SubstrateDb::<D>::session_for_key(&txn, &id.key).unwrap())
-      }
-      sign::ProcessorMessage::Share { id, .. } => {
-        Some(SubstrateDb::<D>::session_for_key(&txn, &id.key).unwrap())
-      }
+      sign::ProcessorMessage::InvalidParticipant { id, .. } => Some(id.session),
+      sign::ProcessorMessage::Preprocess { id, .. } => Some(id.session),
+      sign::ProcessorMessage::Share { id, .. } => Some(id.session),
       // While the Processor's Scanner will always emit Completed, that's routed through the
       // Signer and only becomes a ProcessorMessage::Completed if the Signer is present and
       // confirms it
       sign::ProcessorMessage::Completed { session, .. } => Some(*session),
     },
     ProcessorMessage::Coordinator(inner_msg) => match inner_msg {
-      // This is a special case as it's relevant to *all* Tributaries for this network
+      // This is a special case as it's relevant to *all* Tributaries for this network we're
+      // signing in
       // It doesn't return a Tributary to become `relevant_tributary` though
       coordinator::ProcessorMessage::SubstrateBlockAck { block, plans } => {
         // Get the sessions for these keys
-        let keys = plans.iter().map(|plan| plan.key.clone()).collect::<HashSet<_>>();
-        let mut sessions = vec![];
-        for key in keys {
-          let session = SubstrateDb::<D>::session_for_key(&txn, &key).unwrap();
-          // Only keep them if we're in the Tributary AND they haven't been retied
-          let set = ValidatorSet { network, session };
-          if MainDb::<D>::in_tributary(&txn, set) && (!MainDb::<D>::is_tributary_retired(&txn, set))
-          {
-            sessions.push((session, key));
-          }
-        }
+        let sessions = plans
+          .iter()
+          .map(|plan| plan.session)
+          .filter(|session| {
+            !MainDb::<D>::is_tributary_retired(&txn, ValidatorSet { network, session: *session })
+          })
+          .collect::<HashSet<_>>();
 
         // Ensure we have the Tributaries
-        for (session, _) in &sessions {
+        for session in &sessions {
           if !tributaries.contains_key(session) {
             return false;
           }
         }
 
-        for (session, key) in sessions {
+        for session in sessions {
           let tributary = &tributaries[&session];
           let plans = plans
             .iter()
-            .filter_map(|plan| Some(plan.id).filter(|_| plan.key == key))
+            .filter_map(|plan| Some(plan.id).filter(|_| plan.session == session))
             .collect::<Vec<_>>();
           PlanIds::set(&mut txn, &tributary.spec.genesis(), *block, &plans);
 

coordinator/src/substrate/db.rs

@@ -1,12 +1,12 @@
-use scale::{Encode, Decode};
-
-pub use serai_db::*;
+use scale::Encode;
 
 use serai_client::{
   primitives::NetworkId,
-  validator_sets::primitives::{Session, ValidatorSet, KeyPair},
+  validator_sets::primitives::{Session, ValidatorSet},
 };
 
+pub use serai_db::*;
+
 create_db! {
   SubstrateDb {
     CosignTriggered: () -> (),
@@ -86,24 +86,6 @@ impl<D: Db> SubstrateDb<D> {
     txn.put(Self::event_key(&id, index), []);
   }
 
-  fn session_key(key: &[u8]) -> Vec<u8> {
-    Self::substrate_key(b"session", key)
-  }
-  pub fn session_for_key<G: Get>(getter: &G, key: &[u8]) -> Option<Session> {
-    getter.get(Self::session_key(key)).map(|bytes| Session::decode(&mut bytes.as_ref()).unwrap())
-  }
-  pub fn save_session_for_keys(txn: &mut D::Transaction<'_>, key_pair: &KeyPair, session: Session) {
-    let session = session.encode();
-    let key_0 = Self::session_key(&key_pair.0);
-    let existing = txn.get(&key_0);
-    // This may trigger if 100% of a DKG are malicious, and they create a key equivalent to a prior
-    // key. Since it requires 100% maliciousness, not just 67% maliciousness, this will only assert
-    // in a modified-to-be-malicious stack, making it safe
-    assert!(existing.is_none() || (existing.as_ref() == Some(&session)));
-    txn.put(key_0, session.clone());
-    txn.put(Self::session_key(&key_pair.1), session);
-  }
-
   fn batch_instructions_key(network: NetworkId, id: u32) -> Vec<u8> {
     Self::substrate_key(b"batch", (network, id).encode())
   }

coordinator/src/substrate/mod.rs

@@ -30,7 +30,7 @@ use tokio::{sync::mpsc, time::sleep};
 use crate::{
   Db,
   processors::Processors,
-  tributary::{TributarySpec, SeraiBlockNumber, KeyPairDb},
+  tributary::{TributarySpec, SeraiBlockNumber},
 };
 
 mod db;
@@ -116,19 +116,13 @@ async fn handle_new_set<D: Db>(
   Ok(())
 }
 
-async fn handle_key_gen<D: Db, Pro: Processors>(
-  db: &mut D,
+async fn handle_key_gen<Pro: Processors>(
   processors: &Pro,
   serai: &Serai,
   block: &Block,
   set: ValidatorSet,
   key_pair: KeyPair,
 ) -> Result<(), SeraiError> {
-  // This has to be saved *before* we send ConfirmKeyPair
-  let mut txn = db.txn();
-  SubstrateDb::<D>::save_session_for_keys(&mut txn, &key_pair, set.session);
-  txn.commit();
-
   processors
     .send(
       set.network,
@@ -290,13 +284,7 @@ async fn handle_block<D: Db, Pro: Processors>(
     if !SubstrateDb::<D>::handled_event(&db.0, hash, event_id) {
       log::info!("found fresh key gen event {:?}", key_gen);
       if let ValidatorSetsEvent::KeyGen { set, key_pair } = key_gen {
-        // Immediately ensure this key pair is accessible to the tributary, before we fire any
-        // events off of it
-        let mut txn = db.0.txn();
-        KeyPairDb::set(&mut txn, set, &key_pair);
-        txn.commit();
-
-        handle_key_gen(&mut db.0, processors, serai, &block, set, key_pair).await?;
+        handle_key_gen(processors, serai, &block, set, key_pair).await?;
       } else {
         panic!("KeyGen event wasn't KeyGen: {key_gen:?}");
       }

coordinator/src/tributary/db.rs

@@ -5,7 +5,7 @@ use zeroize::Zeroizing;
 use ciphersuite::{Ciphersuite, Ristretto, group::GroupEncoding};
 use frost::Participant;
 
-use serai_client::validator_sets::primitives::{ValidatorSet, KeyPair};
+use serai_client::validator_sets::primitives::KeyPair;
 
 use processor_messages::coordinator::SubstrateSignableId;
 
@@ -50,7 +50,6 @@ create_db!(
     PlanIds: (genesis: &[u8], block: u64) -> Vec<[u8; 32]>,
     ConfirmationNonces: (genesis: [u8; 32], attempt: u32) -> HashMap<Participant, Vec<u8>>,
     CurrentlyCompletingKeyPair: (genesis: [u8; 32]) -> KeyPair,
-    KeyPairDb: (set: ValidatorSet) -> KeyPair,
     AttemptDb: (genesis: [u8; 32], topic: &Topic) -> u32,
     DataReceived: (genesis: [u8; 32], data_spec: &DataSpecification) -> u16,
     DataDb: (genesis: [u8; 32], data_spec: &DataSpecification, signer_bytes: &[u8; 32]) -> Vec<u8>,

coordinator/src/tributary/handle.rs

@@ -30,7 +30,7 @@ use crate::{
     nonce_decider::NonceDecider,
     dkg_confirmer::DkgConfirmer,
     scanner::{RecognizedIdType, RIDTrait},
-    FatallySlashed, DkgShare, PlanIds, ConfirmationNonces, KeyPairDb, AttemptDb, DataDb,
+    FatallySlashed, DkgShare, PlanIds, ConfirmationNonces, AttemptDb, DataDb,
   },
 };
 
@@ -633,7 +633,6 @@ pub(crate) async fn handle_application_tx<
       let Ok(_) = check_sign_data_len::<D>(txn, spec, data.signed.signer, data.data.len()) else {
         return;
       };
-      let key_pair = KeyPairDb::get(txn, spec.set());
       match handle(
         txn,
         &DataSpecification {
@@ -651,14 +650,7 @@ pub(crate) async fn handle_application_tx<
             .send(
               spec.set().network,
               sign::CoordinatorMessage::Preprocesses {
-                id: SignId {
-                  key: key_pair
-                    .expect("completed SignPreprocess despite not setting the key pair")
-                    .1
-                    .into(),
-                  id: data.plan,
-                  attempt: data.attempt,
-                },
+                id: SignId { session: spec.set().session, id: data.plan, attempt: data.attempt },
                 preprocesses,
               },
             )
@@ -672,7 +664,6 @@ pub(crate) async fn handle_application_tx<
       let Ok(_) = check_sign_data_len::<D>(txn, spec, data.signed.signer, data.data.len()) else {
         return;
       };
-      let key_pair = KeyPairDb::get(txn, spec.set());
       match handle(
         txn,
         &DataSpecification {
@@ -689,14 +680,7 @@ pub(crate) async fn handle_application_tx<
             .send(
               spec.set().network,
               sign::CoordinatorMessage::Shares {
-                id: SignId {
-                  key: key_pair
-                    .expect("completed SignShares despite not setting the key pair")
-                    .1
-                    .into(),
-                  id: data.plan,
-                  attempt: data.attempt,
-                },
+                id: SignId { session: spec.set().session, id: data.plan, attempt: data.attempt },
                 shares,
               },
             )
@@ -724,13 +708,15 @@ pub(crate) async fn handle_application_tx<
       };
 
       // TODO: Confirm this signer hasn't prior published a completion
-      let Some(key_pair) = KeyPairDb::get(txn, spec.set()) else {
-        panic!("SignCompleted for recognized plan ID despite not having a key pair for this set")
-      };
+
       processors
         .send(
           spec.set().network,
-          sign::CoordinatorMessage::Completed { key: key_pair.1.to_vec(), id: plan, tx: tx_hash },
+          sign::CoordinatorMessage::Completed {
+            session: spec.set().session,
+            id: plan,
+            tx: tx_hash,
+          },
         )
         .await;
     }

processor/messages/src/lib.rs

@@ -106,7 +106,7 @@ pub mod sign {
 
   #[derive(Clone, PartialEq, Eq, Hash, Debug, Encode, Decode, BorshSerialize, BorshDeserialize)]
   pub struct SignId {
-    pub key: Vec<u8>,
+    pub session: Session,
     pub id: [u8; 32],
     pub attempt: u32,
   }
@@ -120,20 +120,20 @@ pub mod sign {
     // Re-attempt a signing protocol.
     Reattempt { id: SignId },
     // Completed a signing protocol already.
-    Completed { key: Vec<u8>, id: [u8; 32], tx: Vec<u8> },
+    Completed { session: Session, id: [u8; 32], tx: Vec<u8> },
   }
 
   impl CoordinatorMessage {
     pub fn required_block(&self) -> Option<BlockHash> {
       None
     }
 
-    pub fn key(&self) -> &[u8] {
+    pub fn session(&self) -> Session {
       match self {
-        CoordinatorMessage::Preprocesses { id, .. } => &id.key,
-        CoordinatorMessage::Shares { id, .. } => &id.key,
-        CoordinatorMessage::Reattempt { id } => &id.key,
-        CoordinatorMessage::Completed { key, .. } => key,
+        CoordinatorMessage::Preprocesses { id, .. } => id.session,
+        CoordinatorMessage::Shares { id, .. } => id.session,
+        CoordinatorMessage::Reattempt { id } => id.session,
+        CoordinatorMessage::Completed { session, .. } => *session,
       }
     }
   }
@@ -204,7 +204,7 @@ pub mod coordinator {
 
   #[derive(Clone, PartialEq, Eq, Debug, BorshSerialize, BorshDeserialize)]
   pub struct PlanMeta {
-    pub key: Vec<u8>,
+    pub session: Session,
     pub id: [u8; 32],
   }