Move ConfirmKeyPair from key_gen to substrate

Clarifies the emitter and accordingly why its mutations are justified.

processor/src/key_gen.rs

@@ -16,18 +16,15 @@ use frost::{
 use log::info;
 
 use serai_client::{primitives::BlockHash, validator_sets::primitives::ValidatorSet};
-use messages::key_gen::*;
+use messages::{SubstrateContext, key_gen::*};
 
 use crate::{Get, DbTxn, Db, coins::Coin};
 
 #[derive(Debug)]
-pub enum KeyGenEvent<C: Ciphersuite> {
-  KeyConfirmed {
-    activation_block: BlockHash,
-    substrate_keys: ThresholdKeys<Ristretto>,
-    coin_keys: ThresholdKeys<C>,
-  },
-  ProcessorMessage(ProcessorMessage),
+pub struct KeyConfirmed<C: Ciphersuite> {
+  pub activation_block: BlockHash,
+  pub substrate_keys: ThresholdKeys<Ristretto>,
+  pub coin_keys: ThresholdKeys<C>,
 }
 
 #[derive(Clone, Debug)]
@@ -145,7 +142,7 @@ impl<C: Coin, D: Db> KeyGen<C, D> {
     self.db.keys(key)
   }
 
-  pub async fn handle(&mut self, msg: CoordinatorMessage) -> KeyGenEvent<C::Curve> {
+  pub async fn handle(&mut self, msg: CoordinatorMessage) -> ProcessorMessage {
     let context = |id: &KeyGenId| {
       // TODO2: Also embed the chain ID/genesis block
       format!(
@@ -192,7 +189,7 @@ impl<C: Coin, D: Db> KeyGen<C, D> {
         serialized.extend(commitments.1.serialize());
         self.active_commit.insert(id.set, machines);
 
-        KeyGenEvent::ProcessorMessage(ProcessorMessage::Commitments { id, commitments: serialized })
+        ProcessorMessage::Commitments { id, commitments: serialized }
       }
 
       CoordinatorMessage::Commitments { id, commitments } => {
@@ -263,7 +260,7 @@ impl<C: Coin, D: Db> KeyGen<C, D> {
         KeyGenDb::<C, D>::save_commitments(&mut txn, &id, &commitments);
         txn.commit();
 
-        KeyGenEvent::ProcessorMessage(ProcessorMessage::Shares { id, shares })
+        ProcessorMessage::Shares { id, shares }
       }
 
       CoordinatorMessage::Shares { id, shares } => {
@@ -346,31 +343,35 @@ impl<C: Coin, D: Db> KeyGen<C, D> {
 
         let mut coin_keys = ThresholdKeys::new(coin_keys);
         C::tweak_keys(&mut coin_keys);
-        KeyGenEvent::ProcessorMessage(ProcessorMessage::GeneratedKeyPair {
+        ProcessorMessage::GeneratedKeyPair {
           id,
           substrate_key: substrate_keys.group_key().to_bytes(),
           coin_key: coin_keys.group_key().to_bytes().as_ref().to_vec(),
-        })
-      }
-
-      CoordinatorMessage::ConfirmKeyPair { context, id } => {
-        let mut txn = self.db.0.txn();
-        let (substrate_keys, coin_keys) = KeyGenDb::<C, D>::confirm_keys(&mut txn, &id);
-        txn.commit();
-
-        info!(
-          "Confirmed key pair {} {} from {:?}",
-          hex::encode(substrate_keys.group_key().to_bytes()),
-          hex::encode(coin_keys.group_key().to_bytes()),
-          id
-        );
-
-        KeyGenEvent::KeyConfirmed {
-          activation_block: context.coin_latest_finalized_block,
-          substrate_keys,
-          coin_keys,
         }
       }
     }
   }
+
+  pub async fn confirm(
+    &mut self,
+    context: SubstrateContext,
+    id: KeyGenId,
+  ) -> KeyConfirmed<C::Curve> {
+    let mut txn = self.db.0.txn();
+    let (substrate_keys, coin_keys) = KeyGenDb::<C, D>::confirm_keys(&mut txn, &id);
+    txn.commit();
+
+    info!(
+      "Confirmed key pair {} {} from {:?}",
+      hex::encode(substrate_keys.group_key().to_bytes()),
+      hex::encode(coin_keys.group_key().to_bytes()),
+      id
+    );
+
+    KeyConfirmed {
+      activation_block: context.coin_latest_finalized_block,
+      substrate_keys,
+      coin_keys,
+    }
+  }
 }

processor/src/main.rs

@@ -42,7 +42,7 @@ use coins::Bitcoin;
 use coins::Monero;
 
 mod key_gen;
-use key_gen::{KeyGenEvent, KeyGen};
+use key_gen::{KeyConfirmed, KeyGen};
 
 mod signer;
 use signer::{SignerEvent, Signer};
@@ -288,11 +288,24 @@ async fn handle_coordinator_msg<D: Db, C: Coin, Co: Coordinator>(
 
   match msg.msg.clone() {
     CoordinatorMessage::KeyGen(msg) => {
-      match tributary_mutable.key_gen.handle(msg).await {
-        // This should only occur when Substrate confirms a key, enabling access of
-        // substrate_mutable
-        // TODO: Move this under Substrate accordingly
-        KeyGenEvent::KeyConfirmed { activation_block, substrate_keys, coin_keys } => {
+      // TODO: This may be fired multiple times. What's our plan for that?
+      coordinator.send(ProcessorMessage::KeyGen(tributary_mutable.key_gen.handle(msg).await)).await;
+    }
+
+    CoordinatorMessage::Sign(msg) => {
+      tributary_mutable.signers.get_mut(msg.key()).unwrap().handle(msg).await;
+    }
+
+    CoordinatorMessage::Coordinator(msg) => {
+      tributary_mutable.substrate_signers.get_mut(msg.key()).unwrap().handle(msg).await;
+    }
+
+    CoordinatorMessage::Substrate(msg) => {
+      match msg {
+        messages::substrate::CoordinatorMessage::ConfirmKeyPair { context, id } => {
+          // See TributaryMutable's struct definition for why this block is safe
+          let KeyConfirmed { activation_block, substrate_keys, coin_keys } =
+            tributary_mutable.key_gen.confirm(context, id).await;
           tributary_mutable.substrate_signers.insert(
             substrate_keys.group_key().to_bytes().to_vec(),
             SubstrateSigner::new(raw_db.clone(), substrate_keys),
@@ -312,29 +325,13 @@ async fn handle_coordinator_msg<D: Db, C: Coin, Co: Coordinator>(
           substrate_mutable
             .schedulers
             .insert(key.to_bytes().as_ref().to_vec(), Scheduler::<C>::new(key));
+
           tributary_mutable.signers.insert(
             key.to_bytes().as_ref().to_vec(),
             Signer::new(raw_db.clone(), coin.clone(), coin_keys),
           );
         }
 
-        // TODO: This may be fired multiple times. What's our plan for that?
-        KeyGenEvent::ProcessorMessage(msg) => {
-          coordinator.send(ProcessorMessage::KeyGen(msg)).await;
-        }
-      }
-    }
-
-    CoordinatorMessage::Sign(msg) => {
-      tributary_mutable.signers.get_mut(msg.key()).unwrap().handle(msg).await;
-    }
-
-    CoordinatorMessage::Coordinator(msg) => {
-      tributary_mutable.substrate_signers.get_mut(msg.key()).unwrap().handle(msg).await;
-    }
-
-    CoordinatorMessage::Substrate(msg) => {
-      match msg {
         messages::substrate::CoordinatorMessage::SubstrateBlock {
           context,
           key: key_vec,

processor/src/tests/key_gen.rs

@@ -17,7 +17,7 @@ use serai_client::{
 use messages::{SubstrateContext, key_gen::*};
 use crate::{
   coins::Coin,
-  key_gen::{KeyGenEvent, KeyGen},
+  key_gen::{KeyConfirmed, KeyGen},
 };
 
 const ID: KeyGenId =
@@ -38,14 +38,13 @@ pub async fn test_key_gen<C: Coin>() {
   let mut all_commitments = HashMap::new();
   for i in 1 ..= 5 {
     let key_gen = key_gens.get_mut(&i).unwrap();
-    if let KeyGenEvent::ProcessorMessage(ProcessorMessage::Commitments { id, commitments }) =
-      key_gen
-        .handle(CoordinatorMessage::GenerateKey {
-          id: ID,
-          params: ThresholdParams::new(3, 5, Participant::new(u16::try_from(i).unwrap()).unwrap())
-            .unwrap(),
-        })
-        .await
+    if let ProcessorMessage::Commitments { id, commitments } = key_gen
+      .handle(CoordinatorMessage::GenerateKey {
+        id: ID,
+        params: ThresholdParams::new(3, 5, Participant::new(u16::try_from(i).unwrap()).unwrap())
+          .unwrap(),
+      })
+      .await
     {
       assert_eq!(id, ID);
       all_commitments.insert(Participant::new(u16::try_from(i).unwrap()).unwrap(), commitments);
@@ -68,7 +67,7 @@ pub async fn test_key_gen<C: Coin>() {
   for i in 1 ..= 5 {
     let key_gen = key_gens.get_mut(&i).unwrap();
     let i = Participant::new(u16::try_from(i).unwrap()).unwrap();
-    if let KeyGenEvent::ProcessorMessage(ProcessorMessage::Shares { id, shares }) = key_gen
+    if let ProcessorMessage::Shares { id, shares } = key_gen
       .handle(CoordinatorMessage::Commitments {
         id: ID,
         commitments: clone_without(&all_commitments, &i),
@@ -90,11 +89,7 @@ pub async fn test_key_gen<C: Coin>() {
   for i in 1 ..= 5 {
     let key_gen = key_gens.get_mut(&i).unwrap();
     let i = Participant::new(u16::try_from(i).unwrap()).unwrap();
-    if let KeyGenEvent::ProcessorMessage(ProcessorMessage::GeneratedKeyPair {
-      id,
-      substrate_key,
-      coin_key,
-    }) = key_gen
+    if let ProcessorMessage::GeneratedKeyPair { id, substrate_key, coin_key } = key_gen
       .handle(CoordinatorMessage::Shares {
         id: ID,
         shares: all_shares
@@ -120,27 +115,17 @@ pub async fn test_key_gen<C: Coin>() {
 
   for i in 1 ..= 5 {
     let key_gen = key_gens.get_mut(&i).unwrap();
-    if let KeyGenEvent::KeyConfirmed { activation_block, substrate_keys, coin_keys } = key_gen
-      .handle(CoordinatorMessage::ConfirmKeyPair {
-        context: SubstrateContext { coin_latest_finalized_block: BlockHash([0x11; 32]) },
-        id: ID,
-      })
-      .await
-    {
-      assert_eq!(activation_block, BlockHash([0x11; 32]));
-      let params =
-        ThresholdParams::new(3, 5, Participant::new(u16::try_from(i).unwrap()).unwrap()).unwrap();
-      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()
-      );
-    } else {
-      panic!("didn't get key back");
-    }
+    let KeyConfirmed { activation_block, substrate_keys, coin_keys } = key_gen
+      .confirm(SubstrateContext { coin_latest_finalized_block: BlockHash([0x11; 32]) }, ID)
+      .await;
+    assert_eq!(activation_block, BlockHash([0x11; 32]));
+    let params =
+      ThresholdParams::new(3, 5, Participant::new(u16::try_from(i).unwrap()).unwrap()).unwrap();
+    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()
+    );
   }
 }