Просмотр исходного кода

drk: handle wallet freezes properly

skoupidi 1 год назад
Родитель
Сommit
e3e2850039
2 измененных файлов с 35 добавлено и 7 удалено
  1. 26 7
      bin/drk/src/money.rs
  2. 9 0
      bin/drk/src/rpc.rs

+ 26 - 7
bin/drk/src/money.rs

@@ -820,10 +820,29 @@ impl Drk {
 
     /// Auxiliary function to handle freezes from a transaction money
     /// call.
-    async fn handle_money_call_freezes(&self, freezes: &[TokenId]) -> Result<()> {
+    /// Returns a flag indicating if provided freezes refer to our own
+    /// wallet.
+    async fn handle_money_call_freezes(
+        &self,
+        own_tokens: &[TokenId],
+        freezes: &[TokenId],
+    ) -> Result<bool> {
         // Check if we have any freezes to process
         if freezes.is_empty() {
-            return Ok(())
+            return Ok(false)
+        }
+
+        // Find our own tokens that got frozen
+        let mut own_freezes = Vec::with_capacity(freezes.len());
+        for freeze in freezes {
+            if own_tokens.contains(freeze) {
+                own_freezes.push(freeze);
+            }
+        }
+
+        // Check if we need to freeze anything
+        if own_freezes.is_empty() {
+            return Ok(false)
         }
 
         // This is the SQL query we'll be executing to update frozen tokens into the wallet
@@ -832,7 +851,7 @@ impl Drk {
             *MONEY_TOKENS_TABLE, MONEY_TOKENS_COL_IS_FROZEN, MONEY_TOKENS_COL_TOKEN_ID,
         );
 
-        for token_id in freezes {
+        for token_id in own_freezes {
             // Grab token record key
             let key = serialize_async(token_id).await;
 
@@ -844,7 +863,7 @@ impl Drk {
             }
         }
 
-        Ok(())
+        Ok(true)
     }
 
     /// Append data related to Money contract transactions into the
@@ -879,14 +898,14 @@ impl Drk {
         self.handle_money_call_owncoins(&mut scan_cache.owncoins_nullifiers, &owncoins).await?;
 
         // Handle freezes
-        // TODO: this should return flag if we have frozen tokens indeed
-        self.handle_money_call_freezes(&freezes).await?;
+        let wallet_freezes =
+            self.handle_money_call_freezes(&scan_cache.own_tokens, &freezes).await?;
 
         if self.fun && !owncoins.is_empty() {
             kaching().await;
         }
 
-        Ok((update_tree, wallet_spent_coins || !owncoins.is_empty() || !freezes.is_empty()))
+        Ok((update_tree, wallet_spent_coins || !owncoins.is_empty() || wallet_freezes))
     }
 
     /// Auxiliary function to  grab all the nullifiers from a transaction money call.

+ 9 - 0
bin/drk/src/rpc.rs

@@ -37,6 +37,7 @@ use darkfi::{
     Error, Result,
 };
 use darkfi_dao_contract::model::{DaoBulla, DaoProposalBulla};
+use darkfi_money_contract::model::TokenId;
 use darkfi_sdk::{
     crypto::{
         smt::{PoseidonFp, EMPTY_NODES_FP},
@@ -65,6 +66,8 @@ pub struct ScanCache {
     pub notes_secrets: Vec<SecretKey>,
     /// Our own coins nullifiers
     pub owncoins_nullifiers: BTreeMap<[u8; 32], [u8; 32]>,
+    /// Our own tokens to track freezes
+    pub own_tokens: Vec<TokenId>,
     /// The DAO Merkle tree containing DAO bullas
     pub dao_daos_tree: MerkleTree,
     /// The DAO Merkle tree containing proposals bullas
@@ -87,6 +90,11 @@ impl Drk {
         for coin in self.get_coins(true).await? {
             owncoins_nullifiers.insert(coin.0.nullifier().to_bytes(), coin.0.coin.to_bytes());
         }
+        let mint_authorities = self.get_mint_authorities().await?;
+        let mut own_tokens = Vec::with_capacity(mint_authorities.len());
+        for (token, _, _, _) in mint_authorities {
+            own_tokens.push(token);
+        }
         let (dao_daos_tree, dao_proposals_tree) = self.get_dao_trees().await?;
         let mut own_daos = HashMap::new();
         for dao in self.get_daos().await? {
@@ -108,6 +116,7 @@ impl Drk {
             money_smt,
             notes_secrets,
             owncoins_nullifiers,
+            own_tokens,
             dao_daos_tree,
             dao_proposals_tree,
             own_daos,