From a3649b206280f3af468d96dbf86e45f470668bb6 Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Sun, 27 Aug 2023 18:23:45 -0400
Subject: [PATCH] Move where we trigger KeyGen to avoid a race condition

We only expect processor messages when we have the relevant Tributary. We
queued Tributary creation, yet then kicked off processor messages. We need to
wait until the Tributary is actually created to kick off processor messages.
---
 coordinator/src/main.rs          | 23 ++++++++++++++++++++++-
 coordinator/src/substrate/mod.rs | 25 +------------------------
 2 files changed, 23 insertions(+), 25 deletions(-)

diff --git a/coordinator/src/main.rs b/coordinator/src/main.rs
index 183c010f..1c49d650 100644
--- a/coordinator/src/main.rs
+++ b/coordinator/src/main.rs
@@ -179,6 +179,27 @@ pub async fn scan_tributaries<
         )
         .await;
 
+        // Trigger a DKG for the newly added Tributary
+        let set = spec.set();
+        processors
+          .send(
+            set.network,
+            processor_messages::CoordinatorMessage::KeyGen(
+              processor_messages::key_gen::CoordinatorMessage::GenerateKey {
+                id: processor_messages::key_gen::KeyGenId { set, attempt: 0 },
+                params: frost::ThresholdParams::new(
+                  spec.t(),
+                  spec.n(),
+                  spec
+                    .i(Ristretto::generator() * key.deref())
+                    .expect("adding a tribuary for a set we aren't in set for"),
+                )
+                .unwrap(),
+              },
+            ),
+          )
+          .await;
+
         tributary_readers.push((spec, reader));
       }
     }
@@ -438,7 +459,7 @@ pub async fn handle_processors<D: Db, Pro: Processors, P: P2p>(
           break;
         }
       }
-      spec.unwrap()
+      spec.expect("received message from processor we don't have a tributary for")
     };
 
     let genesis = spec.genesis();
diff --git a/coordinator/src/substrate/mod.rs b/coordinator/src/substrate/mod.rs
index 19d76ff9..b6bb3608 100644
--- a/coordinator/src/substrate/mod.rs
+++ b/coordinator/src/substrate/mod.rs
@@ -4,7 +4,6 @@ use std::collections::{HashSet, HashMap};
 use zeroize::Zeroizing;
 
 use ciphersuite::{group::GroupEncoding, Ciphersuite, Ristretto};
-use frost::ThresholdParams;
 
 use serai_client::{
   SeraiError, Block, Serai,
@@ -19,7 +18,7 @@ use serai_client::{
 
 use serai_db::DbTxn;
 
-use processor_messages::{SubstrateContext, key_gen::KeyGenId, CoordinatorMessage};
+use processor_messages::{SubstrateContext, CoordinatorMessage};
 
 use tokio::time::sleep;
 
@@ -86,28 +85,6 @@ async fn handle_new_set<
 
     let spec = TributarySpec::new(block.hash(), time, set, set_data);
     create_new_tributary(db, spec.clone()).await;
-
-    // Trigger a DKG
-    // TODO: Check how the processor handles this 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?
-    processors
-      .send(
-        set.network,
-        CoordinatorMessage::KeyGen(processor_messages::key_gen::CoordinatorMessage::GenerateKey {
-          id: KeyGenId { set, attempt: 0 },
-          params: ThresholdParams::new(
-            spec.t(),
-            spec.n(),
-            spec
-              .i(Ristretto::generator() * key.deref())
-              .expect("In set for a set we aren't in set for"),
-          )
-          .unwrap(),
-        }),
-      )
-      .await;
   } else {
     log::info!("not present in set {:?}", set);
   }