refactor: wsts-related cleanup

Cargo.toml

@@ -28,7 +28,7 @@ stackslib = { git = "https://github.com/stacks-network/stacks-core", rev = "4977
 stacks-common = { git = "https://github.com/stacks-network/stacks-core", rev = "49777d3fd73a6dbb610be80c376b7d9389c9871a", default-features = false, features = ["canonical"] }
 
 # Trust Machines Dependencies
-wsts = { git = "https://github.com/Trust-Machines/wsts.git", rev = "53ae23f5f35def420877ccc8c0fe3662e64e38a1" }
+wsts = { git = "https://github.com/Trust-Machines/wsts.git", rev = "a7bf38bd54cddf0b78c0f7c521a5ed1537a684fa" }
 
 # Crates.io
 aquamarine = { version = "0.6.0", default-features = false }

protobufs/stacks/signer/v1/messages.proto

@@ -36,7 +36,7 @@ message SignerMessage {
 // A wsts message.
 message WstsMessage {
   // The transaction ID this message relates to, will be a dummy ID for DKG messages
-  bitcoin.BitcoinTxid txid = 1;
+  bitcoin.BitcoinTxid txid = 1 [deprecated = true];
   // The wsts message
   oneof inner {
     // Tell signers to begin DKG by sending DKG public shares
@@ -60,6 +60,10 @@ message WstsMessage {
     // Tell coordinator signature shares
     crypto.wsts.SignatureShareResponse signature_share_response = 11;
   }
+  oneof id {
+    bitcoin.BitcoinTxid id_bitcoin_txid = 12;
+    crypto.PublicKey id_rotate_key = 13;
+  }
 }
 
 // Wraps an inner type with a public key and a signature,

signer/src/message.rs

@@ -1,14 +1,18 @@
 //! Signer message definition for network communication
 
+use rand::rngs::OsRng;
+use rand::Rng;
 use secp256k1::ecdsa::RecoverableSignature;
 
 use crate::bitcoin::utxo::Fees;
 use crate::bitcoin::validation::TxRequestIds;
+use crate::keys::PrivateKey;
 use crate::keys::PublicKey;
 use crate::stacks::contracts::ContractCall;
 use crate::stacks::contracts::StacksTx;
 use crate::storage::model;
 use crate::storage::model::BitcoinBlockHash;
+use crate::storage::model::BitcoinTxId;
 use crate::storage::model::StacksTxId;
 
 /// Messages exchanged between signers
@@ -224,16 +228,81 @@ pub struct BitcoinPreSignRequest {
 #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
 pub struct BitcoinPreSignAck;
 
+/// The identifier for a WSTS message.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum WstsMessageId {
+    /// The WSTS message is related to a Bitcoin transaction.
+    BitcoinTxid(bitcoin::Txid),
+    /// The WSTS message is related to a rotate key operation.
+    RotateKey(PublicKey),
+}
+
+impl From<bitcoin::Txid> for WstsMessageId {
+    fn from(txid: bitcoin::Txid) -> Self {
+        Self::BitcoinTxid(txid)
+    }
+}
+
+impl From<crate::storage::model::BitcoinTxId> for WstsMessageId {
+    fn from(txid: crate::storage::model::BitcoinTxId) -> Self {
+        Self::BitcoinTxid(txid.into())
+    }
+}
+
+impl WstsMessageId {
+    /// Generate a random [`WstsMessageId::BitcoinTxid`] WSTS message ID
+    pub fn random_bitcoin_txid() -> Self {
+        let random_bytes: [u8; 32] = OsRng.gen();
+        let txid = BitcoinTxId::from(random_bytes).into();
+        Self::BitcoinTxid(txid)
+    }
+
+    /// Generate a random [`WstsMessageId::AggregateKey`] WSTS message ID
+    pub fn random_rotate_key() -> Self {
+        let private = PrivateKey::new(&mut OsRng);
+        let public = PublicKey::from_private_key(&private);
+        Self::RotateKey(public)
+    }
+}
+
+impl std::fmt::Display for WstsMessageId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            WstsMessageId::BitcoinTxid(txid) => write!(f, "bitcoin-txid({})", txid),
+            WstsMessageId::RotateKey(aggregate_key) => {
+                write!(f, "rotate-key({})", aggregate_key)
+            }
+        }
+    }
+}
+
 /// A wsts message.
 #[derive(Debug, Clone, PartialEq)]
 pub struct WstsMessage {
-    /// The transaction ID this message relates to,
-    /// will be a dummy ID for DKG messages
-    pub txid: bitcoin::Txid,
+    /// The id of the wsts message.
+    pub id: WstsMessageId,
     /// The wsts message
     pub inner: wsts::net::Message,
 }
 
