fix: run DKG only once (#1115)

* Use the x-only coordinates when searching for DKG shares
* Always load the correct DKG shares for each sighash
* Prevent DKG from running multiple times

signer/src/error.rs

@@ -387,7 +387,7 @@ pub enum Error {
 
     /// Missing dkg shares
     #[error("missing dkg shares for the given aggregate key: {0}")]
-    MissingDkgShares(crate::keys::PublicKey),
+    MissingDkgShares(crate::keys::PublicKeyXOnly),
 
     /// Missing public key
     #[error("missing public key")]
@@ -411,6 +411,11 @@ pub enum Error {
     #[error("DKG has not been run")]
     NoDkgShares,
 
+    /// TODO: We don't want to be able to run DKG more than once. Soon this
+    /// restriction will be lifted.
+    #[error("DKG has already been run, can only run once")]
+    DkgHasAlreadyRun,
+
     /// Too many signer utxos
     #[error("too many signer utxos")]
     TooManySignerUtxos,

signer/src/keys.rs

@@ -290,6 +290,12 @@ impl From<&PublicKey> for PublicKeyXOnly {
     }
 }
 
+impl std::fmt::Display for PublicKeyXOnly {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
 impl PublicKeyXOnly {
     /// Creates a public key directly from a slice.
     pub fn from_slice(data: &[u8]) -> Result<Self, Error> {

signer/src/storage/in_memory.rs

@@ -85,7 +85,7 @@ pub struct Store {
         HashMap<model::BitcoinBlockHash, Vec<model::StacksBlockHash>>,
 
     /// Encrypted DKG shares
-    pub encrypted_dkg_shares: BTreeMap<PublicKey, (OffsetDateTime, model::EncryptedDkgShares)>,
+    pub encrypted_dkg_shares: BTreeMap<PublicKeyXOnly, (OffsetDateTime, model::EncryptedDkgShares)>,
 
     /// Rotate keys transactions
     pub rotate_keys_transactions: HashMap<model::StacksTxId, model::RotateKeysTransaction>,
@@ -574,15 +574,18 @@ impl super::DbRead for SharedStore {
             .contains_key(&block_id.into()))
     }
 
-    async fn get_encrypted_dkg_shares(
+    async fn get_encrypted_dkg_shares<X>(
         &self,
-        aggregate_key: &PublicKey,
-    ) -> Result<Option<model::EncryptedDkgShares>, Error> {
+        aggregate_key: X,
+    ) -> Result<Option<model::EncryptedDkgShares>, Error>
+    where
+        X: Into<PublicKeyXOnly> + Send,
+    {
         Ok(self
             .lock()
             .await
             .encrypted_dkg_shares
-            .get(aggregate_key)
+            .get(&aggregate_key.into())
             .map(|(_, shares)| shares.clone()))
     }
 
@@ -1109,7 +1112,7 @@ impl super::DbWrite for SharedStore {
         shares: &model::EncryptedDkgShares,
     ) -> Result<(), Error> {
         self.lock().await.encrypted_dkg_shares.insert(
-            shares.aggregate_key,
+            shares.aggregate_key.into(),
             (time::OffsetDateTime::now_utc(), shares.clone()),
         );
 

signer/src/storage/mod.rs

@@ -22,6 +22,7 @@ use crate::bitcoin::validation::DepositRequestReport;
 use crate::bitcoin::validation::WithdrawalRequestReport;
 use crate::error::Error;
 use crate::keys::PublicKey;
+use crate::keys::PublicKeyXOnly;
 use crate::stacks::events::CompletedDepositEvent;
 use crate::stacks::events::WithdrawalAcceptEvent;
 use crate::stacks::events::WithdrawalCreateEvent;
@@ -204,10 +205,12 @@ pub trait DbRead {
 
     /// Return the applicable DKG shares for the
     /// given aggregate key
-    fn get_encrypted_dkg_shares(
+    fn get_encrypted_dkg_shares<X>(
         &self,
-        aggregate_key: &PublicKey,
-    ) -> impl Future<Output = Result<Option<model::EncryptedDkgShares>, Error>> + Send;
+        aggregate_key: X,
+    ) -> impl Future<Output = Result<Option<model::EncryptedDkgShares>, Error>> + Send
+    where
+        X: Into<PublicKeyXOnly> + Send;
 
     /// Return the most recent DKG shares, and return None if the table is
     /// empty.

signer/src/storage/postgres.rs

@@ -1452,10 +1452,18 @@ impl super::DbRead for PgStore {
         .map_err(Error::SqlxQuery)
     }
 
-    async fn get_encrypted_dkg_shares(
+    async fn get_encrypted_dkg_shares<X>(
         &self,
-        aggregate_key: &PublicKey,
-    ) -> Result<Option<model::EncryptedDkgShares>, Error> {
+        aggregate_key: X,
+    ) -> Result<Option<model::EncryptedDkgShares>, Error>
+    where
+        X: Into<PublicKeyXOnly> + Send,
+    {
+        // The aggregate_key column stores compressed public keys, which
+        // always include a parity byte. Since the input here is an x-only
+        // public key we don't have a parity byte, so we lop it off when
+        // filtering.
+        let key: PublicKeyXOnly = aggregate_key.into();
         sqlx::query_as::<_, model::EncryptedDkgShares>(
             r#"
             SELECT
@@ -1467,10 +1475,10 @@ impl super::DbRead for PgStore {
               , signer_set_public_keys
               , signature_share_threshold
             FROM sbtc_signer.dkg_shares
-            WHERE aggregate_key = $1;
+            WHERE substring(aggregate_key FROM 2) = $1;
             "#,
         )
-        .bind(aggregate_key)
+        .bind(key)
         .fetch_optional(&self.0)
         .await
         .map_err(Error::SqlxQuery)
@@ -1637,7 +1645,7 @@ impl super::DbRead for PgStore {
         // If we've swept funds before, then will have a signer output and
         // a minimum UTXO height, so let's try that first.
         let Some(min_block_height) = self.minimum_utxo_height().await? else {
-            // If the above functon returns None then we know that there
+            // If the above function returns None then we know that there
             // have been no confirmed sweep transactions thus far, so let's
             // try looking for a donation UTXO.
             return self.get_donation_utxo(chain_tip).await;
@@ -1646,7 +1654,7 @@ impl super::DbRead for PgStore {
         // transaction. Let's look for the UTXO in a block after our
         // min_block_height. Note that `Self::get_utxo` returns `None` only
         // when a reorg has affected all sweep transactions. If this
-        // happens we try seaching for a donation.
+        // happens we try searching for a donation.
         let output_type = model::TxOutputType::SignersOutput;
         let fut = self.get_utxo(chain_tip, output_type, min_block_height);
         match fut.await? {
@@ -1850,7 +1858,7 @@ impl super::DbRead for PgStore {
         _chain_tip: &model::BitcoinBlockHash,
         _context_window: u16,
     ) -> Result<Vec<model::SweptWithdrawalRequest>, Error> {
-        // TODO: This can use a similiar query to
+        // TODO: This can use a similar query to
         // `get_swept_deposit_requests()`, but using withdrawal tables instead
         // of deposit.
         unimplemented!()
@@ -2829,9 +2837,9 @@ mod tests {
         // address in the transaction, then we will consider it a relevant
         // transaction. Now what if someone tries to pull a fast one by
         // deploying their own modified version of the sBTC smart contracts
-        // and creating contract calls against that? We'll the address of
+        // and creating contract calls against that? Well the address of
         // these contract calls won't match the ones that we are interested
-        // in and we will filter them out. We test that now,
+        // in, and we will filter them out. We test that now,
         let contract_call = TransactionContractCall {
             // This is the address of the poser that deployed their own
             // versions of the sBTC smart contracts.

signer/src/transaction_coordinator.rs

@@ -561,13 +561,8 @@ where
         .await?;
 
         for mut transaction in transaction_package {
-            self.sign_and_broadcast(
-                bitcoin_chain_tip,
-                aggregate_key,
-                signer_public_keys,
-                &mut transaction,
-            )
-            .await?;
+            self.sign_and_broadcast(bitcoin_chain_tip, signer_public_keys, &mut transaction)
+                .await?;
 
             // TODO: if this (considering also fallback clients) fails, we will
             // need to handle the inconsistency of having the sweep tx confirmed
@@ -909,13 +904,12 @@ where
     async fn sign_and_broadcast(
         &mut self,
         bitcoin_chain_tip: &model::BitcoinBlockHash,
-        aggregate_key: &PublicKey,
         signer_public_keys: &BTreeSet<PublicKey>,
         transaction: &mut utxo::UnsignedTransaction<'_>,
     ) -> Result<(), Error> {
         let mut coordinator_state_machine = CoordinatorStateMachine::load(
             &mut self.context.get_storage_mut(),
-            *aggregate_key,
+            transaction.signer_utxo.public_key,
             signer_public_keys.clone(),
             self.threshold,
             self.private_key,
@@ -959,7 +953,7 @@ where
 
             let mut coordinator_state_machine = CoordinatorStateMachine::load(
                 &mut self.context.get_storage_mut(),
-                *aggregate_key,
+                deposit.signers_public_key,
                 signer_public_keys.clone(),
                 self.threshold,
                 self.private_key,

signer/src/transaction_signer.rs

@@ -232,16 +232,6 @@ where
                 .await?;
             }
 
-            (
-                message::Payload::BitcoinTransactionSignRequest(request),
-                true,
-                ChainTipStatus::Canonical,
-            ) => {
-                tracing::debug!("handling bitcoin transaction sign request");
-                self.handle_bitcoin_transaction_sign_request(request, &msg.bitcoin_chain_tip)
-                    .await?;
-            }
-
             (message::Payload::WstsMessage(wsts_msg), _, _) => {
                 self.handle_wsts_message(
                     wsts_msg,
@@ -282,6 +272,7 @@ where
             }
             // Message types ignored by the transaction signer
             (message::Payload::StacksTransactionSignature(_), _, _)
+            | (message::Payload::BitcoinTransactionSignRequest(_), _, _)
             | (message::Payload::BitcoinTransactionSignAck(_), _, _)
             | (message::Payload::SignerDepositDecision(_), _, _)
             | (message::Payload::SignerWithdrawalDecision(_), _, _) => (),
@@ -390,63 +381,6 @@ where
         Ok(())
     }
 
-    #[tracing::instrument(skip_all)]
-    async fn handle_bitcoin_transaction_sign_request(
-        &mut self,
-        request: &message::BitcoinTransactionSignRequest,
-        bitcoin_chain_tip: &model::BitcoinBlockHash,
-    ) -> Result<(), Error> {
-        let is_valid_sign_request = self
-            .is_valid_bitcoin_transaction_sign_request(request)
-            .await?;
-
-        if is_valid_sign_request {
-            let new_state_machine = SignerStateMachine::load(
-                &self.context.get_storage_mut(),
-                request.aggregate_key,
-                self.threshold,
-                self.signer_private_key,
-            )
-            .await?;
-
-            let txid = request.tx.compute_txid();
-
-            self.wsts_state_machines.insert(txid, new_state_machine);
-
-            let msg = message::BitcoinTransactionSignAck {
-                txid: request.tx.compute_txid(),
-            };
-
-            self.send_message(msg, bitcoin_chain_tip).await?;
-        } else {
-            tracing::warn!("received invalid sign request");
-        }
-
-        Ok(())
-    }
-
-    async fn is_valid_bitcoin_transaction_sign_request(
-        &self,
-        _request: &message::BitcoinTransactionSignRequest,
-    ) -> Result<bool, Error> {
-        let signer_pub_key = self.signer_public_key();
-        let _accepted_deposit_requests = self
-            .context
-            .get_storage()
-            .get_accepted_deposit_requests(&signer_pub_key)
-            .await?;
-
-        // TODO(286): Validate transaction
-        // - Ensure all inputs are either accepted deposit requests
-        //    or directly spendable by the signers.
-        // - Ensure all outputs are either accepted withdrawals
-        //    or pays to an approved signer set.
-        // - Ensure the transaction fee is lower than the minimum
-        //    `max_fee` of any request.
-
-        Ok(true)
-    }
-
     #[tracing::instrument(skip_all)]
     async fn handle_stacks_transaction_sign_request(
         &mut self,
@@ -506,7 +440,7 @@ where
         let public_key = self.signer_public_key();
 
         let Some(shares) = db.get_encrypted_dkg_shares(&request.aggregate_key).await? else {
-            return Err(Error::MissingDkgShares(request.aggregate_key));
+            return Err(Error::MissingDkgShares(request.aggregate_key.into()));
         };
         // There is one check that applies to all Stacks transactions, and
         // that check is that the current signer is in the signing set
@@ -568,6 +502,17 @@ where
                     return Ok(());
                 }
 
+                let dkg_shares = self
+                    .context
+                    .get_storage()
+                    .get_latest_encrypted_dkg_shares()
+                    .await?;
+
+                if dkg_shares.is_some() {
+                    tracing::warn!("we do not support running DKG more than once");
+                    return Err(Error::DkgHasAlreadyRun);
+                }
+
                 let signer_public_keys = self.get_signer_public_keys(bitcoin_chain_tip).await?;
 
                 let state_machine = SignerStateMachine::new(

signer/src/wsts_state_machine.rs

@@ -8,6 +8,7 @@ use crate::error;
 use crate::error::Error;
 use crate::keys::PrivateKey;
 use crate::keys::PublicKey;
+use crate::keys::PublicKeyXOnly;
 use crate::keys::SignerScriptPubKey as _;
 use crate::storage;
 use crate::storage::model;
@@ -91,7 +92,7 @@ impl SignerStateMachine {
         let encrypted_shares = storage
             .get_encrypted_dkg_shares(&aggregate_key)
             .await?
-            .ok_or(error::Error::MissingDkgShares(aggregate_key))?;
+            .ok_or(error::Error::MissingDkgShares(aggregate_key.into()))?;
 
         let decrypted = wsts::util::decrypt(
             &signer_private_key.to_bytes(),
@@ -232,21 +233,23 @@ impl CoordinatorStateMachine {
     /// where you can either start a signing round or start DKG. This
     /// function is for loading the state with the assumption that DKG has
     /// already been successfully completed.
-    pub async fn load<I, S>(
+    pub async fn load<I, S, X>(
         storage: &mut S,
-        aggregate_key: PublicKey,
+        aggregate_key: X,
         signers: I,
         threshold: u16,
         message_private_key: PrivateKey,
     ) -> Result<Self, Error>
     where
         I: IntoIterator<Item = PublicKey>,
         S: storage::DbRead + storage::DbWrite,
+        X: Into<PublicKeyXOnly>,
     {
+        let x_only_aggregate_key: PublicKeyXOnly = aggregate_key.into();
         let encrypted_shares = storage
-            .get_encrypted_dkg_shares(&aggregate_key)
+            .get_encrypted_dkg_shares(x_only_aggregate_key)
             .await?
-            .ok_or(Error::MissingDkgShares(aggregate_key))?;
+            .ok_or(Error::MissingDkgShares(x_only_aggregate_key))?;
 
         let public_dkg_shares: BTreeMap<u32, wsts::net::DkgPublicShares> =
             BTreeMap::decode(encrypted_shares.public_shares.as_slice())?;
@@ -257,8 +260,9 @@ impl CoordinatorStateMachine {
 
         let mut coordinator = Self::new(signers, threshold, message_private_key);
 
+        let aggregate_key = encrypted_shares.aggregate_key.into();
         coordinator
-            .set_key_and_party_polynomials(aggregate_key.into(), party_polynomials)
+            .set_key_and_party_polynomials(aggregate_key, party_polynomials)
             .map_err(Error::wsts_coordinator)?;
         coordinator.current_dkg_id = 1;
 

signer/tests/integration/stacks_events_observer.rs

@@ -109,7 +109,6 @@ async fn test_new_blocks_sends_update_deposits_to_emily() {
         .expect("Wiping Emily database in test setup failed.");
 
     let body = COMPLETED_DEPOSIT_WEBHOOK.to_string();
-    let new_block_event = serde_json::from_str::<NewBlockEvent>(&body).unwrap();
     let deposit_completed_event = get_registry_event_from_webhook(&body, |event| match event {
         RegistryEvent::CompletedDeposit(event) => Some(event),
         _ => panic!("Expected CompletedDeposit event"),