Embed public and private shares into DkgBegin messages

protobufs/crypto/wsts/wsts.proto

@@ -47,6 +47,8 @@ message DkgPrivateBegin {
   repeated uint32 signer_ids = 2;
   // Key IDs who responded in time for this DKG round
   repeated uint32 key_ids = 3;
+  // Public shares from all signers in this DKG round
+  DkgPublicShares dkg_public_shares = 4;
 }
 
 // DKG private shares message from signer to all signers and coordinator
@@ -89,6 +91,8 @@ message DkgEndBegin {
   repeated uint32 signer_ids = 2;
   // Key IDs who responded in time for this DKG round
   repeated uint32 key_ids = 3;
+  // Private shares from all signers in this DKG round
+  map<uint32, DkgPrivateShares> dkg_private_shares = 4;
 }
 
 // DKG end message from signers to coordinator

signer/Cargo.toml

@@ -50,7 +50,7 @@ tracing-attributes.workspace = true
 tracing-subscriber = { workspace = true }
 url.workspace = true
 # wsts.workspace = true
-wsts = { git = "https://github.com/Trust-Machines/wsts.git", rev = "ebd7d7775ad5e44cdbf4f5c1fb468bdf6c467265" }
+wsts = { git = "https://github.com/Trust-Machines/wsts.git", rev = "e78101f923f517440e389f9da7bbf0e89bd901c8" }
 zeromq.workspace = true
 hex.workspace = true
 cfg-if = "1.0"

signer/src/main.rs

@@ -320,7 +320,7 @@ async fn run_transaction_signer(ctx: impl Context) -> Result<(), Error> {
         rng: rand::thread_rng(),
         signer_private_key: config.signer.private_key,
         wsts_state_machines: HashMap::new(),
-        dkg_begin_pause: Some(Duration::from_secs(10)),
+        dkg_begin_pause: None, //Some(Duration::from_secs(10)),
     };
 
     signer.run().await

signer/src/proto/convert.rs

@@ -14,6 +14,7 @@ use bitcoin::OutPoint;
 use bitvec::array::BitArray;
 use clarity::codec::StacksMessageCodec as _;
 use clarity::vm::types::PrincipalData;
+use hashbrown::HashMap;
 use p256k1::point::Point;
 use p256k1::scalar::Scalar;
 use polynomial::Polynomial;
