fix: stop infinite process of presignature message missing triples

node/Cargo.toml

@@ -19,6 +19,7 @@ cait-sith = { git = "https://github.com/LIT-Protocol/cait-sith.git", features =
     "k256",
 ], rev = "8ad2316"}
 clap = { version = "4.2", features = ["derive", "env"] }
+chrono = "0.4.24"
 google-datastore1 = "5"
 google-secretmanager1 = "5"
 hex = "0.4.3"

node/src/protocol/message.rs

@@ -5,6 +5,7 @@ use super::triple::TripleId;
 use crate::gcp::error::SecretStorageError;
 use crate::http_client::SendError;
 use crate::mesh::Mesh;
+use crate::util;
 
 use async_trait::async_trait;
 use cait_sith::protocol::{InitializationError, MessageData, Participant, ProtocolError};
@@ -43,6 +44,8 @@ pub struct TripleMessage {
     pub epoch: u64,
     pub from: Participant,
     pub data: MessageData,
+    // UNIX timestamp as seconds since the epoch
+    pub timestamp: u64,
 }
 
 #[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
@@ -53,6 +56,8 @@ pub struct PresignatureMessage {
     pub epoch: u64,
     pub from: Participant,
     pub data: MessageData,
+    // UNIX timestamp as seconds since the epoch
+    pub timestamp: u64,
 }
 
 #[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
@@ -66,6 +71,8 @@ pub struct SignatureMessage {
     pub epoch: u64,
     pub from: Participant,
     pub data: MessageData,
+    // UNIX timestamp as seconds since the epoch
+    pub timestamp: u64,
 }
 
 #[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
