Tweak ConfirmKeyPair to alleviate database requirements of coordinator

coordinator/src/substrate.rs

@@ -124,8 +124,8 @@ async fn handle_key_gen<D: Db, Pro: Processor>(
               .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!() },
+          set,
+          key_pair,
         },
       ))
       .await;

processor/messages/src/lib.rs

@@ -9,7 +9,7 @@ use dkg::{Participant, ThresholdParams};
 use serai_primitives::BlockHash;
 use in_instructions_primitives::SignedBatch;
 use tokens_primitives::OutInstructionWithBalance;
-use validator_sets_primitives::ValidatorSet;
+use validator_sets_primitives::{ValidatorSet, KeyPair};
 
 #[derive(Clone, Copy, PartialEq, Eq, Debug, Zeroize, Serialize, Deserialize)]
 pub struct SubstrateContext {
@@ -141,11 +141,12 @@ pub mod coordinator {
 pub mod substrate {
   use super::*;
 
-  #[derive(Clone, PartialEq, Eq, Debug, Zeroize, Serialize, Deserialize)]
+  #[derive(Clone, PartialEq, Eq, Debug, Serialize, Deserialize)]
   pub enum CoordinatorMessage {
     ConfirmKeyPair {
       context: SubstrateContext,
-      id: key_gen::KeyGenId,
+      set: ValidatorSet,
+      key_pair: KeyPair,
     },
     SubstrateBlock {
       context: SubstrateContext,

processor/src/key_gen.rs

@@ -15,7 +15,10 @@ use frost::{
 
 use log::info;
 
-use serai_client::{primitives::BlockHash, validator_sets::primitives::ValidatorSet};
+use serai_client::{
+  primitives::BlockHash,
+  validator_sets::primitives::{ValidatorSet, KeyPair},
+};
 use messages::{SubstrateContext, key_gen::*};
 
 use crate::{Get, DbTxn, Db, coins::Coin};
@@ -65,8 +68,8 @@ impl<C: Coin, D: Db> KeyGenDb<C, D> {
     .unwrap()
   }
 
-  fn generated_keys_key(id: &KeyGenId) -> Vec<u8> {
-    Self::key_gen_key(b"generated_keys", bincode::serialize(id).unwrap())
+  fn generated_keys_key(set: ValidatorSet, key_pair: (&[u8], &[u8])) -> Vec<u8> {
+    Self::key_gen_key(b"generated_keys", bincode::serialize(&(set, key_pair)).unwrap())
   }
   fn save_keys(
     txn: &mut D::Transaction<'_>,
@@ -76,7 +79,13 @@ impl<C: Coin, D: Db> KeyGenDb<C, D> {
   ) {
     let mut keys = substrate_keys.serialize();
     keys.extend(coin_keys.serialize().iter());
-    txn.put(Self::generated_keys_key(id), keys);
+    txn.put(
+      Self::generated_keys_key(
+        id.set,
+        (substrate_keys.group_key().to_bytes().as_ref(), coin_keys.group_key().to_bytes().as_ref()),
+      ),
+      keys,
+    );
   }
 
   fn keys_key(key: &<C::Curve as Ciphersuite>::G) -> Vec<u8> {
@@ -96,9 +105,13 @@ impl<C: Coin, D: Db> KeyGenDb<C, D> {
   }
   fn confirm_keys(
     txn: &mut D::Transaction<'_>,
-    id: &KeyGenId,
+    set: ValidatorSet,
+    key_pair: KeyPair,
   ) -> (ThresholdKeys<Ristretto>, ThresholdKeys<C::Curve>) {
-    let (keys_vec, keys) = Self::read_keys(txn, &Self::generated_keys_key(id));
+    let (keys_vec, keys) = Self::read_keys(
+      txn,
+      &Self::generated_keys_key(set, (key_pair.0.as_ref(), key_pair.1.as_ref())),
+    );
     txn.put(Self::keys_key(&keys.1.group_key()), keys_vec);
     keys
   }
@@ -352,15 +365,16 @@ impl<C: Coin, D: Db> KeyGen<C, D> {
     &mut self,
     txn: &mut D::Transaction<'_>,
     context: SubstrateContext,
-    id: KeyGenId,
+    set: ValidatorSet,
+    key_pair: KeyPair,
   ) -> KeyConfirmed<C::Curve> {
-    let (substrate_keys, coin_keys) = KeyGenDb::<C, D>::confirm_keys(txn, &id);
+    let (substrate_keys, coin_keys) = KeyGenDb::<C, D>::confirm_keys(txn, set, key_pair);
 
     info!(
-      "Confirmed key pair {} {} from {:?}",
+      "Confirmed key pair {} {} for set {:?}",
       hex::encode(substrate_keys.group_key().to_bytes()),
       hex::encode(coin_keys.group_key().to_bytes()),
-      id
+      set,
     );
 
     KeyConfirmed {

processor/src/main.rs

@@ -301,10 +301,10 @@ async fn handle_coordinator_msg<D: Db, C: Coin, Co: Coordinator>(
 
     CoordinatorMessage::Substrate(msg) => {
       match msg {
-        messages::substrate::CoordinatorMessage::ConfirmKeyPair { context, id } => {
+        messages::substrate::CoordinatorMessage::ConfirmKeyPair { context, set, key_pair } => {
           // See TributaryMutable's struct definition for why this block is safe
           let KeyConfirmed { activation_block, substrate_keys, coin_keys } =
-            tributary_mutable.key_gen.confirm(txn, context, id).await;
+            tributary_mutable.key_gen.confirm(txn, context, set, key_pair).await;
           tributary_mutable.substrate_signers.insert(
             substrate_keys.group_key().to_bytes().to_vec(),
             SubstrateSigner::new(substrate_keys),

processor/src/tests/key_gen.rs

@@ -9,6 +9,7 @@ use frost::{Participant, ThresholdParams, tests::clone_without};
 
 use serai_db::{DbTxn, Db, MemDb};
 
+use sp_application_crypto::sr25519;
 use serai_client::{
   primitives::{MONERO_NET_ID, BlockHash},
   validator_sets::primitives::{Session, ValidatorSet},
@@ -124,6 +125,7 @@ pub async fn test_key_gen<C: Coin>() {
     }
     txn.commit();
   }
+  let res = res.unwrap();
 
   // Rebuild 1 and 4
   rebuild(&mut key_gens, &dbs, 1);
@@ -136,7 +138,8 @@ pub async fn test_key_gen<C: Coin>() {
       .confirm(
         &mut txn,
         SubstrateContext { coin_latest_finalized_block: BlockHash([0x11; 32]) },
-        ID,
+        ID.set,
+        (sr25519::Public(res.0), res.1.clone().try_into().unwrap()),
       )
       .await;
     txn.commit();
@@ -147,8 +150,8 @@ pub async fn test_key_gen<C: Coin>() {
     assert_eq!(substrate_keys.params(), params);
     assert_eq!(coin_keys.params(), params);
     assert_eq!(
-      &(substrate_keys.group_key().to_bytes(), coin_keys.group_key().to_bytes().as_ref().to_vec()),
-      res.as_ref().unwrap()
+      (substrate_keys.group_key().to_bytes(), coin_keys.group_key().to_bytes().as_ref().to_vec()),
+      res
     );
   }
 }