enable batching multiple sign calls in one batch transaction (#909)

* enable batching multiple sign calls in one batch transaction

* remove receipt_id from anywhere but indexer

chain-signatures/contract/src/errors/mod.rs

@@ -6,7 +6,7 @@ pub enum SignError {
     #[error("Signature request has timed out.")]
     Timeout,
     #[error("Signature request has already been submitted. Please try again later.")]
-    PayloadCollision,
+    RequestCollision,
     #[error(
         "This key version is not supported. Call latest_key_version() to get the latest supported version."
     )]

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,
@@ -174,7 +181,7 @@ impl VersionedMpcContract {
             };
             Ok(Self::ext(env::current_account_id()).sign_helper(contract_signature_request))
         } else {
-            Err(SignError::PayloadCollision.into())
+            Err(SignError::RequestCollision.into())
         }
     }
 
@@ -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/indexer.rs

@@ -239,7 +239,7 @@ async fn handle_block(
                     key_version: arguments.request.key_version,
                 };
                 pending_requests.push(SignRequest {
-                    receipt_id,
+                    request_id: receipt_id.0,
                     request,
                     epsilon,
                     entropy,

chain-signatures/node/src/kdf.rs

@@ -9,12 +9,12 @@ use sha3::Sha3_256;
 // that we generate different random scalars as delta tweaks.
 // Receipt ID should be unique inside of a block, so it serves us as the request identifier.
 pub fn derive_delta(
-    receipt_id: CryptoHash,
+    request_id: [u8; 32],
     entropy: [u8; 32],
     presignature_big_r: AffinePoint,
 ) -> Scalar {
     let hk = Hkdf::<Sha3_256>::new(None, &entropy);
-    let info = format!("{DELTA_DERIVATION_PREFIX}:{}", receipt_id);
+    let info = format!("{DELTA_DERIVATION_PREFIX}:{}", CryptoHash(request_id));
     let mut okm = [0u8; 32];
     hk.expand(info.as_bytes(), &mut okm).unwrap();
     hk.expand(

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;
@@ -13,7 +14,6 @@ use cait_sith::protocol::{InitializationError, MessageData, Participant, Protoco
 use k256::Scalar;
 use mpc_keys::hpke::{self, Ciphered};
 use near_crypto::Signature;
-use near_primitives::hash::CryptoHash;
 use serde::{Deserialize, Serialize};
 use std::collections::{HashMap, VecDeque};
 use std::sync::Arc;
@@ -63,7 +63,7 @@ pub struct PresignatureMessage {
 
 #[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
 pub struct SignatureMessage {
-    pub receipt_id: CryptoHash,
+    pub request_id: [u8; 32],
     pub proposer: Participant,
     pub presignature_id: PresignatureId,
     pub request: ContractSignRequest,
@@ -103,7 +103,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 +133,11 @@ impl MpcMessageQueue {
                 .signature_bins
                 .entry(message.epoch)
                 .or_default()
-                .entry(message.receipt_id)
+                .entry(SignRequestIdentifier::new(
+                    message.request_id,
+                    message.epsilon,
+                    message.request.payload,
+                ))
                 .or_default()
                 .push_back(message),
         }
@@ -366,7 +370,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 +383,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 +418,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.request_id,
                 *proposer,
                 *presignature_id,
                 request,
@@ -437,15 +441,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 +463,7 @@ impl MessageHandler for RunningState {
                 }
                 Err(err) => {
                     tracing::warn!(
-                        ?receipt_id,
+                        ?sign_request_identifier,
                         ?err,
                         "Unexpected error encounted while generating signature"
                     );