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,18 @@ message WstsMessage {
     // Tell coordinator signature shares
     crypto.wsts.SignatureShareResponse signature_share_response = 11;
   }
+  oneof id {
+    // If this WSTS message is related to a Bitcoin signing round, this field
+    // will be set to the related Bitcoin transaction ID.
+    bitcoin.BitcoinTxid id_bitcoin_txid = 12;
+    // If this WSTS message is related to a rotate-keys transaction, this field
+    // will be set to the _new_ aggregate public key being verified.
+    crypto.PublicKey id_rotate_key = 13;
+    // If this WSTS message is related to a DKG round, this field will be set
+    // to the 32-byte id determined based on the coordinator public key and
+    // block hash, set by the coordinator.
+    crypto.Uint256 id_dkg = 14;
+  }
 }
 
 // Wraps an inner type with a public key and a signature,

signer/src/message.rs

@@ -224,16 +224,70 @@ 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 signing round.
+    BitcoinTxid(bitcoin::Txid),
+    /// The WSTS message is related to a rotate key verification operation.
+    RotateKey(PublicKey),
+    /// The WSTS message is related to a DKG round.
+    Dkg([u8; 32]),
+}
+
+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 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)
+            }
+            WstsMessageId::Dkg(id) => {
+                write!(f, "dkg({})", hex::encode(id))
+            }
+        }
+    }
+}
+
 /// 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,32 @@ 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::Dkg(_) => None,
+                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()),
+                WstsMessageId::Dkg(id) => wsts_message::Id::IdDkg(id.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 +1152,23 @@ 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)?)
+                    }
+                    wsts_message::Id::IdDkg(id) => WstsMessageId::Dkg(id.into()),
+                },
+                None => WstsMessageId::BitcoinTxid(
+                    BitcoinTxId::try_from(value.txid.required()?)?.into(),
+                ),
+            },
             inner,
         })
     }
@@ -1737,21 +1773,24 @@ mod tests {
         U: From<T>,
         E: std::fmt::Debug,
     {
-        // The type T originates from a signer. Let's create a random
-        // instance of one.
-        let original: T = Faker.fake_with_rng(&mut OsRng);
-        // The type U is a protobuf type. Before sending it to other
-        // signers, we convert our internal type into it's protobuf
-        // counterpart. We can always infallibly create U from T.
-        let proto_original = U::from(original.clone());
-
-        // Some other signer receives an instance of U. This could be a
-        // malicious actor or a modified version of the signer binary
-        // where they made some mistake, so converting back to T can fail.
-        let original_from_proto = T::try_from(proto_original).unwrap();
-        // In this case, we know U was created from T correctly, so we
-        // should be able to convert back without issues.
-        assert_eq!(original, original_from_proto);
+        // TODO: proptest
+        for _ in 0..25 {
+            // The type T originates from a signer. Let's create a random
+            // instance of one.
+            let original: T = Faker.fake_with_rng(&mut OsRng);
+            // The type U is a protobuf type. Before sending it to other
+            // signers, we convert our internal type into it's protobuf
+            // counterpart. We can always infallibly create U from T.
+            let proto_original = U::from(original.clone());
+
+            // Some other signer receives an instance of U. This could be a
+            // malicious actor or a modified version of the signer binary
+            // where they made some mistake, so converting back to T can fail.
+            let original_from_proto = T::try_from(proto_original).unwrap();
+            // In this case, we know U was created from T correctly, so we
+            // should be able to convert back without issues.
+            assert_eq!(original, original_from_proto);
+        }
     }
 
     /// This test is identical to [`convert_protobuf_types`] tests above,
@@ -1792,11 +1831,14 @@ mod tests {
         U: From<T>,
         E: std::fmt::Debug,
     {
-        let original: T = Unit.fake_with_rng(&mut OsRng);
-        let proto_original = U::from(original.clone());
+        // TODO: proptest
+        for _ in 0..10 {
+            let original: T = Unit.fake_with_rng(&mut OsRng);
+            let proto_original = U::from(original.clone());
 
-        let original_from_proto = T::try_from(proto_original).unwrap();
-        assert_eq!(original, original_from_proto);
+            let original_from_proto = T::try_from(proto_original).unwrap();
+            assert_eq!(original, original_from_proto);
+        }
     }
 
     // The following are tests for structs that do not derive eq
@@ -1842,11 +1884,14 @@ mod tests {
         U: From<T>,
         E: std::fmt::Debug,
     {
-        let original: T = Unit.fake_with_rng(&mut OsRng);
-        let proto_original = U::from(original.clone());
+        // TODO: proptest
+        for _ in 0..25 {
+            let original: T = Unit.fake_with_rng(&mut OsRng);
+            let proto_original = U::from(original.clone());
 
-        let original_from_proto = T::try_from(proto_original).unwrap();
-        assert_eq!(wrapper(original), wrapper(original_from_proto));
+            let original_from_proto = T::try_from(proto_original).unwrap();
+            assert_eq!(wrapper(original), wrapper(original_from_proto));
+        }
     }
 
     #[test]

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, 14")]
+    pub id: ::core::option::Option<wsts_message::Id>,
 }
 /// Nested message and enum types in `WstsMessage`.
 pub mod wsts_message {
@@ -335,6 +338,22 @@ pub mod wsts_message {
             super::super::super::super::crypto::wsts::SignatureShareResponse,
         ),
     }
+    #[derive(Clone, Copy, PartialEq, ::prost::Oneof)]
+    pub enum Id {
+        /// If this WSTS message is related to a Bitcoin signing round, this field
+        /// will be set to the related Bitcoin transaction ID.
+        #[prost(message, tag = "12")]
+        IdBitcoinTxid(super::super::super::super::bitcoin::BitcoinTxid),
+        /// If this WSTS message is related to a rotate-keys transaction, this field
+        /// will be set to the _new_ aggregate public key being verified.
+        #[prost(message, tag = "13")]
+        IdRotateKey(super::super::super::super::crypto::PublicKey),
+        /// If this WSTS message is related to a DKG round, this field will be set
+        /// to the 32-byte id determined based on the coordinator public key and
+        /// block hash, set by the coordinator.
+        #[prost(message, tag = "14")]
+        IdDkg(super::super::super::super::crypto::Uint256),
+    }
 }
 /// 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),
         }
     }