+impl WstsMessage {
+    /// Returns the type of the message as a &str.
+    pub fn type_id(&self) -> &'static str {
+        match self.inner {
+            wsts::net::Message::DkgBegin(_) => "dkg-begin",
+            wsts::net::Message::DkgEndBegin(_) => "dkg-end-begin",
+            wsts::net::Message::DkgEnd(_) => "dkg-end",
+            wsts::net::Message::DkgPrivateBegin(_) => "dkg-private-begin",
+            wsts::net::Message::DkgPrivateShares(_) => "dkg-private-shares",
+            wsts::net::Message::DkgPublicShares(_) => "dkg-public-shares",
+            wsts::net::Message::NonceRequest(_) => "nonce-request",
+            wsts::net::Message::NonceResponse(_) => "nonce-response",
+            wsts::net::Message::SignatureShareRequest(_) => "signature-share-request",
+            wsts::net::Message::SignatureShareResponse(_) => "signature-share-response",
+        }
+    }
+}
+
 /// Convenient type aliases
 type StacksBlockHash = [u8; 32];
 

signer/src/proto/convert.rs

@@ -54,6 +54,7 @@ use crate::message::SignerWithdrawalDecision;
 use crate::message::StacksTransactionSignRequest;
 use crate::message::StacksTransactionSignature;
 use crate::message::WstsMessage;
+use crate::message::WstsMessageId;
 use crate::proto;
 use crate::stacks::contracts::AcceptWithdrawalV1;
 use crate::stacks::contracts::CompleteDepositV1;
@@ -69,6 +70,8 @@ use crate::storage::model::StacksBlockHash;
 use crate::storage::model::StacksPrincipal;
 use crate::storage::model::StacksTxId;
 
