test: more tests of wallet's handling of guesser-UTXOs

Verify that wallet-restoration handles guesser-UTXOs correctly.

src/models/blockchain/transaction/utxo.rs

@@ -129,6 +129,12 @@ impl Utxo {
         self.lock_script_hash
     }
 
+    /// Returns true iff this UTXO is a lock script with the preimage provided
+    /// as input argument.
+    pub(crate) fn is_lockscript_with_preimage(&self, preimage: Digest) -> bool {
+        self.lock_script_hash == LockScript::hash_lock_from_preimage(preimage).hash()
+    }
+
     pub fn new_native_currency(lock_script: LockScript, amount: NeptuneCoins) -> Self {
         Self::new(lock_script, vec![Coin::new_native_currency(amount)])
     }

src/models/state/mod.rs

@@ -924,6 +924,14 @@ impl GlobalState {
         let incoming_utxo_count = incoming_utxos.len();
         info!("Checking {} incoming UTXOs", incoming_utxo_count);
 
+        let existing_nonce_preimages = self
+            .wallet_state
+            .wallet_db
+            .nonce_preimages()
+            .get_all()
+            .await;
+        let mut new_nonce_preimages = vec![];
+
         let mut recovery_data_for_missing_mutxos = vec![];
         {
             // Two UTXOs are considered the same iff their AOCL index and
@@ -983,6 +991,14 @@ impl GlobalState {
                 warn!("Cannot restore UTXO with AOCL index {} because it is in the future from our tip. Current AOCL leaf count is {current_aocl_leaf_count}. Maybe this UTXO can be recovered once more blocks are downloaded from peers?", incoming_utxo.aocl_index);
                 continue;
             }
+
+            // Check if UTXO is guesser-reward and associated key doesn't already exist.
+            if incoming_utxo.is_guesser_fee() {
+                if !existing_nonce_preimages.contains(&incoming_utxo.receiver_preimage) {
+                    new_nonce_preimages.push(incoming_utxo.receiver_preimage);
+                }
+            }
+
             let ms_item = Hash::hash(&incoming_utxo.utxo);
             let restored_msmp_res = ams_ref
                 .ams()
@@ -997,7 +1013,7 @@ impl GlobalState {
                 Ok(msmp) => {
                     // Verify that the restored MSMP is valid
                     if !ams_ref.ams().verify(ms_item, &msmp).await {
-                        warn!("Restored MSMP is invalid. Skipping restoration of UTXO with AOCL index {}. Maybe this UTXO is on an abandoned chain?", incoming_utxo.aocl_index);
+                        warn!("Restored MSMP is invalid. Skipping restoration of UTXO with AOCL index {}. Maybe this UTXO is on an abandoned chain? Or maybe it was spent?", incoming_utxo.aocl_index);
                         continue;
                     }
 
@@ -1037,6 +1053,11 @@ impl GlobalState {
             restored_mutxos += 1;
         }
 
+        // Update state with all nonce-preimage keys from guesser-fee UTXOs
+        for new_nonce_preimage in new_nonce_preimages {
+            self.wallet_state.add_raw_hash_key(new_nonce_preimage).await;
+        }
+
         self.wallet_state.wallet_db.persist().await;
         info!("Successfully restored {restored_mutxos} monitored UTXOs to wallet database");
 
@@ -1682,6 +1703,7 @@ mod global_state_tests {
     use crate::models::blockchain::block::Block;
     use crate::tests::shared::fake_valid_successor_for_tests;
     use crate::tests::shared::make_mock_block;
+    use crate::tests::shared::make_mock_block_with_nonce_preimage_and_guesser_fraction;
     use crate::tests::shared::mock_genesis_global_state;
     use crate::tests::shared::wallet_state_has_all_valid_mps;
 
@@ -1872,33 +1894,83 @@ mod global_state_tests {
         let mut global_state_lock =
             mock_genesis_global_state(network, 2, wallet, cli_args::Args::default()).await;
         let genesis_block = Block::genesis_block(network);
-        let (block1, expected_utxos_block1) =
-            make_mock_block(&genesis_block, None, own_key, rng.gen()).await;
+        let nonce_preimage = rng.gen();
+        let (block1, composer_utxos) = make_mock_block_with_nonce_preimage_and_guesser_fraction(
+            &genesis_block,
+            None,
+            own_key,
+            rng.gen(),
+            0.5,
+            nonce_preimage,
+        )
+        .await;
+        let guesser_utxos = block1.guesser_fee_expected_utxos(nonce_preimage);
+        let all_mining_rewards = [composer_utxos, guesser_utxos].concat();
+
         global_state_lock
-            .lock_guard_mut()
+            .set_new_self_mined_tip(block1.clone(), all_mining_rewards)
             .await
-            .wallet_state
-            .add_expected_utxos(expected_utxos_block1)
-            .await;
-        global_state_lock.set_new_tip(block1.clone()).await.unwrap();
+            .unwrap();
 
-        // Delete everything from monitored UTXO (premined UTXO and block-1
-        // composer coinbases).
+        // Delete everything from monitored UTXO and from raw-hash keys.
         let mut global_state = global_state_lock.lock_guard_mut().await;
         {
             let monitored_utxos = global_state.wallet_state.wallet_db.monitored_utxos_mut();
             assert_eq!(
-                3,
+                5,
                 monitored_utxos.len().await,
-                "MUTXO must have genesis element and composer rewards"
+                "MUTXO must have genesis element, composer rewards, and guesser rewards"
             );
             monitored_utxos.pop().await;
             monitored_utxos.pop().await;
             monitored_utxos.pop().await;
+            monitored_utxos.pop().await;
+            monitored_utxos.pop().await;
+
+            let nonce_preimage_keys = global_state.wallet_state.wallet_db.nonce_preimages();
+            assert_eq!(
+                1,
+                nonce_preimage_keys.len().await,
+                "Exactly Nonce-preimage must be stored to DB"
+            );
+            global_state.wallet_state.clear_raw_hash_keys().await;
+            global_state
+                .wallet_state
+                .wallet_db
+                .expected_utxos_mut()
+                .clear()
+                .await;
 
             assert!(
-                monitored_utxos.is_empty().await,
-                "MUTXO must be empty after clearing"
+                global_state
+                    .wallet_state
+                    .wallet_db
+                    .monitored_utxos()
+                    .is_empty()
+                    .await
+            );
+            assert!(
+                global_state
+                    .wallet_state
+                    .wallet_db
+                    .expected_utxos()
+                    .is_empty()
+                    .await
+            );
+            assert!(
+                global_state
+                    .wallet_state
+                    .wallet_db
+                    .nonce_preimages()
+                    .is_empty()
+                    .await
+            );
+            assert_eq!(
+                0,
+                global_state
+                    .wallet_state
+                    .get_known_raw_hash_lock_keys()
+                    .count()
             );
         }
 
@@ -1911,9 +1983,9 @@ mod global_state_tests {
                 .unwrap();
             let monitored_utxos = global_state.wallet_state.wallet_db.monitored_utxos();
             assert_eq!(
-                3,
+                5,
                 monitored_utxos.len().await,
-                "MUTXO must have genesis element and premine after recovery"
+                "MUTXO must have genesis elements and premine after recovery"
             );
 
             let mutxos = monitored_utxos.get_all().await;
@@ -1924,22 +1996,23 @@ mod global_state_tests {
                     genesis_block.header().height
                 )),
                 mutxos[0].confirmed_in_block,
-                "Historical information must be restored for premine TX"
+                "Historical information must be restored for premine UTXO"
             );
 
-            for (i, mutxo) in mutxos.iter().enumerate().skip(1).take((1..=2).count()) {
+            for (i, mutxo) in mutxos.iter().enumerate().skip(1).take((1..=4).count()) {
                 assert_eq!(
                     Some((
                         block1.hash(),
                         block1.header().timestamp,
                         block1.header().height
                     )),
                     mutxo.confirmed_in_block,
-                    "Historical information must be restored for composer TX, i={i}"
+                    "Historical information must be restored for composer and guesser UTXOs, i={i}"
                 );
             }
 
-            // Verify that the restored MUTXOs have MSMPs
+            // Verify that the restored MUTXOs have MSMPs, and that they're
+            // valid.
             for mutxo in mutxos {
                 let ms_item = Hash::hash(&mutxo.utxo);
                 assert!(global_state
@@ -1956,6 +2029,30 @@ mod global_state_tests {
                     "MUTXO must have the correct latest block digest value"
                 );
             }
+
+            // Verify that guesser-fee UTXO keys have also been restored.
+            let cached_hash_lock_keys = global_state
+                .wallet_state
+                .get_known_raw_hash_lock_keys()
+                .collect_vec();
+            assert_eq!(
+                vec![SpendingKey::RawHashLock {
+                    preimage: nonce_preimage
+                }],
+                cached_hash_lock_keys,
+                "Cached hash lock keys must match expected value after recovery"
+            );
+            let persisted_hash_lock_keys = global_state
+                .wallet_state
+                .wallet_db
+                .nonce_preimages()
+                .get_all()
+                .await;
+            assert_eq!(
+                vec![nonce_preimage],
+                persisted_hash_lock_keys,
+                "Persisted hash lock keys must match expected value after recovery"
+            );
         }
     }
 

src/models/state/wallet/incoming_utxo.rs

@@ -4,7 +4,6 @@ use tasm_lib::prelude::Digest;
 
 use super::expected_utxo::UtxoNotifier;
 use super::utxo_notification::UtxoNotificationPayload;
-use crate::models::blockchain::transaction::lock_script::LockScript;
 use crate::models::blockchain::transaction::utxo::Utxo;
 use crate::models::state::ExpectedUtxo;
 use crate::models::state::Tip5;
@@ -48,9 +47,10 @@ impl IncomingUtxo {
         )
     }
 
+    /// Returns true iff this UTXO is a guesser reward.
     pub(crate) fn is_guesser_fee(&self) -> bool {
-        self.utxo.lock_script_hash()
-            == LockScript::hash_lock_from_preimage(self.receiver_preimage).hash()
+        self.utxo
+            .is_lockscript_with_preimage(self.receiver_preimage)
     }
 
     pub(crate) fn from_utxo_notification_payload(

src/models/state/wallet/wallet_state.rs

@@ -103,6 +103,12 @@ impl IncomingUtxoRecoveryData {
         let item = Tip5::hash(&self.utxo);
         commit(item, self.sender_randomness, self.receiver_preimage.hash())
     }
+
+    /// Returns true iff this UTXO is a guesser reward.
+    pub(crate) fn is_guesser_fee(&self) -> bool {
+        self.utxo
+            .is_lockscript_with_preimage(self.receiver_preimage)
+    }
 }
 
 impl TryFrom<&MonitoredUtxo> for IncomingUtxoRecoveryData {
@@ -538,6 +544,27 @@ impl WalletState {
         }
     }
 
+    /// Add a RawHashLock key to the wallet's state. If key is already stored,
+    /// this is a no-op.
+    ///
+    /// Assumes that the cache agrees with the database.
+    pub(crate) async fn add_raw_hash_key(&mut self, nonce_preimage: Digest) {
+        let as_key = SpendingKey::RawHashLock {
+            preimage: nonce_preimage,
+        };
+        if self.known_raw_hash_lock_keys.contains(&as_key) {
+            return;
+        }
+
+        self.wallet_db
+            .nonce_preimages_mut()
+            .push(nonce_preimage)
+            .await;
+        self.known_raw_hash_lock_keys.push(as_key);
+
+        return;
+    }
+
     /// Return a list of UTXOs spent by this wallet in the transaction
     async fn scan_for_spent_utxos(
         &self,
@@ -1215,14 +1242,7 @@ impl WalletState {
 
         // Write nonce-preimage for guesser fee UTXOs to DB, and cache.
         if let Some(nonce_preimage) = nonce_preimage {
-            self.wallet_db
-                .nonce_preimages_mut()
-                .push(nonce_preimage)
-                .await;
-            self.known_raw_hash_lock_keys
-                .push(SpendingKey::RawHashLock {
-                    preimage: nonce_preimage,
-                });
+            self.add_raw_hash_key(nonce_preimage).await;
         }
 
         // Mark all expected UTXOs that were received in this block as received
@@ -1444,6 +1464,14 @@ mod tests {
     use crate::tests::shared::mock_genesis_wallet_state;
     use crate::tests::shared::wallet_state_has_all_valid_mps;
 
+    impl WalletState {
+        /// Delete all nonce-preimage keys from database and cache.
+        pub(crate) async fn clear_raw_hash_keys(&mut self) {
+            self.known_raw_hash_lock_keys.clear();
+            self.wallet_db.nonce_preimages_mut().clear().await;
+        }
+    }
+
     #[tokio::test]
     #[traced_test]
     async fn find_monitored_utxo_test() {

src/tests/shared.rs

@@ -727,18 +727,15 @@ pub(crate) async fn make_mock_block(
     seed: [u8; 32],
 ) -> (Block, Vec<ExpectedUtxo>) {
     let nonce_preimage = Digest::default();
-    let (block, composer_expected_utxos) =
-        make_mock_block_with_nonce_preimage_and_guesser_fraction(
-            previous_block,
-            block_timestamp,
-            composer_key,
-            seed,
-            0f64,
-            nonce_preimage,
-        )
-        .await;
-
-    (block, composer_expected_utxos)
+    make_mock_block_with_nonce_preimage_and_guesser_fraction(
+        previous_block,
+        block_timestamp,
+        composer_key,
+        seed,
+        0f64,
+        nonce_preimage,
+    )
+    .await
 }
 
 /// Return a dummy-wallet used for testing. The returned wallet is populated with