refactor: wsts-related cleanup (#1287)

* some wsts cleanup/refactor

* think that's it

* add message support for a specific dkg wstsmessageid

* add a dkg wstsmessageid variant

* fix merge artifacts

* tracing fields

* storage mut

* lovely cascading changes

* remove wstsmessage.txid infavor of id

* remove save and prefer trait default impls

* remove unused trait methods

* logging stuff

* rename message ids

* remove unneeded allow(deprecated)

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

@@ -35,8 +35,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;
+  reserved 1;
   // The wsts message
   oneof inner {
     // Tell signers to begin DKG by sending DKG public shares
@@ -60,6 +59,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 sweep = 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 dkg_verification = 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 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.
+    Sweep(bitcoin::Txid),
+    /// The WSTS message is related to a rotate key verification operation.
+    DkgVerification(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::Sweep(txid)
+    }
+}
+
+impl From<crate::storage::model::BitcoinTxId> for WstsMessageId {
+    fn from(txid: crate::storage::model::BitcoinTxId) -> Self {
+        Self::Sweep(txid.into())
+    }
+}
+
+impl std::fmt::Display for WstsMessageId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            WstsMessageId::Sweep(txid) => write!(f, "sweep({})", txid),
+            WstsMessageId::DkgVerification(aggregate_key) => {
+                write!(f, "dkg-verification({})", 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 {
@@ -1089,15 +1092,25 @@ impl From<WstsMessage> for proto::WstsMessage {
                 proto::wsts_message::Inner::SignatureShareResponse(inner.into())
             }
         };
+
         proto::WstsMessage {
-            txid: Some(BitcoinTxId::from(value.txid).into()),
+            id: Some(match value.id {
+                WstsMessageId::Sweep(txid) => wsts_message::Id::Sweep(proto::BitcoinTxid {
+                    txid: Some(proto::Uint256::from(BitcoinTxId::from(txid).into_bytes())),
+                }),
+                WstsMessageId::DkgVerification(pubkey) => {
+                    wsts_message::Id::DkgVerification(pubkey.into())
+                }
+                WstsMessageId::Dkg(id) => wsts_message::Id::Dkg(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 +1144,17 @@ impl TryFrom<proto::WstsMessage> for WstsMessage {
                 wsts::net::Message::SignatureShareResponse(inner.try_into()?)
             }
         };
+
         Ok(WstsMessage {
-            txid: BitcoinTxId::try_from(value.txid.required()?)?.into(),
+            id: match value.id.required()? {
+                wsts_message::Id::Sweep(txid) => {
+                    WstsMessageId::Sweep(BitcoinTxId::try_from(txid)?.into())
+                }
+                wsts_message::Id::DkgVerification(pubkey) => {
+                    WstsMessageId::DkgVerification(PublicKey::try_from(pubkey)?)
+                }
+                wsts_message::Id::Dkg(id) => WstsMessageId::Dkg(id.into()),
+            },
             inner,
         })
     }
@@ -1737,21 +1759,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 +1817,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 +1870,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

@@ -286,12 +286,11 @@ pub mod signer_message {
 /// A wsts message.
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct WstsMessage {
-    /// The transaction ID this message relates to, will be a dummy ID for DKG messages
-    #[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 +334,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")]
+        Sweep(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")]
+        DkgVerification(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")]
+        Dkg(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),
         }
     }

signer/src/testing/transaction_coordinator.rs

@@ -1019,8 +1019,9 @@ where
             .into();
 
         let dkg_txid = testing::dummy::txid(&fake::Faker, rng);
-        let (aggregate_key, all_dkg_shares) =
-            signer_set.run_dkg(bitcoin_chain_tip, dkg_txid, rng).await;
+        let (aggregate_key, all_dkg_shares) = signer_set
+            .run_dkg(bitcoin_chain_tip, dkg_txid.into(), rng)
+            .await;
 
         let encrypted_dkg_shares = all_dkg_shares.first().unwrap();
 

signer/src/testing/transaction_signer.rs

@@ -213,7 +213,9 @@ where
             coordinator_signer_info,
             self.signing_threshold,
         );
-        let aggregate_key = coordinator.run_dkg(bitcoin_chain_tip, dummy_txid).await;
+        let aggregate_key = coordinator
+            .run_dkg(bitcoin_chain_tip, dummy_txid.into())
+            .await;
 
         for handle in event_loop_handles.into_iter() {
             assert!(handle
@@ -260,7 +262,7 @@ async fn run_dkg_and_store_results_for_signers<'s: 'r, 'r, S, Rng>(
     let dkg_txid = testing::dummy::txid(&fake::Faker, rng);
     let bitcoin_chain_tip = *chain_tip;
     let (_, all_dkg_shares) = testing_signer_set
-        .run_dkg(bitcoin_chain_tip, dkg_txid, rng)
+        .run_dkg(bitcoin_chain_tip, dkg_txid.into(), rng)
         .await;
 
     for (storage, encrypted_dkg_shares) in stores.into_iter().zip(all_dkg_shares) {