From 03855facd1c728d09bd558f62965b50ca334bba0 Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Sun, 26 Nov 2023 09:52:03 -0500
Subject: [PATCH] Finish replacing key with session

---
 coordinator/src/db.rs               |  8 ------
 coordinator/src/main.rs             | 41 +++++++++++------------------
 coordinator/src/substrate/db.rs     | 26 +++---------------
 coordinator/src/substrate/mod.rs    | 18 +++----------
 coordinator/src/tributary/db.rs     |  3 +--
 coordinator/src/tributary/handle.rs | 32 +++++++---------------
 processor/messages/src/lib.rs       | 16 +++++------
 processor/src/main.rs               | 27 ++++++++++---------
 processor/src/signer.rs             |  4 +--
 processor/src/tests/addresses.rs    |  4 ++-
 processor/src/tests/signer.rs       | 10 +++----
 processor/src/tests/wallet.rs       |  7 +++--
 tests/coordinator/src/tests/sign.rs | 24 +++++------------
 tests/processor/src/tests/send.rs   | 13 +++++----
 14 files changed, 82 insertions(+), 151 deletions(-)

diff --git a/coordinator/src/db.rs b/coordinator/src/db.rs
index 7be798b62..d8ec4356e 100644
--- a/coordinator/src/db.rs
+++ b/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 {
diff --git a/coordinator/src/main.rs b/coordinator/src/main.rs
index 54f212026..04f57c8ae 100644
--- a/coordinator/src/main.rs
+++ b/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);
 
diff --git a/coordinator/src/substrate/db.rs b/coordinator/src/substrate/db.rs
index b1c4c1c2b..02fe65cf9 100644
--- a/coordinator/src/substrate/db.rs
+++ b/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())
   }
diff --git a/coordinator/src/substrate/mod.rs b/coordinator/src/substrate/mod.rs
index 13b14e6b5..d304db704 100644
--- a/coordinator/src/substrate/mod.rs
+++ b/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:?}");
       }
diff --git a/coordinator/src/tributary/db.rs b/coordinator/src/tributary/db.rs
index 6b6d01ede..c48a2311f 100644
--- a/coordinator/src/tributary/db.rs
+++ b/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>,
diff --git a/coordinator/src/tributary/handle.rs b/coordinator/src/tributary/handle.rs
index 79b7a8e11..bb9152c9d 100644
--- a/coordinator/src/tributary/handle.rs
+++ b/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;
     }
diff --git a/processor/messages/src/lib.rs b/processor/messages/src/lib.rs
index 731fe71ba..fe4f8f892 100644
--- a/processor/messages/src/lib.rs
+++ b/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,7 +120,7 @@ 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 {
@@ -128,12 +128,12 @@ pub mod sign {
       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],
   }
 
diff --git a/processor/src/main.rs b/processor/src/main.rs
index 2ff30bfac..799e528f6 100644
--- a/processor/src/main.rs
+++ b/processor/src/main.rs
@@ -44,7 +44,7 @@ mod coordinator;
 pub use coordinator::*;
 
 mod key_gen;
-use key_gen::{KeyConfirmed, KeyGen};
+use key_gen::{SessionDb, KeyConfirmed, KeyGen};
 
 mod signer;
 use signer::Signer;
@@ -84,7 +84,7 @@ struct TributaryMutable<N: Network, D: Db> {
   // invalidating the Tributary's mutable borrow. The signer is coded to allow for attempted usage
   // of a dropped task.
   key_gen: KeyGen<N, D>,
-  signers: HashMap<Vec<u8>, Signer<N, D>>,
+  signers: HashMap<Session, Signer<N, D>>,
 
   // This is also mutably borrowed by the Scanner.
   // The Scanner starts new sign tasks.
@@ -206,7 +206,7 @@ async fn handle_coordinator_msg<D: Db, N: Network, Co: Coordinator>(
       }
       tributary_mutable
         .signers
-        .insert(key_pair.1.into(), Signer::new(network.clone(), session, network_keys));
+        .insert(session, Signer::new(network.clone(), session, network_keys));
     }
 
     substrate_mutable.add_key(txn, activation_number, network_key).await;
