feat: stacks signer able to save multiple dkg shares and load it where appropriate

.github/workflows/bitcoin-tests.yml

@@ -84,6 +84,7 @@ jobs:
           - tests::nakamoto_integrations::forked_tenure_is_ignored
           - tests::signer::stackerdb_dkg
           - tests::signer::stackerdb_sign_request_rejected
+          - tests::signer::stackerdb_recover_old_dkg_key
           - tests::signer::stackerdb_block_proposal
           - tests::signer::stackerdb_filter_bad_transactions
           - tests::signer::stackerdb_mine_2_nakamoto_reward_cycles

stacks-signer/src/runloop.rs

@@ -22,6 +22,7 @@ use blockstack_lib::chainstate::stacks::boot::SIGNERS_NAME;
 use blockstack_lib::util_lib::boot::boot_code_id;
 use hashbrown::HashMap;
 use libsigner::{SignerEntries, SignerEvent, SignerRunLoop};
+use rand_core::OsRng;
 use slog::{slog_debug, slog_error, slog_info, slog_warn};
 use stacks_common::types::chainstate::StacksAddress;
 use stacks_common::{debug, error, info, warn};
@@ -213,7 +214,7 @@ impl RunLoop {
     }
 
     /// Refresh signer configuration for a specific reward cycle
-    fn refresh_signer_config(&mut self, reward_cycle: u64) {
+    fn refresh_signer_config(&mut self, reward_cycle: u64) -> Result<(), ClientError> {
         let reward_index = reward_cycle % 2;
         if let Some(new_signer_config) = self.get_signer_config(reward_cycle) {
             let signer_id = new_signer_config.signer_id;
@@ -235,12 +236,33 @@ impl RunLoop {
                     }
                 }
             }
-            let new_signer = Signer::from(new_signer_config);
+            let mut new_signer = Signer::from(new_signer_config);
+
+            let dkg_id = retry_with_exponential_backoff(|| {
+                self.stacks_client
+                    .get_last_round(reward_cycle)
+                    .map_err(backoff::Error::transient)
+            })?
+            .unwrap_or(0);
+
+            let approved_aggregate_key = retry_with_exponential_backoff(|| {
+                self.stacks_client
+                    .get_approved_aggregate_key(reward_cycle)
+                    .map_err(backoff::Error::transient)
+            })?;
+
+            new_signer.state_machine.reset(dkg_id, &mut OsRng);
+            new_signer.approved_aggregate_public_key = approved_aggregate_key;
+            new_signer
+                .load_saved_state()
+                .expect("Failed to load signer state");
             info!("{new_signer} initialized.");
             self.stacks_signers.insert(reward_index, new_signer);
         } else {
             warn!("Signer is not registered for reward cycle {reward_cycle}. Waiting for confirmed registration...");
         }
+
+        Ok(())
     }
 
     fn initialize_runloop(&mut self) -> Result<(), ClientError> {
@@ -251,10 +273,10 @@ impl RunLoop {
                 .map_err(backoff::Error::transient)
         })?;
         let current_reward_cycle = reward_cycle_info.reward_cycle;