@@ -251,6 +258,14 @@ impl MessageHandler for RunningState {
         for (id, queue) in queue.presignature_bins.entry(self.epoch).or_default() {
             let mut leftover_messages = Vec::new();
             while let Some(message) = queue.pop_front() {
+                // Skip message if it already timed out
+                if util::is_elapsed_longer_than_timeout(
+                    message.timestamp,
+                    crate::types::PROTOCOL_PRESIG_TIMEOUT,
+                ) {
+                    continue;
+                }
+
                 match presignature_manager
                     .get_or_generate(
                         participants,
@@ -267,8 +282,12 @@ impl MessageHandler for RunningState {
                     Err(presignature::GenerationError::AlreadyGenerated) => {
                         tracing::info!(id, "presignature already generated, nothing left to do")
                     }
+                    Err(presignature::GenerationError::TripleIsGenerating(_)) => {
+                        // Store the message until triple gets generated
+                        leftover_messages.push(message)
+                    }
                     Err(presignature::GenerationError::TripleIsMissing(_)) => {
-                        // Store the message until we are ready to process it
+                        // Store the message until triple is ready
                         leftover_messages.push(message)
                     }
                     Err(presignature::GenerationError::CaitSithInitializationError(error)) => {
@@ -293,6 +312,13 @@ impl MessageHandler for RunningState {
         for (receipt_id, queue) in queue.signature_bins.entry(self.epoch).or_default() {
             let mut leftover_messages = Vec::new();
             while let Some(message) = queue.pop_front() {
+                // Skip message if it already timed out
+                if util::is_elapsed_longer_than_timeout(
+                    message.timestamp,
+                    crate::types::PROTOCOL_SIGNATURE_TIMEOUT,
+                ) {
+                    continue;
+                }
                 tracing::info!(
                     presignature_id = message.presignature_id,
                     "new signature message"

node/src/protocol/presignature.rs

@@ -7,6 +7,7 @@ use crate::types::{PresignatureProtocol, PublicKey, SecretKeyShare};
 use crate::util::AffinePointExt;
 use cait_sith::protocol::{Action, InitializationError, Participant, ProtocolError};
 use cait_sith::{KeygenOutput, PresignArguments, PresignOutput};
+use chrono::Utc;
 use k256::Secp256k1;
 use near_lake_primitives::AccountId;
 use std::collections::hash_map::Entry;
@@ -86,6 +87,8 @@ pub enum GenerationError {
     CaitSithInitializationError(#[from] InitializationError),
     #[error("datastore storage error: {0}")]
     DatastoreStorageError(#[from] DatastoreStorageError),
+    #[error("triple {0} is generating")]
+    TripleIsGenerating(TripleId),
 }
 
 /// Abstracts how triples are generated by providing a way to request a new triple that will be
@@ -102,7 +105,6 @@ pub struct PresignatureManager {
     /// The set of presignatures that were already taken. This will be maintained for at most
     /// presignature timeout period just so messages are cycled through the system.
     taken: HashMap<PresignatureId, Instant>,
-
     me: Participant,
     threshold: usize,
     epoch: u64,
@@ -311,16 +313,31 @@ impl PresignatureManager {
                     tracing::info!(id, "joining protocol to generate a new presignature");
                     let (triple0, triple1) = match triple_manager.take_two(triple0, triple1).await {
                         Ok(result) => result,
-                        Err(error) => {
-                            tracing::warn!(
-                                ?error,
-                                id,
-                                triple0,
-                                triple1,
-                                "could not initiate non-introduced presignature: triple might not have completed for this node yet"
-                            );
-                            return Err(error);
-                        }
+                        Err(error) => match error {
+                            GenerationError::TripleIsGenerating(_) => {
+                                tracing::warn!(
+                                        ?error,
+                                        id,
+                                        triple0,
+                                        triple1,
+                                        "could not initiate non-introduced presignature: one triple is generating"
+                                    );
+                                return Err(error);
+                            }
+                            GenerationError::TripleIsMissing(_) => {
+                                tracing::warn!(
+                                        ?error,
+                                        id,
+                                        triple0,
+                                        triple1,
+                                        "could not initiate non-introduced presignature: one triple is missing"
+                                    );
+                                return Err(error);
+                            }
+                            _ => {
+                                return Err(error);
+                            }
+                        },
                     };
                     let generator = Self::generate_internal(
                         participants,
@@ -396,6 +413,7 @@ impl PresignatureManager {
                                     epoch: self.epoch,
                                     from: self.me,
                                     data: data.clone(),
+                                    timestamp: Utc::now().timestamp() as u64
                                 },
                             ))
                         }
@@ -409,6 +427,7 @@ impl PresignatureManager {
                             epoch: self.epoch,
                             from: self.me,
                             data,
+                            timestamp: Utc::now().timestamp() as u64
                         },
                     )),
                     Action::Return(output) => {

node/src/protocol/signature.rs

@@ -6,6 +6,7 @@ use crate::types::{PublicKey, SignatureProtocol};
 use crate::util::{AffinePointExt, ScalarExt};
 use cait_sith::protocol::{Action, InitializationError, Participant, ProtocolError};
 use cait_sith::{FullSignature, PresignOutput};
+use chrono::Utc;
 use k256::{Scalar, Secp256k1};
 use near_crypto::Signer;
 use near_fetch::signer::ExposeAccountId;
@@ -386,6 +387,7 @@ impl SignatureManager {
                                     epoch: self.epoch,
                                     from: self.me,
                                     data: data.clone(),
+                                    timestamp: Utc::now().timestamp() as u64
                                 },
                             ))
                         }
@@ -402,6 +404,7 @@ impl SignatureManager {
                             epoch: self.epoch,
                             from: self.me,
                             data,
+                            timestamp: Utc::now().timestamp() as u64
                         },
                     )),
                     Action::Return(output) => {

node/src/protocol/triple.rs

@@ -9,6 +9,7 @@ use crate::types::TripleProtocol;
 use crate::util::AffinePointExt;
 use cait_sith::protocol::{Action, InitializationError, Participant, ProtocolError};
 use cait_sith::triples::{TripleGenerationOutput, TriplePub, TripleShare};
+use chrono::Utc;
 use highway::{HighwayHash, HighwayHasher};
 use k256::elliptic_curve::group::GroupEncoding;
 use k256::Secp256k1;
@@ -245,9 +246,17 @@ impl TripleManager {
         id1: TripleId,
     ) -> Result<(Triple, Triple), GenerationError> {
         if !self.triples.contains_key(&id0) {
-            Err(GenerationError::TripleIsMissing(id0))
+            if self.generators.contains_key(&id0) {
+                Err(GenerationError::TripleIsGenerating(id0))
+            } else {
+                Err(GenerationError::TripleIsMissing(id0))
+            }
         } else if !self.triples.contains_key(&id1) {
-            Err(GenerationError::TripleIsMissing(id1))
+            if self.generators.contains_key(&id1) {
+                Err(GenerationError::TripleIsGenerating(id1))
+            } else {
+                Err(GenerationError::TripleIsMissing(id1))
+            }
         } else {
             self.delete_triple_from_storage(id0).await?;
             self.delete_triple_from_storage(id1).await?;
@@ -405,6 +414,7 @@ impl TripleManager {
                                     epoch: self.epoch,
                                     from: self.me,
                                     data: data.clone(),
+                                    timestamp: Utc::now().timestamp() as u64,
                                 },
                             ))
                         }
@@ -416,6 +426,7 @@ impl TripleManager {
                             epoch: self.epoch,
                             from: self.me,
                             data,
+                            timestamp: Utc::now().timestamp() as u64,
                         },
                     )),
                     Action::Return(output) => {

node/src/util.rs

@@ -1,4 +1,5 @@
 use crate::types::PublicKey;
+use chrono::{DateTime, LocalResult, TimeZone, Utc};
 use k256::elliptic_curve::scalar::FromUintUnchecked;
 use k256::elliptic_curve::sec1::{FromEncodedPoint, ToEncodedPoint};
 use k256::{AffinePoint, EncodedPoint, Scalar, U256};
@@ -83,3 +84,16 @@ pub fn get_triple_timeout() -> Duration {
         .unwrap_or_default()
         .unwrap_or(crate::types::PROTOCOL_TRIPLE_TIMEOUT)
 }
+
+pub fn is_elapsed_longer_than_timeout(timestamp_sec: u64, timeout: Duration) -> bool {
+    if let LocalResult::Single(msg_timestamp) = Utc.timestamp_opt(timestamp_sec as i64, 0) {
+        let now_datetime: DateTime<Utc> = Utc::now();
+        // Calculate the difference in seconds
+        let elapsed_duration = now_datetime.signed_duration_since(msg_timestamp);
+        let timeout = chrono::Duration::seconds(timeout.as_secs() as i64)
+            + chrono::Duration::nanoseconds(timeout.subsec_nanos() as i64);
+        elapsed_duration > timeout
+    } else {
+        false
+    }
+}