Make sure triples and presignatures are not reused

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

@@ -187,6 +187,17 @@ impl PresignatureManager {
 
     pub async fn insert(&mut self, presignature: Presignature) {
         tracing::debug!(id = ?presignature.id, "inserting presignature");
+        if self.contains(&presignature.id).await {
+            tracing::error!(id = presignature.id, "mine presignature already inserted");
+            return;
+        }
+        if self.contains_used(&presignature.id).await {
+            tracing::error!(
+                id = presignature.id,
+                "tried to insert used mine presignature"
+            );
+            return;
+        }
         // Remove from taken list if it was there
         self.gc.remove(&presignature.id);
         if let Err(e) = self.presignature_storage.insert(presignature).await {
@@ -196,13 +207,31 @@ impl PresignatureManager {
 
     pub async fn insert_mine(&mut self, presignature: Presignature) {
         tracing::debug!(id = ?presignature.id, "inserting mine presignature");
+        if self.contains(&presignature.id).await {
+            tracing::error!(id = presignature.id, "mine presignature already inserted");
+            return;
+        }
+        if self.contains_used(&presignature.id).await {
+            tracing::error!(
+                id = presignature.id,
+                "tried to insert used mine presignature"
+            );
+            return;
+        }
         // Remove from taken list if it was there
         self.gc.remove(&presignature.id);
         if let Err(e) = self.presignature_storage.insert_mine(presignature).await {
             tracing::error!(?e, "failed to insert mine presignature");
         }
     }
 
+    pub async fn insert_used(&mut self, id: PresignatureId) {
+        tracing::debug!(id, "inserting presignature to used");
+        if let Err(e) = self.presignature_storage.insert_used(id).await {
+            tracing::error!(?e, "failed to insert presignature to used");
+        }
+    }
+
     /// Returns true if the presignature with the given id is already generated
     pub async fn contains(&self, id: &PresignatureId) -> bool {
         self.presignature_storage
@@ -225,11 +254,20 @@ impl PresignatureManager {
             .unwrap_or(false)
     }
 
+    pub async fn contains_used(&self, id: &PresignatureId) -> bool {
+        self.presignature_storage
+            .contains_used(id)
+            .await
+            .map_err(|e| tracing::warn!(?e, "failed to check if presignature was used"))
+            .unwrap_or(false)
+    }
+
     pub async fn take(&mut self, id: PresignatureId) -> Result<Presignature, GenerationError> {
         if let Some(presignature) = self.presignature_storage.take(&id).await.map_err(|e| {
             tracing::error!(?e, "failed to look for presignature");
             GenerationError::PresignatureIsMissing(id)
         })? {
+            self.insert_used(presignature.id).await;
             self.gc.insert(id, Instant::now());
             tracing::debug!(id, "took presignature");
             return Ok(presignature);
@@ -257,6 +295,7 @@ impl PresignatureManager {
             })
             .ok()?
         {
+            self.insert_used(presignature.id).await;
             tracing::debug!(id = ?presignature.id, "took presignature of mine");
             return Some(presignature);
         }
@@ -285,6 +324,16 @@ impl PresignatureManager {
             .unwrap_or(0)
     }
 
+    pub async fn len_used(&self) -> usize {
+        self.presignature_storage
+            .len_used()
+            .await
+            .map_err(|e| {
+                tracing::error!(?e, "failed to count used presignatures");
+            })
+            .unwrap_or(0)
+    }
+
     /// Returns if there are unspent presignatures available in the manager.
     pub async fn is_empty(&self) -> bool {
         self.len_generated().await == 0
@@ -433,11 +482,6 @@ impl PresignatureManager {
                         participants = ?presig_participants.keys_vec(),
                         "running: the intersection of participants is less than the threshold"
                     );
-
-                    // Insert back the triples to be used later since this active set of
-                    // participants were not able to make use of these triples.
-                    triple_manager.insert_mine(triple0).await;
-                    triple_manager.insert_mine(triple1).await;
                 } else {
                     self.generate(
                         &presig_participants,

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

@@ -626,7 +626,6 @@ impl SignatureManager {
             );
             return;
         }
-        let mut failed_presigs = Vec::new();
         while let Some(mut presignature) = {
             if self.failed.is_empty() && my_requests.is_empty() {
                 None
@@ -640,7 +639,6 @@ impl SignatureManager {
                     participants = ?sig_participants.keys_vec(),
                     "intersection of stable participants and presignature participants is less than threshold"
                 );
-                failed_presigs.push(presignature);
                 continue;
             }
             let presig_id = presignature.id;
@@ -650,7 +648,7 @@ impl SignatureManager {
             // when the request made it into the NEAR network.
             // issue: https://github.com/near/mpc-recovery/issues/596
             if let Some((sign_request_identifier, failed_req)) = self.failed.pop_front() {
-                if let Err((presignature, InitializationError::BadParameters(err))) = self
+                if let Err((_, InitializationError::BadParameters(err))) = self
                     .retry_failed_generation(
                         sign_request_identifier.clone(),
                         failed_req,
@@ -665,7 +663,6 @@ impl SignatureManager {
                         ?err,
                         "failed to retry signature generation: trashing presignature"
                     );
-                    failed_presigs.push(presignature);
                     continue;
                 }
 
@@ -677,11 +674,10 @@ impl SignatureManager {
             }
 
             let Some(my_request) = my_requests.pop_front() else {
-                failed_presigs.push(presignature);
                 continue;
             };
 
-            if let Err((presignature, InitializationError::BadParameters(err))) = self.generate(
+            if let Err((_, InitializationError::BadParameters(err))) = self.generate(
                 &sig_participants,
                 my_request.request_id,
                 presignature,
@@ -691,17 +687,10 @@ impl SignatureManager {
                 my_request.time_added,
                 cfg,
             ) {
-                failed_presigs.push(presignature);
                 tracing::warn!(request_id = ?CryptoHash(my_request.request_id), presig_id, ?err, "failed to start signature generation: trashing presignature");
                 continue;
             }
         }
-
-        // add back the failed presignatures that were incompatible to be made into
-        // signatures due to failures or lack of participants.
-        for presignature in failed_presigs {
-            presignature_manager.insert_mine(presignature).await;
-        }
     }
 
     pub async fn publish<T: SignerExt>(

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

@@ -146,6 +146,14 @@ impl TripleManager {
 
     pub async fn insert(&mut self, triple: Triple) {
         tracing::debug!(id = triple.id, "inserting triple");
+        if self.contains(&triple.id).await {
+            tracing::error!(id = triple.id, "triple already inserted");
+            return;
+        }
+        if self.contains_used(&triple.id).await {
+            tracing::error!(id = triple.id, "tried to insert used triple");
+            return;
+        }
         self.gc.remove(&triple.id);
         if let Err(e) = self.triple_storage.insert(triple).await {
             tracing::warn!(?e, "failed to insert triple");
@@ -154,12 +162,27 @@ impl TripleManager {
 
     pub async fn insert_mine(&mut self, triple: Triple) {
         tracing::debug!(id = triple.id, "inserting mine triple");
+        if self.contains(&triple.id).await {
+            tracing::error!(id = triple.id, "mine triple already inserted");
+            return;
+        }
+        if self.contains_used(&triple.id).await {
+            tracing::error!(id = triple.id, "tried to insert used mine triple");
+            return;
+        }
         self.gc.remove(&triple.id);
         if let Err(e) = self.triple_storage.insert_mine(triple).await {
             tracing::warn!(?e, "failed to insert mine triple");
         }
     }
 
+    pub async fn insert_used(&mut self, id: TripleId) {
+        tracing::debug!(id, "inserting triple to used");
+        if let Err(e) = self.triple_storage.insert_used(id).await {
+            tracing::warn!(?e, "failed to insert tripel to used");
+        }
+    }
+
     pub async fn contains(&self, id: &TripleId) -> bool {
         self.triple_storage
             .contains(id)
@@ -176,6 +199,14 @@ impl TripleManager {
             .unwrap_or(false)
     }
 
+    pub async fn contains_used(&self, id: &TripleId) -> bool {
+        self.triple_storage
+            .contains_used(id)
+            .await
+            .map_err(|e| tracing::warn!(?e, "failed to check if triple was used"))
+            .unwrap_or(false)
+    }
+
     /// Take two unspent triple by theirs id with no way to return it. Only takes
     /// if both of them are present.
     /// It is very important to NOT reuse the same triple twice for two different
@@ -232,6 +263,9 @@ impl TripleManager {
             }
         };
 
+        self.insert_used(triple_0.id).await;
+        self.insert_used(triple_1.id).await;
+
         self.gc.insert(id0, Instant::now());
         self.gc.insert(id1, Instant::now());
 
@@ -279,6 +313,9 @@ impl TripleManager {
             }
         };
 
+        self.insert_used(triple_0.id).await;
+        self.insert_used(triple_1.id).await;
+
         self.gc.insert(triple_0.id, Instant::now());
         self.gc.insert(triple_1.id, Instant::now());
 
@@ -297,6 +334,11 @@ impl TripleManager {
         self.triple_storage.len_mine().await.unwrap_or(0)
     }
 
+    /// Returns the number of used triples
+    pub async fn len_used(&self) -> usize {
+        self.triple_storage.len_used().await.unwrap_or(0)
+    }
+
     /// Returns if there's any unspent triple in the manager.
     pub async fn is_empty(&self) -> bool {
         self.len_generated().await == 0

chain-signatures/node/src/storage/presignature_storage.rs

@@ -8,7 +8,7 @@ use crate::protocol::presignature::{Presignature, PresignatureId};
 type PresigResult<T> = std::result::Result<T, anyhow::Error>;
 
 // Can be used to "clear" redis storage in case of a breaking change
-const PRESIGNATURE_STORAGE_VERSION: &str = "v1";
+const PRESIGNATURE_STORAGE_VERSION: &str = "v2";
 
 pub fn init(pool: &Pool, node_account_id: &AccountId) -> PresignatureRedisStorage {
     PresignatureRedisStorage {
@@ -45,6 +45,14 @@ impl PresignatureRedisStorage {
         Ok(())
     }
 
+    pub async fn insert_used(&self, id: PresignatureId) -> PresigResult<()> {
+        let mut connection = self.redis_pool.get().await?;
+        connection
+            .sadd::<&str, PresignatureId, ()>(&self.used_key(), id)
+            .await?;
+        Ok(())
+    }
+
     pub async fn contains(&self, id: &PresignatureId) -> PresigResult<bool> {
         let mut connection = self.redis_pool.get().await?;
         let result: bool = connection.hexists(self.presig_key(), id).await?;
@@ -57,6 +65,12 @@ impl PresignatureRedisStorage {
         Ok(result)
     }
 
+    pub async fn contains_used(&self, id: &PresignatureId) -> PresigResult<bool> {
+        let mut connection = self.redis_pool.get().await?;
+        let result: bool = connection.sismember(self.used_key(), id).await?;
+        Ok(result)
+    }
+
     pub async fn take(&self, id: &PresignatureId) -> PresigResult<Option<Presignature>> {
         let mut connection = self.redis_pool.get().await?;
         if self.contains_mine(id).await? {
@@ -96,10 +110,17 @@ impl PresignatureRedisStorage {
         Ok(result)
     }
 
+    pub async fn len_used(&self) -> PresigResult<usize> {
+        let mut connection = self.redis_pool.get().await?;
+        let result: usize = connection.scard(self.used_key()).await?;
+        Ok(result)
+    }
+
     pub async fn clear(&self) -> PresigResult<()> {
         let mut connection = self.redis_pool.get().await?;
         connection.del::<&str, ()>(&self.presig_key()).await?;
         connection.del::<&str, ()>(&self.mine_key()).await?;
+        connection.del::<&str, ()>(&self.used_key()).await?;
         Ok(())
     }
 
@@ -116,6 +137,13 @@ impl PresignatureRedisStorage {
             PRESIGNATURE_STORAGE_VERSION, self.node_account_id
         )
     }
+
+    fn used_key(&self) -> String {
+        format!(
+            "presignatures_used:{}:{}",
+            PRESIGNATURE_STORAGE_VERSION, self.node_account_id
+        )
+    }
 }
 
 impl ToRedisArgs for Presignature {

chain-signatures/node/src/storage/triple_storage.rs

@@ -8,7 +8,7 @@ use near_account_id::AccountId;
 type TripleResult<T> = std::result::Result<T, anyhow::Error>;
 
 // Can be used to "clear" redis storage in case of a breaking change
-const TRIPLE_STORAGE_VERSION: &str = "v1";
+const TRIPLE_STORAGE_VERSION: &str = "v2";
 
 pub fn init(pool: &Pool, account_id: &AccountId) -> TripleRedisStorage {
     TripleRedisStorage {
@@ -39,6 +39,13 @@ impl TripleRedisStorage {
         Ok(())
     }
 
+    pub async fn insert_used(&self, id: TripleId) -> TripleResult<()> {
+        let mut conn = self.redis_pool.get().await?;
+        conn.sadd::<&str, TripleId, ()>(&self.used_key(), id)
+            .await?;
+        Ok(())
+    }
+
     pub async fn contains(&self, id: &TripleId) -> TripleResult<bool> {
         let mut conn = self.redis_pool.get().await?;
         let result: bool = conn.hexists(self.triple_key(), id).await?;
@@ -51,6 +58,12 @@ impl TripleRedisStorage {
         Ok(result)
     }
 
+    pub async fn contains_used(&self, id: &TripleId) -> TripleResult<bool> {
+        let mut conn = self.redis_pool.get().await?;
+        let result: bool = conn.sismember(self.used_key(), id).await?;
+        Ok(result)
+    }
+
     pub async fn take(&self, id: &TripleId) -> TripleResult<Option<Triple>> {
         let mut conn = self.redis_pool.get().await?;
         if self.contains_mine(id).await? {
@@ -89,10 +102,17 @@ impl TripleRedisStorage {
         Ok(result)
     }
 
+    pub async fn len_used(&self) -> TripleResult<usize> {
+        let mut conn = self.redis_pool.get().await?;
+        let result: usize = conn.scard(self.used_key()).await?;
+        Ok(result)
+    }
+
     pub async fn clear(&self) -> TripleResult<()> {
         let mut conn = self.redis_pool.get().await?;
         conn.del::<&str, ()>(&self.triple_key()).await?;
         conn.del::<&str, ()>(&self.mine_key()).await?;
+        conn.del::<&str, ()>(&self.used_key()).await?;
         Ok(())
     }
 
@@ -109,6 +129,13 @@ impl TripleRedisStorage {
             TRIPLE_STORAGE_VERSION, self.node_account_id
         )
     }
+
+    fn used_key(&self) -> String {
+        format!(
+            "triples_used:{}:{}",
+            TRIPLE_STORAGE_VERSION, self.node_account_id
+        )
+    }
 }
 
 impl ToRedisArgs for Triple {