-        self.refresh_signer_config(current_reward_cycle);
+        self.refresh_signer_config(current_reward_cycle)?;
         // We should only attempt to initialize the next reward cycle signer if we are in the prepare phase of the next reward cycle
         if reward_cycle_info.is_in_prepare_phase(reward_cycle_info.last_burnchain_block_height) {
-            self.refresh_signer_config(current_reward_cycle.saturating_add(1));
+            self.refresh_signer_config(current_reward_cycle.saturating_add(1))?;
         }
         self.current_reward_cycle_info = Some(reward_cycle_info);
         if self.stacks_signers.is_empty() {
@@ -290,7 +312,7 @@ impl RunLoop {
                 .unwrap_or(true)
             {
                 info!("Received a new burnchain block height ({current_burn_block_height}) in the prepare phase of the next reward cycle ({next_reward_cycle}). Checking for signer registration...");
-                self.refresh_signer_config(next_reward_cycle);
+                self.refresh_signer_config(next_reward_cycle)?;
             }
         }
         self.cleanup_stale_signers(current_reward_cycle);

stacks-signer/src/signer.rs

@@ -49,14 +49,20 @@ use wsts::state_machine::coordinator::{
 };
 use wsts::state_machine::signer::Signer as SignerStateMachine;
 use wsts::state_machine::{OperationResult, SignError};
-use wsts::traits::Signer as _;
+use wsts::traits::{Signer as _, SignerState};
 use wsts::v2;
 
 use crate::client::{ClientError, StackerDB, StacksClient};
 use crate::config::SignerConfig;
 use crate::coordinator::CoordinatorSelector;
 use crate::signerdb::SignerDb;
 
+/// The number of previous DKG shares persisted
+const NUM_STORED_DKG_SHARES: usize = 2;
+
+/// The persisted signer states
+type StoredSignerStates = VecDeque<SignerState>;
+
 /// The signer StackerDB slot ID, purposefully wrapped to prevent conflation with SignerID
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Copy, PartialOrd, Ord)]
 pub struct SignerSlotID(pub u32);
