feat: implement presignature generation (#342)

Co-authored-by: Phuong Nguyen <ChaoticTempest@users.noreply.github.com>

.github/workflows/multichain-integration.yml

@@ -40,6 +40,7 @@ jobs:
         run: |
           docker pull ghcr.io/near/os-relayer:12ba6e35690df3979fce0b36a41d0ca0db9c0ab4
           docker pull ghcr.io/near/near-lake-indexer:e6519c922435f3d18b5f2ddac5d1ec171ef4dd6b
+          docker pull localstack/localstack:latest
 
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v2

integration-tests/src/multichain/mod.rs

@@ -23,6 +23,17 @@ pub enum Nodes<'a> {
 }
 
 impl Nodes<'_> {
+    pub fn len(&self) -> usize {
+        match self {
+            Nodes::Local { nodes, .. } => nodes.len(),
+            Nodes::Docker { nodes, .. } => nodes.len(),
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
     pub fn ctx(&self) -> &Context {
         match self {
             Nodes::Local { ctx, .. } => ctx,

integration-tests/tests/lib.rs

@@ -243,6 +243,35 @@ mod wait_for {
             .retry(&ExponentialBuilder::default().with_max_times(6))
             .await
     }
+
+    pub async fn has_at_least_presignatures<'a>(
+        ctx: &MultichainTestContext<'a>,
+        id: usize,
+        expected_presignature_count: usize,
+    ) -> anyhow::Result<StateView> {
+        let is_enough_presignatures = || async {
+            let state_view: StateView = ctx
+                .http_client
+                .get(format!("{}/state", ctx.nodes.url(id)))
+                .send()
+                .await?
+                .json()
+                .await?;
+
+            match state_view {
+                StateView::Running {
+                    presignature_count, ..
+                } if presignature_count >= expected_presignature_count => Ok(state_view),
+                StateView::Running { .. } => {
+                    anyhow::bail!("node does not have enough presignatures yet")
+                }
+                StateView::NotRunning => anyhow::bail!("node is not running"),
+            }
+        };
+        is_enough_presignatures
+            .retry(&ExponentialBuilder::default().with_max_times(6))
+            .await
+    }
 }
 
 trait MpcCheck {

integration-tests/tests/multichain/mod.rs

@@ -30,10 +30,19 @@ async fn test_multichain_reshare() -> anyhow::Result<()> {
 }
 
 #[test(tokio::test)]
-async fn test_triples() -> anyhow::Result<()> {
+async fn test_triples_and_presignatures() -> anyhow::Result<()> {
     with_multichain_nodes(3, |ctx| {
         Box::pin(async move {
-            wait_for::has_at_least_triples(&ctx, 0, 2).await?;
+            // Wait for network to complete key generation
+            let state_0 = wait_for::running_mpc(&ctx, 0).await?;
+            assert_eq!(state_0.participants.len(), 3);
+
+            for i in 0..ctx.nodes.len() {
+                wait_for::has_at_least_triples(&ctx, i, 2).await?;
+            }
+            for i in 0..ctx.nodes.len() {
+                wait_for::has_at_least_presignatures(&ctx, i, 2).await?;
+            }
 
             Ok(())
         })

node/src/protocol/consensus.rs

@@ -3,6 +3,7 @@ use super::state::{
     JoiningState, NodeState, PersistentNodeData, RunningState, StartedState,
     WaitingForConsensusState,
 };
+use crate::protocol::presignature::PresignatureManager;
 use crate::protocol::state::{GeneratingState, ResharingState};
 use crate::protocol::triple::TripleManager;
 use crate::types::PrivateKeyShare;
@@ -86,7 +87,7 @@ impl ConsensusProtocol for StartedState {
                                 tracing::info!(
                                     "contract state is running and we are already a participant"
                                 );
-                                let participants_vec =
+                                let participants_vec: Vec<Participant> =
                                     contract_state.participants.keys().cloned().collect();
                                 Ok(NodeState::Running(RunningState {
                                     epoch,
@@ -95,6 +96,12 @@ impl ConsensusProtocol for StartedState {
                                     private_share,
                                     public_key,
                                     triple_manager: TripleManager::new(
+                                        participants_vec.clone(),
+                                        ctx.me(),
+                                        contract_state.threshold,
+                                        epoch,
+                                    ),
+                                    presignature_manager: PresignatureManager::new(
                                         participants_vec,
                                         ctx.me(),
                                         contract_state.threshold,
@@ -261,14 +268,21 @@ impl ConsensusProtocol for WaitingForConsensusState {
                     if contract_state.public_key != self.public_key {
                         return Err(ConsensusError::MismatchedPublicKey);
                     }
-                    let participants_vec = self.participants.keys().cloned().collect();
+                    let participants_vec: Vec<Participant> =
+                        self.participants.keys().cloned().collect();
                     Ok(NodeState::Running(RunningState {
                         epoch: self.epoch,
                         participants: self.participants,
                         threshold: self.threshold,
                         private_share: self.private_share,
                         public_key: self.public_key,
                         triple_manager: TripleManager::new(
+                            participants_vec.clone(),
+                            ctx.me(),
+                            self.threshold,
+                            self.epoch,
+                        ),
+                        presignature_manager: PresignatureManager::new(
                             participants_vec,
                             ctx.me(),
                             self.threshold,

node/src/protocol/cryptography.rs

@@ -109,7 +109,7 @@ impl CryptographicProtocol for ResharingState {
     ) -> Result<NodeState, CryptographicError> {
         tracing::info!("progressing key reshare");
         loop {
-            let action = self.protocol.poke().unwrap();
+            let action = self.protocol.poke()?;
             match action {
                 Action::Wait => {
                     tracing::debug!("waiting");
@@ -173,13 +173,39 @@ impl CryptographicProtocol for RunningState {
         mut self,
         ctx: C,
     ) -> Result<NodeState, CryptographicError> {
-        if self.triple_manager.potential_len() < 2 {
+        if self.triple_manager.my_len() < 2 {
             self.triple_manager.generate()?;
         }
         for (p, msg) in self.triple_manager.poke()? {
             let url = self.participants.get(&p).unwrap();
             http_client::message(ctx.http_client(), url.clone(), MpcMessage::Triple(msg)).await?;
         }
+
+        if self.presignature_manager.potential_len() < 2 {
+            // To ensure there is no contention between different nodes we are only using triples
+            // that we proposed. This way in a non-BFT environment we are guaranteed to never try
+            // to use the same triple as any other node.
+            if let Some((triple0, triple1)) = self.triple_manager.take_mine_twice() {
+                self.presignature_manager.generate(
+                    triple0,
+                    triple1,
+                    &self.public_key,
+                    &self.private_share,
+                )?;
+            } else {
+                tracing::debug!("we don't have enough triples to generate a presignature");
+            }
+        }
+        for (p, msg) in self.presignature_manager.poke()? {
+            let url = self.participants.get(&p).unwrap();
+            http_client::message(
+                ctx.http_client(),
+                url.clone(),
+                MpcMessage::Presignature(msg),
+            )
+            .await?;
+        }
+
         Ok(NodeState::Running(self))
     }
 }

node/src/protocol/message.rs

@@ -29,18 +29,30 @@ pub struct TripleMessage {
     pub data: MessageData,
 }
 
+#[derive(Serialize, Deserialize, Debug)]
+pub struct PresignatureMessage {
+    pub id: u64,
+    pub triple0: u64,
+    pub triple1: u64,
+    pub epoch: u64,
+    pub from: Participant,
+    pub data: MessageData,
+}
+
 #[derive(Serialize, Deserialize, Debug)]
 pub enum MpcMessage {
     Generating(GeneratingMessage),
     Resharing(ResharingMessage),
     Triple(TripleMessage),
+    Presignature(PresignatureMessage),
 }
 
 #[derive(Default)]
 pub struct MpcMessageQueue {
     generating: VecDeque<GeneratingMessage>,
     resharing_bins: HashMap<u64, VecDeque<ResharingMessage>>,
     triple_bins: HashMap<u64, HashMap<u64, VecDeque<TripleMessage>>>,
+    presignature_bins: HashMap<u64, HashMap<u64, VecDeque<PresignatureMessage>>>,
 }
 
 impl MpcMessageQueue {
@@ -59,6 +71,13 @@ impl MpcMessageQueue {
                 .entry(message.id)
                 .or_default()
                 .push_back(message),
+            MpcMessage::Presignature(message) => self
+                .presignature_bins
+                .entry(message.epoch)
+                .or_default()
+                .entry(message.id)
+                .or_default()
+                .push_back(message),
         }
     }
 }
@@ -121,6 +140,20 @@ impl MessageHandler for RunningState {
                 }
             }
         }
+        for (id, queue) in queue.presignature_bins.entry(self.epoch).or_default() {
+            while let Some(message) = queue.pop_front() {
+                if let Some(protocol) = self.presignature_manager.get_or_generate(
+                    *id,
+                    message.triple0,
+                    message.triple1,
+                    &mut self.triple_manager,
+                    &self.public_key,
+                    &self.private_share,
+                )? {
+                    protocol.message(message.from, message.data);
+                }
+            }
+        }
         Ok(())
     }
 }

node/src/protocol/mod.rs

@@ -2,6 +2,7 @@ mod consensus;
 mod contract;
 mod cryptography;
 mod message;
+mod presignature;
 mod state;
 mod triple;