Make key_gen a gadget, add ConfirmKeyPair

tests/processor/src/tests/mod.rs

@@ -1,13 +1,15 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, time::SystemTime};
 
 use zeroize::Zeroizing;
 
 use ciphersuite::{Ciphersuite, Ristretto};
 use dkg::{Participant, ThresholdParams, tests::clone_without};
 
-use serai_primitives::NetworkId;
+use serai_primitives::{NetworkId, BlockHash, PublicKey};
 use serai_validator_sets_primitives::{Session, ValidatorSet};
 
+use messages::{SubstrateContext, key_gen::KeyGenId, CoordinatorMessage, ProcessorMessage};
+
 use serai_message_queue::{Service, Metadata, client::MessageQueue};
 
 use dockertest::{DockerOperations, DockerTest};
@@ -17,97 +19,85 @@ use crate::*;
 const COORDINATORS: usize = 4;
 const THRESHOLD: usize = ((COORDINATORS * 2) / 3) + 1;
 
+struct Coordinator {
+  next_send_id: u64,
+  next_recv_id: u64,
+  queue: MessageQueue,
+}
+
 fn coordinator_queue(
   ops: &DockerOperations,
   handle: String,
   coord_key: <Ristretto as Ciphersuite>::F,
-) -> MessageQueue {
+) -> Coordinator {
   let rpc = ops.handle(&handle).host_port(2287).unwrap();
   let rpc = rpc.0.to_string() + ":" + &rpc.1.to_string();
-  MessageQueue::new(Service::Coordinator, rpc, Zeroizing::new(coord_key))
+  Coordinator {
+    next_send_id: 0,
+    next_recv_id: 0,
+    queue: MessageQueue::new(Service::Coordinator, rpc, Zeroizing::new(coord_key)),
+  }
 }
 