+use super::wsts_message;
+
 /// This trait is to make it easy to handle fields of protobuf structs that
 /// are `None`, when they should be `Some(_)`.
 trait RequiredField: Sized {
@@ -1058,6 +1061,7 @@ impl TryFrom<proto::SignatureShareResponse> for SignatureShareResponse {
     }
 }
 impl From<WstsMessage> for proto::WstsMessage {
+    #[allow(deprecated)]
     fn from(value: WstsMessage) -> Self {
         let inner = match value.inner {
             wsts::net::Message::DkgBegin(inner) => {
@@ -1089,15 +1093,30 @@ impl From<WstsMessage> for proto::WstsMessage {
                 proto::wsts_message::Inner::SignatureShareResponse(inner.into())
             }
         };
+
         proto::WstsMessage {
-            txid: Some(BitcoinTxId::from(value.txid).into()),
+            txid: match value.id {
+                WstsMessageId::BitcoinTxid(txid) => {
+                    Some(proto::BitcoinTxid::from(BitcoinTxId::from(txid)))
+                }
+                WstsMessageId::RotateKey(_) => None,
+            },
+            id: Some(match value.id {
+                WstsMessageId::BitcoinTxid(txid) => {
+                    wsts_message::Id::IdBitcoinTxid(proto::BitcoinTxid {
+                        txid: Some(proto::Uint256::from(BitcoinTxId::from(txid).into_bytes())),
+                    })
+                }
+                WstsMessageId::RotateKey(pubkey) => wsts_message::Id::IdRotateKey(pubkey.into()),
+            }),
             inner: Some(inner),
         }
     }
 }
 
 impl TryFrom<proto::WstsMessage> for WstsMessage {
     type Error = Error;
+
     fn try_from(value: proto::WstsMessage) -> Result<Self, Self::Error> {
         let inner = match value.inner.required()? {
             proto::wsts_message::Inner::DkgBegin(inner) => {
@@ -1131,8 +1150,22 @@ impl TryFrom<proto::WstsMessage> for WstsMessage {
                 wsts::net::Message::SignatureShareResponse(inner.try_into()?)
             }
         };
+
+        #[allow(deprecated)]
         Ok(WstsMessage {
-            txid: BitcoinTxId::try_from(value.txid.required()?)?.into(),
+            id: match value.id {
+                Some(id) => match id {
+                    wsts_message::Id::IdBitcoinTxid(txid) => {
+                        WstsMessageId::BitcoinTxid(BitcoinTxId::try_from(txid)?.into())
+                    }
+                    wsts_message::Id::IdRotateKey(pubkey) => {
+                        WstsMessageId::RotateKey(PublicKey::try_from(pubkey)?)
+                    }
+                },
+                None => WstsMessageId::BitcoinTxid(
+                    BitcoinTxId::try_from(value.txid.required()?)?.into(),
+                ),
+            },
             inner,
         })
     }

signer/src/proto/generated/stacks.signer.v1.rs

@@ -287,11 +287,14 @@ pub mod signer_message {
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct WstsMessage {
     /// The transaction ID this message relates to, will be a dummy ID for DKG messages
+    #[deprecated]
     #[prost(message, optional, tag = "1")]
     pub txid: ::core::option::Option<super::super::super::bitcoin::BitcoinTxid>,
     /// The wsts message
     #[prost(oneof = "wsts_message::Inner", tags = "2, 3, 4, 5, 6, 7, 8, 9, 10, 11")]
     pub inner: ::core::option::Option<wsts_message::Inner>,
+    #[prost(oneof = "wsts_message::Id", tags = "12, 13")]
+    pub id: ::core::option::Option<wsts_message::Id>,
 }
 /// Nested message and enum types in `WstsMessage`.
 pub mod wsts_message {
@@ -335,6 +338,13 @@ pub mod wsts_message {
             super::super::super::super::crypto::wsts::SignatureShareResponse,
         ),
     }
+    #[derive(Clone, Copy, PartialEq, ::prost::Oneof)]
+    pub enum Id {
+        #[prost(message, tag = "12")]
+        IdBitcoinTxid(super::super::super::super::bitcoin::BitcoinTxid),
+        #[prost(message, tag = "13")]
+        IdRotateKey(super::super::super::super::crypto::PublicKey),
+    }
 }
 /// Wraps an inner type with a public key and a signature,
 /// allowing easy verification of the integrity of the inner data.

signer/src/testing/message.rs

@@ -113,7 +113,7 @@ impl fake::Dummy<fake::Faker> for message::WstsMessage {
         };
 
         Self {
-            txid: dummy::txid(config, rng),
+            id: dummy::txid(config, rng).into(),
             inner: wsts::net::Message::DkgEndBegin(dkg_end_begin),
         }
     }

signer/src/testing/wsts.rs

@@ -20,6 +20,7 @@ use crate::ecdsa::SignEcdsa as _;
 use crate::keys::PrivateKey;
 use crate::keys::PublicKey;
 use crate::message;
+use crate::message::WstsMessageId;
 use crate::network;
 use crate::network::MessageTransfer as _;
 use crate::storage;
@@ -147,9 +148,10 @@ impl Coordinator {
             .start_public_shares()
             .expect("failed to start public shares");
 
-        self.send_packet(bitcoin_chain_tip, txid, outbound).await;
+        self.send_packet(bitcoin_chain_tip, txid.into(), outbound)
+            .await;
 
-        match self.loop_until_result(bitcoin_chain_tip, txid).await {
+        match self.loop_until_result(bitcoin_chain_tip, txid.into()).await {
             wsts::state_machine::OperationResult::Dkg(aggregate_key) => {
                 PublicKey::try_from(&aggregate_key).expect("Got the point at infinity")
             }
@@ -161,7 +163,7 @@ impl Coordinator {
     pub async fn run_signing_round(
         &mut self,
         bitcoin_chain_tip: model::BitcoinBlockHash,
-        txid: bitcoin::Txid,
+        id: WstsMessageId,
         msg: &[u8],
         signature_type: SignatureType,
     ) -> wsts::taproot::SchnorrProof {
@@ -170,9 +172,9 @@ impl Coordinator {
             .start_signing_round(msg, signature_type)
             .expect("failed to start signing round");
 
-        self.send_packet(bitcoin_chain_tip, txid, outbound).await;
+        self.send_packet(bitcoin_chain_tip, id, outbound).await;
 
-        match self.loop_until_result(bitcoin_chain_tip, txid).await {
+        match self.loop_until_result(bitcoin_chain_tip, id).await {
             wsts::state_machine::OperationResult::SignTaproot(signature)
             | wsts::state_machine::OperationResult::SignSchnorr(signature) => signature,
             _ => panic!("unexpected operation result"),
@@ -182,7 +184,7 @@ impl Coordinator {
     async fn loop_until_result(
         &mut self,
         bitcoin_chain_tip: model::BitcoinBlockHash,
-        txid: bitcoin::Txid,
+        id: WstsMessageId,
     ) -> wsts::state_machine::OperationResult {
         let future = async move {
             loop {
@@ -204,7 +206,7 @@ impl Coordinator {
 
                 if let Some(packet) = outbound_packet {
                     tokio::time::sleep(std::time::Duration::from_secs(1)).await;
-                    self.send_packet(bitcoin_chain_tip, txid, packet).await;
+                    self.send_packet(bitcoin_chain_tip, id, packet).await;
                 }
 
                 if let Some(result) = operation_result {
@@ -272,7 +274,7 @@ impl Signer {
                         .process_inbound_messages(&[packet.clone()])
                         .expect("message processing failed");
 
-                    self.send_packet(bitcoin_chain_tip, wsts_msg.txid, packet.clone())
+                    self.send_packet(bitcoin_chain_tip, wsts_msg.id, packet.clone())
                         .await;
 
                     if let wsts::net::Message::DkgEnd(_) = packet.msg {
@@ -312,7 +314,7 @@ impl Signer {
                         .process_inbound_messages(&[packet.clone()])
                         .expect("message processing failed");
 
-                    self.send_packet(bitcoin_chain_tip, wsts_msg.txid, packet.clone())
+                    self.send_packet(bitcoin_chain_tip, wsts_msg.id, packet.clone())
                         .await;
 
                     if let wsts::net::Message::SignatureShareResponse(_) = packet.msg {
@@ -339,10 +341,10 @@ trait WstsEntity {
     async fn send_packet(
         &mut self,
         bitcoin_chain_tip: model::BitcoinBlockHash,
-        txid: bitcoin::Txid,
+        id: WstsMessageId,
         packet: wsts::net::Packet,
     ) {
-        let payload: message::Payload = message::WstsMessage { txid, inner: packet.msg }.into();
+        let payload: message::Payload = message::WstsMessage { id, inner: packet.msg }.into();
 
         let msg = payload
             .to_message(bitcoin_chain_tip)