@@ -169,6 +175,8 @@ pub struct Signer {
     pub mainnet: bool,
     /// The signer id
     pub signer_id: u32,
+    /// The signer slot id for the current signer
+    pub signer_slot_id: SignerSlotID,
     /// The signer slot ids for the signers in the reward cycle
     pub signer_slot_ids: Vec<SignerSlotID>,
     /// The addresses of other signers
@@ -277,7 +285,7 @@ impl Signer {
 
 impl From<SignerConfig> for Signer {
     fn from(signer_config: SignerConfig) -> Self {
-        let mut stackerdb = StackerDB::from(&signer_config);
+        let stackerdb = StackerDB::from(&signer_config);
 
         let num_signers = signer_config
             .signer_entries
@@ -324,7 +332,7 @@ impl From<SignerConfig> for Signer {
         let signer_db =
             SignerDb::new(&signer_config.db_path).expect("Failed to connect to signer Db");
 
-        let mut state_machine = SignerStateMachine::new(
+        let state_machine = SignerStateMachine::new(
             threshold,
             num_signers,
             num_keys,
@@ -334,20 +342,6 @@ impl From<SignerConfig> for Signer {
             signer_config.signer_entries.public_keys,
         );
 
-        if let Some(state) = load_encrypted_signer_state(
-            &mut stackerdb,
-            signer_config.signer_slot_id,
-            &state_machine.network_private_key,
-        ).or_else(|err| {
-                warn!("Failed to load encrypted signer state from StackerDB, falling back to SignerDB: {err}");
-                load_encrypted_signer_state(
-                    &signer_db,
-                    signer_config.reward_cycle,
-                    &state_machine.network_private_key)
-            }).expect("Failed to load encrypted signer state from both StackerDB and SignerDB") {
-            state_machine.signer = state;
-        };
-
         Self {
             coordinator,
             state_machine,
@@ -361,6 +355,7 @@ impl From<SignerConfig> for Signer {
                 .signer_ids
                 .into_keys()
                 .collect(),
+            signer_slot_id: signer_config.signer_slot_id,
             signer_slot_ids: signer_config.signer_slot_ids.clone(),
             next_signer_slot_ids: vec![],
             next_signer_addresses: vec![],
@@ -1059,6 +1054,21 @@ impl Signer {
 
     /// Process a dkg result by broadcasting a vote to the stacks node
     fn process_dkg(&mut self, stacks_client: &StacksClient, dkg_public_key: &Point) {
+        // Abort if an aggregate key has already been set for the current rewards cycle
+        if let Some(aggregate_key) = stacks_client
+            .get_approved_aggregate_key(self.reward_cycle)
+            .map_err(|e| error!("{self}: Unable to get approved aggregate key: {e}"))
+            .unwrap_or_default()
+        {
+            self.approved_aggregate_public_key = Some(aggregate_key);
+            match self.load_saved_state() {
+                Err(e) => error!("{self}: Failed to load saved state: {e}"),
+                Ok(()) => (),
+            };
+
+            return;
+        }
+
         let mut dkg_results_bytes = vec![];
         debug!(
             "{self}: Received DKG result. Broadcasting vote to the stacks node...";
@@ -1280,8 +1290,30 @@ impl Signer {
     fn save_signer_state(&mut self) -> Result<(), PersistenceError> {
         let rng = &mut OsRng;
 
-        let state = self.state_machine.signer.save();
-        let serialized_state = serde_json::to_vec(&state)?;
+        let mut saved_states = self.load_encrypted_signer_states().unwrap_or_else(|err| {
+            warn!("{self}: Failed to load previous dkg state: {err}");
+            VecDeque::new()
+        });
+
+        let current_dkg_state = self.state_machine.signer.save();
+        let group_key = current_dkg_state.group_key;
+
+        info!("{self}: Saving state for key {group_key}");
+
+        // Ensure we don't save the same state multiple times
+        if saved_states
+            .iter()
+            .find(|state| state.group_key == group_key)
+            .is_none()
+        {
+            saved_states.push_back(current_dkg_state);
+        }
+
+        if saved_states.len() > NUM_STORED_DKG_SHARES {
+            saved_states.pop_front();
+        }
+
+        let serialized_state = serde_json::to_vec(&saved_states)?;
 
         let encrypted_state = encrypt(
             &self.state_machine.network_private_key,
@@ -1328,6 +1360,74 @@ impl Signer {
         Ok(())
     }
 
+    /// Load the saved signer state for the current dkg round
+    pub fn load_saved_state(&mut self) -> Result<(), PersistenceError> {
+        let Some(aggregate_key) = self.approved_aggregate_public_key else {
+            return Ok(());
+        };
+
+        info!("{self}: Loading saved state for key: {aggregate_key}");
+        if let Some(state) = self.load_saved_state_for_aggregate_key(aggregate_key)? {
+            let party_id = state.party_id;
+            let poly_commitment = state.get_poly_commitment(&mut OsRng);
+
+            let party_polynomials = poly_commitment
+                .as_ref()
+                .map(|poly_commitment| (&party_id, poly_commitment));
+
+            let mut dkg_results_bytes = vec![];
+
+            if let Err(e) = SignerMessage::serialize_dkg_result(
+                &mut dkg_results_bytes,
+                &aggregate_key,
+                party_polynomials.into_iter(),
+            ) {
+                error!(
+                    "{self}: Failed to serialize DKGResults message after loading saved state: {e}"
+                );
+            } else if let Err(e) = self
+                .stackerdb
+                .send_message_bytes_with_retry(&MessageSlotID::DkgResults, dkg_results_bytes)
+            {
+                error!("{self}: Failed to send DKGResults message to StackerDB after loading: {e}");
+            };
+
+            self.state_machine.signer = state;
+        } else {
+            warn!("{self}: Signer unable to load state for key {aggregate_key}");
+        };
+
+        Ok(())
+    }
+
+    /// Load the saved state for a particular aggregate key.
+    fn load_saved_state_for_aggregate_key(
+        &mut self,
+        aggregate_key: Point,
+    ) -> Result<Option<v2::Signer>, PersistenceError> {
+        Ok(
+            get_signer_state(self.load_encrypted_signer_states()?, aggregate_key)
+                .map(|state| v2::Signer::load(&state)),
+        )
+    }
+
+    /// Load the entire encrypted signer state
+    fn load_encrypted_signer_states(&mut self) -> Result<StoredSignerStates, PersistenceError> {
+        let loaded_signers = load_encrypted_signer_state(
+            &mut self.stackerdb,
+            self.signer_slot_id.into(),
+            &self.state_machine.network_private_key,
+        ).or_else(|err| {
+                warn!("Failed to load encrypted signer state from StackerDB, falling back to SignerDB: {err}");
+                load_encrypted_signer_state(
+                    &self.signer_db,
+                    self.reward_cycle,
+                    &self.state_machine.network_private_key)
+        })?;
+
+        Ok(loaded_signers)
+    }
+
     /// Send any operation results across the provided channel
     fn send_operation_results(
         &mut self,
@@ -1403,9 +1503,6 @@ impl Signer {
         self.approved_aggregate_public_key =
             stacks_client.get_approved_aggregate_key(self.reward_cycle)?;
         if self.approved_aggregate_public_key.is_some() {
-            // TODO: this will never work as is. We need to have stored our party shares on the side etc for this particular aggregate key.
-            // Need to update state to store the necessary info, check against it to see if we have participated in the winning round and
-            // then overwrite our value accordingly. Otherwise, we will be locked out of the round and should not participate.
             let internal_dkg = self.coordinator.aggregate_public_key;
             if internal_dkg != self.approved_aggregate_public_key {
                 warn!("{self}: we do not support changing the internal DKG key yet. Expected {internal_dkg:?} got {:?}", self.approved_aggregate_public_key);
@@ -1609,15 +1706,22 @@ fn load_encrypted_signer_state<S: SignerStateStorage>(
     storage: S,
     id: S::IdType,
     private_key: &Scalar,
-) -> Result<Option<v2::Signer>, PersistenceError> {
-    if let Some(encrypted_state) = storage.get_encrypted_signer_state(id)? {
-        let serialized_state = decrypt(private_key, &encrypted_state)?;
-        let state = serde_json::from_slice(&serialized_state)
-            .expect("Failed to deserialize decryoted state");
-        Ok(Some(v2::Signer::load(&state)))
-    } else {
-        Ok(None)
-    }
+) -> Result<StoredSignerStates, PersistenceError> {
+    let Some(encrypted_state) = storage.get_encrypted_signer_state(id)? else {
+        return Ok(VecDeque::new());
+    };
+    let serialized_state = decrypt(private_key, &encrypted_state)?;
+
+    Ok(serde_json::from_slice(&serialized_state)?)
+}
+
+fn get_signer_state(
+    loaded_signers: StoredSignerStates,
+    aggregate_key: Point,
+) -> Option<SignerState> {
+    loaded_signers
+        .into_iter()
+        .find(|state| state.group_key == aggregate_key)
 }
 
 trait SignerStateStorage {

testnet/stacks-node/src/nakamoto_node/sign_coordinator.rs

@@ -133,6 +133,8 @@ fn get_signer_commitments(
         MessageSlotID::DkgResults.stacker_db_contract(is_mainnet, reward_cycle);
     let signer_set_len = u32::try_from(reward_set.len())
         .map_err(|_| ChainstateError::InvalidStacksBlock("Reward set length exceeds u32".into()))?;
+
+    let mut all_party_polynomials = HashMap::new();
     for signer_id in 0..signer_set_len {
         let Some(signer_data) = stackerdbs.get_latest_chunk(&commitment_contract, signer_id)?
         else {
@@ -154,6 +156,10 @@ fn get_signer_commitments(
             continue;
         };
 
+        for (party_id, poly_commitment) in party_polynomials.iter() {
+            all_party_polynomials.insert(*party_id, poly_commitment.clone());
+        }
+
         if &aggregate_key != expected_aggregate_key {
             warn!(
                 "Aggregate key in DKG results does not match expected, will look for results from other signers.";
@@ -177,6 +183,16 @@ fn get_signer_commitments(
 
         return Ok(party_polynomials);
     }
+
+    let computed_key = all_party_polynomials
+        .iter()
+        .fold(Point::default(), |s, (_, comm)| s + comm.poly[0]);
+
+    if &computed_key == expected_aggregate_key {
+        debug!("Aggregate key computed from combined DKG results match expected. Using this one");
+        return Ok(all_party_polynomials.into_iter().collect());
+    }
+
     error!(
         "No valid DKG results found for the active signing set, cannot coordinate a group signature";
         "reward_cycle" => reward_cycle,