enable batching multiple sign calls in one batch transaction

chain-signatures/contract/src/lib.rs

@@ -64,17 +64,23 @@ impl Default for VersionedMpcContract {
 #[derive(BorshDeserialize, BorshSerialize, Debug)]
 pub struct MpcContract {
     protocol_state: ProtocolContractState,
-    pending_requests: LookupMap<SignatureRequest, YieldIndex>,
+    pending_requests: LookupMap<SignatureRequest, Option<YieldIndex>>,
     request_counter: u32,
     proposed_updates: ProposedUpdates,
     config: Config,
 }
 
 impl MpcContract {
+    fn mark_request_received(&mut self, request: &SignatureRequest) {
+        if self.pending_requests.insert(request, &None).is_none() {
+            self.request_counter += 1;
+        }
+    }
+
     fn add_request(&mut self, request: &SignatureRequest, data_id: CryptoHash) {
         if self
             .pending_requests
-            .insert(request, &YieldIndex { data_id })
+            .insert(request, &Some(YieldIndex { data_id }))
             .is_none()
         {
             self.request_counter += 1;
@@ -166,6 +172,7 @@ impl VersionedMpcContract {
                 "sign: predecessor={predecessor}, payload={payload:?}, path={path:?}, key_version={key_version}",
             );
             env::log_str(&serde_json::to_string(&near_sdk::env::random_seed_array()).unwrap());
+            self.mark_request_received(&request);
             let contract_signature_request = ContractSignatureRequest {
                 request,
                 requester: predecessor,
@@ -275,7 +282,7 @@ impl VersionedMpcContract {
 
             match self {
                 Self::V0(mpc_contract) => {
-                    if let Some(YieldIndex { data_id }) =
+                    if let Some(Some(YieldIndex { data_id })) =
                         mpc_contract.pending_requests.get(&request)
                     {
                         env::promise_yield_resume(
@@ -803,6 +810,12 @@ impl VersionedMpcContract {
         }
     }
 
+    fn mark_request_received(&mut self, request: &SignatureRequest) {
+        match self {
+            Self::V0(ref mut mpc_contract) => mpc_contract.mark_request_received(request),
+        }
+    }
+
     fn threshold(&self) -> Result<usize, Error> {
         match self {
             Self::V0(contract) => match &contract.protocol_state {

chain-signatures/node/Cargo.toml

@@ -15,6 +15,7 @@ aws-sdk-s3 = "1.29"
 aws-types = "1.2"
 axum = { version = "0.6.19" }
 axum-extra = "0.7"
+borsh = "1.5.0"
 cait-sith = { git = "https://github.com/LIT-Protocol/cait-sith.git", features = [
     "k256",
 ], rev = "8ad2316" }

chain-signatures/node/src/protocol/message.rs

@@ -1,5 +1,6 @@
 use super::cryptography::CryptographicError;
 use super::presignature::{GenerationError, PresignatureId};
+use super::signature::SignRequestIdentifier;
 use super::state::{GeneratingState, NodeState, ResharingState, RunningState};
 use super::triple::TripleId;
 use crate::gcp::error::SecretStorageError;
@@ -103,7 +104,7 @@ pub struct MpcMessageQueue {
     resharing_bins: HashMap<u64, VecDeque<ResharingMessage>>,
     triple_bins: HashMap<u64, HashMap<TripleId, VecDeque<TripleMessage>>>,
     presignature_bins: HashMap<u64, HashMap<PresignatureId, VecDeque<PresignatureMessage>>>,
-    signature_bins: HashMap<u64, HashMap<CryptoHash, VecDeque<SignatureMessage>>>,
+    signature_bins: HashMap<u64, HashMap<SignRequestIdentifier, VecDeque<SignatureMessage>>>,
 }
 
 impl MpcMessageQueue {
@@ -133,7 +134,11 @@ impl MpcMessageQueue {
                 .signature_bins
                 .entry(message.epoch)
                 .or_default()
-                .entry(message.receipt_id)
+                .entry(SignRequestIdentifier::new(
+                    message.receipt_id,
+                    message.epsilon,
+                    message.request.payload,
+                ))
                 .or_default()
                 .push_back(message),
         }
@@ -366,7 +371,7 @@ impl MessageHandler for RunningState {
 
         let mut signature_manager = self.signature_manager.write().await;
         let signature_messages = queue.signature_bins.entry(self.epoch).or_default();
-        signature_messages.retain(|receipt_id, queue| {
+        signature_messages.retain(|sign_request_identifier, queue| {
             // Skip message if it already timed out
             if queue.is_empty()
                 || queue.iter().any(|msg| {
@@ -379,9 +384,9 @@ impl MessageHandler for RunningState {
                 return false;
             }
 
-            !signature_manager.refresh_gc(receipt_id)
+            !signature_manager.refresh_gc(sign_request_identifier)
         });
-        for (receipt_id, queue) in signature_messages {
+        for (sign_request_identifier, queue) in signature_messages {
             // SAFETY: this unwrap() is safe since we have already checked that the queue is not empty.
             let SignatureMessage {
                 proposer,
@@ -414,7 +419,7 @@ impl MessageHandler for RunningState {
             // TODO: Validate that the message matches our sign_queue
             let protocol = match signature_manager.get_or_generate(
                 participants,
-                *receipt_id,
+                sign_request_identifier.receipt_id,
                 *proposer,
                 *presignature_id,
                 request,
@@ -437,15 +442,19 @@ impl MessageHandler for RunningState {
                     // and have the other nodes timeout in the following cases:
                     // - If a presignature is in GC, then it was used already or failed to be produced.
                     // - If a presignature is missing, that means our system cannot process this signature.
-                    tracing::warn!(%receipt_id, ?err, "signature cannot be generated");
+                    tracing::warn!(
+                        ?sign_request_identifier,
+                        ?err,
+                        "signature cannot be generated"
+                    );
                     queue.clear();
                     continue;
                 }
                 Err(GenerationError::CaitSithInitializationError(error)) => {
                     // ignore the whole of the messages since the generation had bad parameters. Also have the other node who
                     // initiated the protocol resend the message or have it timeout on their side.
                     tracing::warn!(
-                        ?receipt_id,
+                        ?sign_request_identifier,
                         presignature_id,
                         ?error,
                         "unable to initialize incoming signature protocol"
@@ -455,7 +464,7 @@ impl MessageHandler for RunningState {
                 }
                 Err(err) => {
                     tracing::warn!(
-                        ?receipt_id,
+                        ?sign_request_identifier,
                         ?err,
                         "Unexpected error encounted while generating signature"
                     );

chain-signatures/node/src/protocol/presignature.rs

@@ -234,7 +234,7 @@ impl PresignatureManager {
             )));
         }
 
-        tracing::info!(id, "starting protocol to generate a new presignature");
+        tracing::debug!(id, "starting protocol to generate a new presignature");
         let generator = Self::generate_internal(
             participants,
             self.me,