,ake sure triples are not reused

near · Nov 19, 2024 · b530b43 · b530b43
1 parent b69ce2f
commit b530b43
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 1 deletion.
diff --git a/chain-signatures/node/src/protocol/triple.rs b/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());
 

diff --git a/chain-signatures/node/src/storage/triple_storage.rs b/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 {