Return the aggregate key in the will sign query

signer/src/storage/in_memory.rs

@@ -896,13 +896,13 @@ impl super::DbRead for SharedStore {
     async fn will_sign_bitcoin_tx_sighash(
         &self,
         sighash: &model::SigHash,
-    ) -> Result<Option<bool>, Error> {
+    ) -> Result<Option<(bool, PublicKeyXOnly)>, Error> {
         Ok(self
             .lock()
             .await
             .bitcoin_sighashes
             .get(sighash)
-            .map(|s| s.will_sign))
+            .map(|s| (s.will_sign, s.aggregate_key)))
     }
 }
 

signer/src/storage/mod.rs

@@ -338,7 +338,7 @@ pub trait DbRead {
     fn will_sign_bitcoin_tx_sighash(
         &self,
         sighash: &model::SigHash,
-    ) -> impl Future<Output = Result<Option<bool>, Error>> + Send;
+    ) -> impl Future<Output = Result<Option<(bool, PublicKeyXOnly)>, Error>> + Send;
 }
 
 /// Represents the ability to write data to the signer storage.

signer/src/storage/postgres.rs

@@ -1896,10 +1896,12 @@ impl super::DbRead for PgStore {
     async fn will_sign_bitcoin_tx_sighash(
         &self,
         sighash: &model::SigHash,
-    ) -> Result<Option<bool>, Error> {
-        sqlx::query_scalar::<_, bool>(
+    ) -> Result<Option<(bool, PublicKeyXOnly)>, Error> {
+        sqlx::query_as::<_, (bool, PublicKeyXOnly)>(
             r#"
-            SELECT will_sign
+            SELECT
+                will_sign
+              , x_only_public_key
             FROM sbtc_signer.bitcoin_tx_sighashes
             WHERE sighash = $1
             "#,

signer/src/transaction_signer.rs

@@ -21,6 +21,7 @@ use crate::ecdsa::SignEcdsa as _;
 use crate::error::Error;
 use crate::keys::PrivateKey;
 use crate::keys::PublicKey;
+use crate::keys::PublicKeyXOnly;
 use crate::message;
 use crate::message::BitcoinPreSignAck;
 use crate::message::StacksTransactionSignRequest;
@@ -608,10 +609,10 @@ where
 
                 let db = self.context.get_storage();
                 let sig_hash = &request.message;
-                let validation_outcome = Self::validate_bitcoin_sign_request(&db, sig_hash).await;
+                let locking_aggregate_key = Self::validate_bitcoin_sign_request(&db, sig_hash).await;
 
-                let validation_status = match &validation_outcome {
-                    Ok(()) => "success",
+                let validation_status = match &locking_aggregate_key {
+                    Ok(_) => "success",
                     Err(Error::SigHashConversion(_)) => "improper-sighash",
                     Err(Error::UnknownSigHash(_)) => "unknown-sighash",
                     Err(Error::InvalidSigHash(_)) => "invalid-sighash",
@@ -627,13 +628,9 @@ where
                 .increment(1);
 
                 if !self.wsts_state_machines.contains_key(&msg.txid) {
-                    let (maybe_aggregate_key, _) = self
-                        .get_signer_set_and_aggregate_key(bitcoin_chain_tip)
-                        .await?;
-
                     let state_machine = SignerStateMachine::load(
                         &db,
-                        maybe_aggregate_key.ok_or(Error::NoDkgShares)?,
+                        locking_aggregate_key?,
                         self.threshold,
                         self.signer_private_key,
                     )
@@ -721,17 +718,17 @@ where
 
     /// Check whether we will sign the message, which is supposed to be a
     /// bitcoin sighash
-    async fn validate_bitcoin_sign_request<D>(db: &D, message: &[u8]) -> Result<(), Error>
+    async fn validate_bitcoin_sign_request<D>(db: &D, msg: &[u8]) -> Result<PublicKeyXOnly, Error>
     where
         D: DbRead,
     {
-        let sighash = TapSighash::from_slice(message)
+        let sighash = TapSighash::from_slice(msg)
             .map_err(Error::SigHashConversion)?
             .into();
 
         match db.will_sign_bitcoin_tx_sighash(&sighash).await? {
-            Some(true) => Ok(()),
-            Some(false) => Err(Error::InvalidSigHash(sighash)),
+            Some((true, public_key)) => Ok(public_key),
+            Some((false, _)) => Err(Error::InvalidSigHash(sighash)),
             None => Err(Error::UnknownSigHash(sighash)),
         }
     }

signer/src/wsts_state_machine.rs

@@ -82,15 +82,15 @@ impl SignerStateMachine {
     /// Create a state machine from loaded DKG shares for the given aggregate key
     pub async fn load<S>(
         storage: &S,
-        aggregate_key: PublicKey,
+        aggregate_key: PublicKeyXOnly,
         threshold: u32,
         signer_private_key: PrivateKey,
     ) -> Result<Self, error::Error>
     where
         S: storage::DbRead,
     {
         let encrypted_shares = storage
-            .get_encrypted_dkg_shares(&aggregate_key)
+            .get_encrypted_dkg_shares(aggregate_key)
             .await?
             .ok_or(error::Error::MissingDkgShares(aggregate_key.into()))?;
 

signer/tests/integration/postgres.rs

@@ -2919,7 +2919,7 @@ async fn can_write_and_get_multiple_bitcoin_txs_sighashes() {
     let results = join_all(withdrawal_outputs_futures).await;
 
     for (output, result) in sighashes.iter().zip(results) {
-        let result = result.unwrap().unwrap();
+        let (result, _) = result.unwrap().unwrap();
         assert_eq!(result, output.will_sign);
     }
     signer::testing::storage::drop_db(db).await;

signer/tests/integration/transaction_signer.rs

@@ -319,14 +319,14 @@ pub async fn assert_should_be_able_to_handle_sbtc_requests() {
 
     // Check that the intentions to sign the requests sighashes
     // are stored in the database
-    let will_sign = db
+    let (will_sign, _) = db
         .will_sign_bitcoin_tx_sighash(&signer_digest.sighash.into())
         .await
         .expect("query to check if signer sighash is stored failed")
         .expect("signer sighash not stored");
 
     assert!(will_sign);
-    let will_sign = db
+    let (will_sign, _) = db
         .will_sign_bitcoin_tx_sighash(&deposit_digest.sighash.into())
         .await
         .expect("query to check if deposit sighash is stored failed")