@@ -649,21 +650,36 @@ impl From<proto::DkgBegin> for DkgBegin {
 
 impl From<DkgPrivateBegin> for proto::DkgPrivateBegin {
     fn from(value: DkgPrivateBegin) -> Self {
+        let shares = value
+            .dkg_public_shares
+            .iter()
+            .map(|(i, s)| (*i, proto::SignerDkgPublicShares::from(s.clone())))
+            .collect();
+        let dkg_public_shares = Some(proto::DkgPublicShares { shares });
         proto::DkgPrivateBegin {
             dkg_id: value.dkg_id,
             signer_ids: value.signer_ids,
             key_ids: value.key_ids,
+            dkg_public_shares,
         }
     }
 }
 
-impl From<proto::DkgPrivateBegin> for DkgPrivateBegin {
-    fn from(value: proto::DkgPrivateBegin) -> Self {
-        DkgPrivateBegin {
+impl TryFrom<proto::DkgPrivateBegin> for DkgPrivateBegin {
+    type Error = Error;
+    fn try_from(value: proto::DkgPrivateBegin) -> Result<Self, Self::Error> {
+        let mut dkg_public_shares = HashMap::new();
+        if let Some(shares) = value.dkg_public_shares {
+            for (id, share) in shares.shares {
+                dkg_public_shares.insert(id, DkgPublicShares::try_from(share)?);
+            }
+        }
+        Ok(DkgPrivateBegin {
             dkg_id: value.dkg_id,
             signer_ids: value.signer_ids,
             key_ids: value.key_ids,
-        }
+            dkg_public_shares,
+        })
     }
 }
 
@@ -720,17 +736,28 @@ impl From<DkgEndBegin> for proto::DkgEndBegin {
             dkg_id: value.dkg_id,
             signer_ids: value.signer_ids,
             key_ids: value.key_ids,
+            dkg_private_shares: value
+                .dkg_private_shares
+                .iter()
+                .map(|(i, s)| (*i, proto::DkgPrivateShares::from(s.clone())))
+                .collect(),
         }
     }
 }
 
-impl From<proto::DkgEndBegin> for DkgEndBegin {
-    fn from(value: proto::DkgEndBegin) -> Self {
-        DkgEndBegin {
+impl TryFrom<proto::DkgEndBegin> for DkgEndBegin {
+    type Error = Error;
+    fn try_from(value: proto::DkgEndBegin) -> Result<Self, Self::Error> {
+        let mut dkg_private_shares = HashMap::new();
+        for (id, shares) in value.dkg_private_shares {
+            dkg_private_shares.insert(id, DkgPrivateShares::try_from(shares.clone())?);
+        }
+        Ok(DkgEndBegin {
             dkg_id: value.dkg_id,
             signer_ids: value.signer_ids,
             key_ids: value.key_ids,
-        }
+            dkg_private_shares,
+        })
     }
 }
 
@@ -1131,13 +1158,13 @@ impl TryFrom<proto::WstsMessage> for WstsMessage {
                 wsts::net::Message::DkgPublicShares(inner.try_into()?)
             }
             proto::wsts_message::Inner::DkgPrivateBegin(inner) => {
-                wsts::net::Message::DkgPrivateBegin(inner.into())
+                wsts::net::Message::DkgPrivateBegin(inner.try_into()?)
             }
             proto::wsts_message::Inner::DkgPrivateShares(inner) => {
                 wsts::net::Message::DkgPrivateShares(inner.try_into()?)
             }
             proto::wsts_message::Inner::DkgEndBegin(inner) => {
-                wsts::net::Message::DkgEndBegin(inner.into())
+                wsts::net::Message::DkgEndBegin(inner.try_into()?)
             }
             proto::wsts_message::Inner::DkgEnd(inner) => {
                 wsts::net::Message::DkgEnd(inner.try_into()?)
@@ -1933,11 +1960,12 @@ mod tests {
         }
     }
     impl Dummy<Unit> for DkgPrivateBegin {
-        fn dummy_with_rng<R: rand::Rng + ?Sized>(_: &Unit, rng: &mut R) -> Self {
+        fn dummy_with_rng<R: rand::Rng + ?Sized>(config: &Unit, rng: &mut R) -> Self {
             DkgPrivateBegin {
                 dkg_id: Faker.fake_with_rng(rng),
                 signer_ids: Faker.fake_with_rng(rng),
                 key_ids: Faker.fake_with_rng(rng),
+                dkg_public_shares: config.fake_with_rng(rng),
             }
         }
     }
@@ -1962,11 +1990,12 @@ mod tests {
     }
 
     impl Dummy<Unit> for DkgEndBegin {
-        fn dummy_with_rng<R: rand::Rng + ?Sized>(_: &Unit, rng: &mut R) -> Self {
+        fn dummy_with_rng<R: rand::Rng + ?Sized>(config: &Unit, rng: &mut R) -> Self {
             DkgEndBegin {
                 dkg_id: Faker.fake_with_rng(rng),
                 signer_ids: Faker.fake_with_rng(rng),
                 key_ids: Faker.fake_with_rng(rng),
+                dkg_private_shares: config.fake_with_rng(rng),
             }
         }
     }
@@ -2217,6 +2246,24 @@ mod tests {
         }
     }
 
+    impl Dummy<Unit> for hashbrown::HashMap<u32, DkgPublicShares> {
+        fn dummy_with_rng<R: rand::Rng + ?Sized>(config: &Unit, rng: &mut R) -> Self {
+            fake::vec![(); 0..20]
+                .into_iter()
+                .map(|_| (Faker.fake_with_rng(rng), config.fake_with_rng(rng)))
+                .collect()
+        }
+    }
+
+    impl Dummy<Unit> for hashbrown::HashMap<u32, DkgPrivateShares> {
+        fn dummy_with_rng<R: rand::Rng + ?Sized>(config: &Unit, rng: &mut R) -> Self {
+            fake::vec![(); 0..20]
+                .into_iter()
+                .map(|_| (Faker.fake_with_rng(rng), config.fake_with_rng(rng)))
+                .collect()
+        }
+    }
+
     #[test]
     fn conversion_between_bytes_and_uint256() {
         let number = proto::Uint256 {

signer/src/proto/generated/crypto.wsts.rs

@@ -137,6 +137,9 @@ pub struct DkgPrivateBegin {
     /// Key IDs who responded in time for this DKG round
     #[prost(uint32, repeated, tag = "3")]
     pub key_ids: ::prost::alloc::vec::Vec<u32>,
+    /// Public shares from all signers in this DKG round
+    #[prost(message, optional, tag = "4")]
+    pub dkg_public_shares: ::core::option::Option<DkgPublicShares>,
 }
 /// DKG private shares message from signer to all signers and coordinator
 /// This maps to this type <<https://github.com/Trust-Machines/wsts/blob/2d6cb87218bb8dd9ed0519356afe57a0b9a697cb/src/net.rs#L185-L195>>
@@ -193,6 +196,9 @@ pub struct DkgEndBegin {
     /// Key IDs who responded in time for this DKG round
     #[prost(uint32, repeated, tag = "3")]
     pub key_ids: ::prost::alloc::vec::Vec<u32>,
+    /// Private shares from all signers in this DKG round
+    #[prost(btree_map = "uint32, message", tag = "4")]
+    pub dkg_private_shares: ::prost::alloc::collections::BTreeMap<u32, DkgPrivateShares>,
 }
 /// DKG end message from signers to coordinator
 /// This maps to this type <<https://github.com/Trust-Machines/wsts/blob/2d6cb87218bb8dd9ed0519356afe57a0b9a697cb/src/net.rs#L246-L255>>

signer/src/testing/message.rs

@@ -132,6 +132,7 @@ impl fake::Dummy<fake::Faker> for message::WstsMessage {
             dkg_id: config.fake_with_rng(rng),
             signer_ids: config.fake_with_rng(rng),
             key_ids: config.fake_with_rng(rng),
+            dkg_private_shares: Default::default(),
         };
 
         Self {

signer/src/testing/wsts.rs

@@ -1,5 +1,6 @@
 //! Test utilities for running a wsts signer and coordinator.
 
+use rand::rngs::OsRng;
 use std::collections::BTreeMap;
 use std::collections::BTreeSet;
 use std::time::Duration;
@@ -122,6 +123,7 @@ impl Coordinator {
             sign_timeout: None,
             signer_key_ids,
             signer_public_keys,
+            embed_public_private_shares: true,
         };
 
         let wsts_coordinator = fire::Coordinator::new(config);
@@ -313,6 +315,7 @@ impl Signer {
     /// Participate in a DKG round and return the result
     pub async fn run_until_dkg_end(mut self) -> Self {
         let future = async move {
+            let mut rng = OsRng;
             loop {
                 let msg = self.network.receive().await.expect("network error");
                 let bitcoin_chain_tip = msg.bitcoin_chain_tip;
@@ -328,12 +331,12 @@ impl Signer {
 
                 let outbound_packets = self
                     .wsts_signer
-                    .process_inbound_messages(&[packet])
+                    .process_inbound_messages(&[packet], &mut rng)
                     .expect("message processing failed");
 
                 for packet in outbound_packets {
                     self.wsts_signer
-                        .process_inbound_messages(&[packet.clone()])
+                        .process_inbound_messages(&[packet.clone()], &mut rng)
                         .expect("message processing failed");
 
                     self.send_packet(bitcoin_chain_tip, wsts_msg.txid, packet.clone())
@@ -353,6 +356,7 @@ impl Signer {
     /// Participate in a signing round and return the result
     pub async fn run_until_signature_share_response(mut self) -> Self {
         let future = async move {
+            let mut rng = OsRng;
             loop {
                 let msg = self.network.receive().await.expect("network error");
                 let bitcoin_chain_tip = msg.bitcoin_chain_tip;
@@ -368,12 +372,12 @@ impl Signer {
 
                 let outbound_packets = self
                     .wsts_signer
-                    .process_inbound_messages(&[packet])
+                    .process_inbound_messages(&[packet], &mut rng)
                     .expect("message processing failed");
 
                 for packet in outbound_packets {
                     self.wsts_signer
-                        .process_inbound_messages(&[packet.clone()])
+                        .process_inbound_messages(&[packet.clone()], &mut rng)
                         .expect("message processing failed");
 
                     self.send_packet(bitcoin_chain_tip, wsts_msg.txid, packet.clone())