@@ -220,7 +220,7 @@ async fn handle_coordinator_msg<D: Db, N: Network, Co: Coordinator>(
     CoordinatorMessage::Sign(msg) => {
       if let Some(msg) = tributary_mutable
         .signers
-        .get_mut(msg.key())
+        .get_mut(&msg.session())
         .expect("coordinator told us to sign with a signer we don't have")
         .handle(txn, msg)
         .await
@@ -415,9 +415,9 @@ async fn handle_coordinator_msg<D: Db, N: Network, Co: Coordinator>(
                 block: substrate_block,
                 plans: to_sign
                   .iter()
-                  .map(|signable| PlanMeta {
-                    key: signable.0.to_bytes().as_ref().to_vec(),
-                    id: signable.1,
+                  .filter_map(|signable| {
+                    SessionDb::get(txn, signable.0.to_bytes().as_ref())
+                      .map(|session| PlanMeta { session, id: signable.1 })
                   })
                   .collect(),
               })
@@ -427,7 +427,8 @@ async fn handle_coordinator_msg<D: Db, N: Network, Co: Coordinator>(
           // See commentary in TributaryMutable for why this is safe
           let signers = &mut tributary_mutable.signers;
           for (key, id, tx, eventuality) in to_sign {
-            if let Some(signer) = signers.get_mut(key.to_bytes().as_ref()) {
+            if let Some(session) = SessionDb::get(txn, key.to_bytes().as_ref()) {
+              let signer = signers.get_mut(&session).unwrap();
               if let Some(msg) = signer.sign_transaction(txn, id, tx, eventuality).await {
                 coordinator.send(msg).await;
               }
@@ -516,7 +517,6 @@ async fn boot<N: Network, D: Db, Co: Coordinator>(
     let mut signer = Signer::new(network.clone(), session, network_keys);
 
     // Sign any TXs being actively signed
-    let key = key.to_bytes();
     for (plan, tx, eventuality) in &actively_signing {
       if plan.key == network_key {
         let mut txn = raw_db.txn();
@@ -530,7 +530,7 @@ async fn boot<N: Network, D: Db, Co: Coordinator>(
       }
     }
 
-    signers.insert(key.as_ref().to_vec(), signer);
+    signers.insert(session, signer);
   }
 
   // Spawn a task to rebroadcast signed TXs yet to be mined into a finalized block
@@ -629,7 +629,9 @@ async fn run<N: Network, D: Db, Co: Coordinator>(mut raw_db: D, network: N, mut
 
             if let Some((retired_key, new_key)) = retired_key_new_key {
               // Safe to mutate since all signing operations are done and no more will be added
-              tributary_mutable.signers.remove(retired_key.to_bytes().as_ref());
+              if let Some(retired_session) = SessionDb::get(&txn, retired_key.to_bytes().as_ref()) {
+                tributary_mutable.signers.remove(&retired_session);
+              }
               tributary_mutable.batch_signer.take();
               let keys = tributary_mutable.key_gen.keys(&new_key);
               if let Some((session, (substrate_keys, _))) = keys {
@@ -639,7 +641,8 @@ async fn run<N: Network, D: Db, Co: Coordinator>(mut raw_db: D, network: N, mut
             }
           },
           MultisigEvent::Completed(key, id, tx) => {
-            if let Some(signer) = tributary_mutable.signers.get_mut(&key) {
+            if let Some(session) = SessionDb::get(&txn, &key) {
+              let signer = tributary_mutable.signers.get_mut(&session).unwrap();
               if let Some(msg) = signer.completed(&mut txn, id, tx) {
                 coordinator.send(msg).await;
               }
diff --git a/processor/src/signer.rs b/processor/src/signer.rs
index 97ffa24e2..3fcb0d70b 100644
--- a/processor/src/signer.rs
+++ b/processor/src/signer.rs
@@ -370,7 +370,7 @@ impl<N: Network, D: Db> Signer<N, D> {
     // Update the attempt number
     self.attempt.insert(id, attempt);
 
-    let id = SignId { key: self.keys[0].group_key().to_bytes().as_ref().to_vec(), id, attempt };
+    let id = SignId { session: self.session, id, attempt };
 
     info!("signing for {} #{}", hex::encode(id.id), id.attempt);
 
@@ -602,7 +602,7 @@ impl<N: Network, D: Db> Signer<N, D> {
 
       CoordinatorMessage::Reattempt { id } => self.attempt(txn, id.id, id.attempt).await,
 
-      CoordinatorMessage::Completed { key: _, id, tx: mut tx_vec } => {
+      CoordinatorMessage::Completed { session: _, id, tx: mut tx_vec } => {
         let mut tx = <N::Transaction as Transaction<N>>::Id::default();
         if tx.as_ref().len() != tx_vec.len() {
           let true_len = tx_vec.len();
diff --git a/processor/src/tests/addresses.rs b/processor/src/tests/addresses.rs
index 6fec102f7..8d7baa9ed 100644
--- a/processor/src/tests/addresses.rs
+++ b/processor/src/tests/addresses.rs
@@ -7,6 +7,8 @@ use frost::{Participant, ThresholdKeys};
 
 use tokio::time::timeout;
 
+use serai_client::validator_sets::primitives::Session;
+
 use serai_db::{DbTxn, MemDb};
 
 use crate::{
@@ -50,7 +52,7 @@ async fn spend<N: Network, D: Db>(
       ),
     );
   }
-  sign(network.clone(), keys_txs).await;
+  sign(network.clone(), Session(0), keys_txs).await;
 
   for _ in 0 .. N::CONFIRMATIONS {
     network.mine_block().await;
diff --git a/processor/src/tests/signer.rs b/processor/src/tests/signer.rs
index 1ef78f0ac..9f76ae85f 100644
--- a/processor/src/tests/signer.rs
+++ b/processor/src/tests/signer.rs
@@ -2,7 +2,6 @@ use std::collections::HashMap;
 
 use rand_core::{RngCore, OsRng};
 
-use ciphersuite::group::GroupEncoding;
 use frost::{
   Participant, ThresholdKeys,
   dkg::tests::{key_gen, clone_without},
@@ -25,16 +24,13 @@ use crate::{
 #[allow(clippy::type_complexity)]
 pub async fn sign<N: Network>(
   network: N,
+  session: Session,
   mut keys_txs: HashMap<
     Participant,
     (ThresholdKeys<N::Curve>, (N::SignableTransaction, N::Eventuality)),
   >,
 ) -> <N::Transaction as Transaction<N>>::Id {
-  let actual_id = SignId {
-    key: keys_txs[&Participant::new(1).unwrap()].0.group_key().to_bytes().as_ref().to_vec(),
-    id: [0xaa; 32],
-    attempt: 0,
-  };
+  let actual_id = SignId { session, id: [0xaa; 32], attempt: 0 };
 
   let mut keys = HashMap::new();
   let mut txs = HashMap::new();
@@ -197,7 +193,7 @@ pub async fn test_signer<N: Network>(network: N) {
 
   // The signer may not publish the TX if it has a connection error
   // It doesn't fail in this case
-  let txid = sign(network.clone(), keys_txs).await;
+  let txid = sign(network.clone(), Session(0), keys_txs).await;
   let tx = network.get_transaction(&txid).await.unwrap();
   assert_eq!(tx.id(), txid);
   // Mine a block, and scan it, to ensure that the TX actually made it on chain
diff --git a/processor/src/tests/wallet.rs b/processor/src/tests/wallet.rs
index c0248c7e2..74a6dd54f 100644
--- a/processor/src/tests/wallet.rs
+++ b/processor/src/tests/wallet.rs
@@ -8,7 +8,10 @@ use tokio::time::timeout;
 
 use serai_db::{DbTxn, Db, MemDb};
 
-use serai_client::primitives::{NetworkId, Coin, Amount, Balance};
+use serai_client::{
+  primitives::{NetworkId, Coin, Amount, Balance},
+  validator_sets::primitives::Session,
+};
 
 use crate::{
   Payment, Plan,
@@ -140,7 +143,7 @@ pub async fn test_wallet<N: Network>(network: N) {
     keys_txs.insert(i, (keys, (signable, eventuality)));
   }
 
-  let txid = sign(network.clone(), keys_txs).await;
+  let txid = sign(network.clone(), Session(0), keys_txs).await;
   let tx = network.get_transaction(&txid).await.unwrap();
   network.mine_block().await;
   let block_number = network.get_latest_block_number().await.unwrap();
diff --git a/tests/coordinator/src/tests/sign.rs b/tests/coordinator/src/tests/sign.rs
index fa974a293..ce4b09189 100644
--- a/tests/coordinator/src/tests/sign.rs
+++ b/tests/coordinator/src/tests/sign.rs
@@ -4,10 +4,9 @@ use std::{
   collections::{HashSet, HashMap},
 };
 
-use zeroize::Zeroizing;
 use rand_core::{RngCore, OsRng};
 
-use ciphersuite::{group::GroupEncoding, Ciphersuite, Secp256k1};
+use ciphersuite::Secp256k1;
 
 use dkg::Participant;
 
@@ -29,18 +28,13 @@ use messages::{coordinator::PlanMeta, sign::SignId, SubstrateContext, Coordinato
 
 use crate::tests::*;
 
-pub async fn sign<C: Ciphersuite>(
+pub async fn sign(
   processors: &mut [Processor],
   processor_is: &[u8],
   session: Session,
-  network_key: &Zeroizing<C::F>,
   plan_id: [u8; 32],
 ) {
-  let id = SignId {
-    key: (C::generator() * **network_key).to_bytes().as_ref().to_vec(),
-    id: plan_id,
-    attempt: 0,
-  };
+  let id = SignId { session, id: plan_id, attempt: 0 };
 
   // Select a random participant to exclude, so we know for sure who *is* participating
   assert_eq!(COORDINATORS - THRESHOLD, 1);
@@ -165,7 +159,7 @@ pub async fn sign<C: Ciphersuite>(
     assert_eq!(
       processor.recv_message().await,
       CoordinatorMessage::Sign(messages::sign::CoordinatorMessage::Completed {
-        key: id.key.clone(),
+        session,
         id: id.id,
         tx: b"signed_tx".to_vec()
       })
@@ -198,8 +192,7 @@ async fn sign_test() {
       }
       let mut processors = new_processors;
 
-      let (participant_is, substrate_key, network_key) =
-        key_gen::<Secp256k1>(&mut processors).await;
+      let (participant_is, substrate_key, _) = key_gen::<Secp256k1>(&mut processors).await;
 
       // 'Send' external coins into Serai
       let serai = processors[0].serai().await;
@@ -346,16 +339,13 @@ async fn sign_test() {
           .send_message(messages::ProcessorMessage::Coordinator(
             messages::coordinator::ProcessorMessage::SubstrateBlockAck {
               block: last_serai_block.number(),
-              plans: vec![PlanMeta {
-                key: (Secp256k1::generator() * *network_key).to_bytes().to_vec(),
-                id: plan_id,
-              }],
+              plans: vec![PlanMeta { session: Session(0), id: plan_id }],
             },
           ))
           .await;
       }
 
-      sign::<Secp256k1>(&mut processors, &participant_is, Session(0), &network_key, plan_id).await;
+      sign(&mut processors, &participant_is, Session(0), plan_id).await;
     })
     .await;
 }
diff --git a/tests/processor/src/tests/send.rs b/tests/processor/src/tests/send.rs
index d6e3ad9ce..986671c16 100644
--- a/tests/processor/src/tests/send.rs
+++ b/tests/processor/src/tests/send.rs
@@ -19,7 +19,7 @@ use crate::{*, tests::*};
 #[allow(unused)]
 pub(crate) async fn recv_sign_preprocesses(
   coordinators: &mut [Coordinator],
-  key: Vec<u8>,
+  session: Session,
   attempt: u32,
 ) -> (SignId, HashMap<Participant, Vec<u8>>) {
   let mut id = None;
@@ -34,7 +34,7 @@ pub(crate) async fn recv_sign_preprocesses(
         preprocesses: mut these_preprocesses,
       }) => {
         if id.is_none() {
-          assert_eq!(&this_id.key, &key);
+          assert_eq!(&this_id.session, &session);
           assert_eq!(this_id.attempt, attempt);
           id = Some(this_id.clone());
         }
@@ -238,8 +238,8 @@ fn send_test() {
 
       // Start signing the TX
       let (mut id, mut preprocesses) =
-        recv_sign_preprocesses(&mut coordinators, key_pair.1.to_vec(), 0).await;
-      assert_eq!(id, SignId { key: key_pair.1.to_vec(), id: plans[0].id, attempt: 0 });
+        recv_sign_preprocesses(&mut coordinators, Session(0), 0).await;
+      assert_eq!(id, SignId { session: Session(0), id: plans[0].id, attempt: 0 });
 
       // Trigger a random amount of re-attempts
       for attempt in 1 ..= u32::try_from(OsRng.next_u64() % 4).unwrap() {
@@ -251,8 +251,7 @@ fn send_test() {
             .send_message(messages::sign::CoordinatorMessage::Reattempt { id: id.clone() })
             .await;
         }
-        (id, preprocesses) =
-          recv_sign_preprocesses(&mut coordinators, key_pair.1.to_vec(), attempt).await;
+        (id, preprocesses) = recv_sign_preprocesses(&mut coordinators, Session(0), attempt).await;
       }
       let participating = preprocesses.keys().cloned().collect::<Vec<_>>();
 
@@ -276,7 +275,7 @@ fn send_test() {
           // Tell them of it as a completion of the relevant signing nodess
           coordinator
             .send_message(messages::sign::CoordinatorMessage::Completed {
-              key: key_pair.1.to_vec(),
+              session: Session(0),
               id: id.id,
               tx: tx_id.clone(),
             })