-// Receive a message from a processor via its coordinator
+// Send a message to a processor via its coordinator
-async fn recv_message(
+async fn send_message(coordinator: &mut Coordinator, network: NetworkId, msg: CoordinatorMessage) {
-  coordinator: &MessageQueue,
-  from: NetworkId,
-  id: u64,
-) -> messages::ProcessorMessage {
-  let msg =
-    tokio::time::timeout(core::time::Duration::from_secs(10), coordinator.next(id)).await.unwrap();
-  assert_eq!(msg.from, Service::Processor(from));
-  assert_eq!(msg.id, id);
-  coordinator.ack(id).await;
-  serde_json::from_slice(&msg.msg).unwrap()
-}
-
-// Perform an interaction with all processors via their coordinators
-async fn interact_with_all<
-  FS: Fn(Participant) -> messages::key_gen::CoordinatorMessage,
-  FR: FnMut(Participant, messages::key_gen::ProcessorMessage),
->(
-  id: u64,
-  coordinators: &[MessageQueue],
-  network: NetworkId,
-  message: FS,
-  mut recv: FR,
-) {
-  for (i, coordinator) in coordinators.iter().enumerate() {
-    let participant = Participant::new(u16::try_from(i + 1).unwrap()).unwrap();
   coordinator
+    .queue
     .queue(
       Metadata {
         from: Service::Coordinator,
         to: Service::Processor(network),
-          intent: id.to_le_bytes().to_vec(),
+        intent: coordinator.next_send_id.to_le_bytes().to_vec(),
       },
-        serde_json::to_string(&messages::CoordinatorMessage::KeyGen(message(participant)))
+      serde_json::to_string(&msg).unwrap().into_bytes(),
-          .unwrap()
-          .into_bytes(),
     )
     .await;
+  coordinator.next_send_id += 1;
+}
 
-    match recv_message(coordinator, network, id).await {
+// Receive a message from a processor via its coordinator
-      messages::ProcessorMessage::KeyGen(msg) => recv(participant, msg),
+async fn recv_message(coordinator: &mut Coordinator, from: NetworkId) -> ProcessorMessage {
+  let msg = tokio::time::timeout(
+    core::time::Duration::from_secs(10),
+    coordinator.queue.next(coordinator.next_recv_id),
+  )
+  .await
+  .unwrap();
+  assert_eq!(msg.from, Service::Processor(from));
+  assert_eq!(msg.id, coordinator.next_recv_id);
+  coordinator.queue.ack(coordinator.next_recv_id).await;
+  coordinator.next_recv_id += 1;
+  serde_json::from_slice(&msg.msg).unwrap()
+}
+
+async fn key_gen(coordinators: &mut [Coordinator], network: NetworkId) {
+  // Perform an interaction with all processors via their coordinators
+  async fn interact_with_all<
+    FS: Fn(Participant) -> messages::key_gen::CoordinatorMessage,
+    FR: FnMut(Participant, messages::key_gen::ProcessorMessage),
+  >(
+    coordinators: &mut [Coordinator],
+    network: NetworkId,
+    message: FS,
+    mut recv: FR,
+  ) {
+    for (i, coordinator) in coordinators.iter_mut().enumerate() {
+      let participant = Participant::new(u16::try_from(i + 1).unwrap()).unwrap();
+      send_message(coordinator, network, CoordinatorMessage::KeyGen(message(participant))).await;
+
+      match recv_message(coordinator, network).await {
+        ProcessorMessage::KeyGen(msg) => recv(participant, msg),
         _ => panic!("processor didn't return KeyGen message"),
       }
     }
-}
-
-#[test]
-fn key_gen() {
-  for network in [NetworkId::Bitcoin, NetworkId::Monero] {
-    let mut coordinators = vec![];
-    let mut test = DockerTest::new();
-    for _ in 0 .. COORDINATORS {
-      let (coord_handle, coord_key, compositions) = processor_stack(network);
-      coordinators.push((coord_handle, coord_key));
-      for composition in compositions {
-        test.add_composition(composition);
   }
-    }
-
-    test.run(|ops| async move {
-      // Sleep for a second for the message-queue to boot
-      // It isn't an error to start immediately, it just silences an error
-      tokio::time::sleep(core::time::Duration::from_secs(1)).await;
-
-      // Connect to the Message Queues as the coordinator
-      let coordinators = coordinators
-        .into_iter()
-        .map(|(handle, key)| coordinator_queue(&ops, handle, key))
-        .collect::<Vec<_>>();
 
   // Order a key gen
-      let id = messages::key_gen::KeyGenId {
+  let id = KeyGenId { set: ValidatorSet { session: Session(0), network }, attempt: 0 };
-        set: ValidatorSet { session: Session(0), network },
-        attempt: 0,
-      };
 
   let mut commitments = HashMap::new();
   interact_with_all(
-        0,
+    coordinators,
-        &coordinators,
     network,
     |participant| messages::key_gen::CoordinatorMessage::GenerateKey {
       id,
@@ -134,8 +124,7 @@ fn key_gen() {
   // Send the commitments to all parties
   let mut shares = HashMap::new();
   interact_with_all(
-        1,
+    coordinators,
-        &coordinators,
     network,
     |participant| messages::key_gen::CoordinatorMessage::Commitments {
       id,
@@ -155,8 +144,7 @@ fn key_gen() {
   let mut substrate_key = None;
   let mut coin_key = None;
   interact_with_all(
-        2,
+    coordinators,
-        &coordinators,
     network,
     |participant| messages::key_gen::CoordinatorMessage::Shares {
       id,
@@ -185,6 +173,57 @@ fn key_gen() {
     },
   )
   .await;
+
+  // Confirm the key pair
+  // TODO: Beter document coin_latest_finalized_block's genesis state, and error if a set claims
+  // [0; 32] was finalized
+  let context = SubstrateContext {
+    serai_time: SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs(),
+    coin_latest_finalized_block: BlockHash([0; 32]),
+  };
+  for coordinator in coordinators {
+    send_message(
+      coordinator,
+      network,
+      CoordinatorMessage::Substrate(messages::substrate::CoordinatorMessage::ConfirmKeyPair {
+        context,
+        set: id.set,
+        key_pair: (
+          PublicKey::from_raw(substrate_key.unwrap()),
+          coin_key.clone().unwrap().try_into().unwrap(),
+        ),
+      }),
+    )
+    .await;
+  }
+  tokio::time::sleep(core::time::Duration::from_secs(5)).await;
+}
+
+#[test]
+fn key_gen_test() {
+  for network in [NetworkId::Bitcoin, NetworkId::Monero] {
+    let mut coordinators = vec![];
+    let mut test = DockerTest::new();
+    for _ in 0 .. COORDINATORS {
+      let (coord_handle, coord_key, compositions) = processor_stack(network);
+      coordinators.push((coord_handle, coord_key));
+      for composition in compositions {
+        test.add_composition(composition);
+      }
+    }
+
+    test.run(|ops| async move {
+      // Sleep for a second for the message-queue to boot
+      // It isn't an error to start immediately, it just silences an error
+      tokio::time::sleep(core::time::Duration::from_secs(1)).await;
+
+      // Connect to the Message Queues as the coordinator
+      let mut coordinators = coordinators
+        .into_iter()
+        .map(|(handle, key)| coordinator_queue(&ops, handle, key))
+        .collect::<Vec<_>>();
+
+      key_gen(&mut coordinators, network).await;
     });
   }
 }