claim_utxo: Handle scenario where UTXO is already spent

When a UTXO is claimed, check if it is already spent and mark it as
such.

src/models/state/archival_state.rs

@@ -33,6 +33,7 @@ use crate::models::database::LastFileRecord;
 use crate::prelude::twenty_first;
 use crate::util_types::mutator_set::addition_record::AdditionRecord;
 use crate::util_types::mutator_set::mutator_set_accumulator::MutatorSetAccumulator;
+use crate::util_types::mutator_set::removal_record::AbsoluteIndexSet;
 use crate::util_types::mutator_set::removal_record::RemovalRecord;
 use crate::util_types::mutator_set::rusty_archival_mutator_set::RustyArchivalMutatorSet;
 
@@ -439,10 +440,11 @@ impl ArchivalState {
     }
 
     /// searches max `max_search_depth` from tip for a matching transaction
-    /// output.
+    /// input.
     ///
-    /// If `max_search_depth` is set to `None`, then all blocks are searched. A
-    /// A `max_search_depth` of `Some(0)` will only consider the tip.
+    /// If `max_search_depth` is set to `None`, then all blocks are searched
+    /// until a match is found. A `max_search_depth` of `Some(0)` will only
+    /// consider the tip.
     pub(crate) async fn find_canonical_block_with_output(
         &self,
         output: AdditionRecord,
@@ -475,6 +477,44 @@ impl ArchivalState {
         }
     }
 
+    /// searches max `max_search_depth` from tip for a matching transaction
+    /// input.
+    ///
+    /// If `max_search_depth` is set to `None`, then all blocks are searched
+    /// until a match is found. A `max_search_depth` of `Some(0)` will only
+    /// consider the tip.
+    pub(crate) async fn find_canonical_block_with_input(
+        &self,
+        input: AbsoluteIndexSet,
+        max_search_depth: Option<u64>,
+    ) -> Option<Block> {
+        let mut block = self.get_tip().await;
+        let mut search_depth = 0;
+
+        loop {
+            if block
+                .body()
+                .transaction_kernel
+                .inputs
+                .iter()
+                .any(|rr| rr.absolute_indices == input)
+            {
+                break Some(block);
+            }
+
+            if max_search_depth.is_some_and(|max| max <= search_depth) {
+                return None;
+            }
+
+            block = self
+                .get_block(block.header().prev_block_digest)
+                .await
+                .ok()??;
+
+            search_depth += 1;
+        }
+    }
+
     /// Return latest block from database, or genesis block if no other block
     /// is known.
     pub async fn get_tip(&self) -> Block {
@@ -1053,6 +1093,7 @@ mod archival_state_tests {
     use crate::tests::shared::unit_test_databases;
     use crate::util_types::mutator_set::addition_record::AdditionRecord;
     use crate::util_types::test_shared::mutator_set::mock_item_mp_rr_for_init_msa;
+    use crate::util_types::test_shared::mutator_set::random_removal_record;
 
     async fn make_test_archival_state(network: Network) -> ArchivalState {
         let (block_index_db, _peer_db_lock, data_dir) = unit_test_databases(network).await.unwrap();
@@ -2315,6 +2356,41 @@ mod archival_state_tests {
         }
     }
 
+    #[traced_test]
+    #[tokio::test]
+    async fn find_canonical_block_with_output_genesis_block_test() {
+        let network = Network::Main;
+        let archival_state = make_test_archival_state(network).await;
+        let genesis_block = Block::genesis_block(network);
+
+        let addition_records = Block::genesis_block(network)
+            .body()
+            .transaction_kernel
+            .outputs
+            .clone();
+
+        for ar in addition_records.iter() {
+            let found_block = archival_state
+                .find_canonical_block_with_output(*ar, None)
+                .await
+                .unwrap();
+            assert_eq!(genesis_block.hash(), found_block.hash());
+        }
+    }
+
+    #[traced_test]
+    #[tokio::test]
+    async fn find_canonical_block_with_input_genesis_block_test() {
+        let network = Network::Main;
+        let archival_state = make_test_archival_state(network).await;
+        let random_index_set: AbsoluteIndexSet = random_removal_record().absolute_indices;
+
+        assert!(archival_state
+            .find_canonical_block_with_input(random_index_set, None)
+            .await
+            .is_none());
+    }
+
     #[traced_test]
     #[tokio::test]
     async fn ms_update_to_tip_fork_depth_1() {

src/models/state/wallet/monitored_utxo.rs

@@ -5,6 +5,7 @@ use serde::Serialize;
 use twenty_first::math::tip5::Digest;
 
 use crate::models::blockchain::block::block_height::BlockHeight;
+use crate::models::blockchain::block::Block;
 use crate::models::blockchain::transaction::utxo::Utxo;
 use crate::models::proof_abstractions::timestamp::Timestamp;
 use crate::models::state::archival_state::ArchivalState;
@@ -71,6 +72,14 @@ impl MonitoredUtxo {
             .map(|x| x.1.clone())
     }
 
+    pub(crate) fn mark_as_spent(&mut self, spending_block: &Block) {
+        self.spent_in_block = Some((
+            spending_block.hash(),
+            spending_block.kernel.header.timestamp,
+            spending_block.kernel.header.height,
+        ));
+    }
+
     /// Get the most recent (block hash, membership proof) entry in the database,
     /// if any.
     pub fn get_latest_membership_proof_entry(&self) -> Option<(Digest, MsMembershipProof)> {

src/models/state/wallet/wallet_state.rs

@@ -1059,11 +1059,7 @@ impl WalletState {
                     );
 
                     let mut spent_mutxo = monitored_utxos.get(*mutxo_list_index).await;
-                    spent_mutxo.spent_in_block = Some((
-                        new_block.hash(),
-                        new_block.kernel.header.timestamp,
-                        new_block.kernel.header.height,
-                    ));
+                    spent_mutxo.mark_as_spent(new_block);
                     monitored_utxos.set(*mutxo_list_index, spent_mutxo).await;
                 }
             }

src/rpc_server.rs

@@ -685,7 +685,32 @@ impl NeptuneRPCServer {
                     block.header().timestamp,
                     block.header().height,
                 ));
-                monitored_utxo.add_membership_proof_for_tip(tip_digest, msmp);
+                monitored_utxo.add_membership_proof_for_tip(tip_digest, msmp.clone());
+
+                // Was UTXO already spent? If so, register it as such.
+                let msa = ams.accumulator().await;
+                if !msa.verify(item, &msmp) {
+                    warn!("Claimed UTXO was already spent. Marking it as such.");
+
+                    if let Some(spending_block) = state
+                        .chain
+                        .archival_state()
+                        .find_canonical_block_with_input(
+                            msmp.compute_indices(item),
+                            max_search_depth,
+                        )
+                        .await
+                    {
+                        warn!(
+                            "Claimed UTXO was spent in block {}; which has height {}",
+                            spending_block.hash(),
+                            spending_block.header().height
+                        );
+                        monitored_utxo.mark_as_spent(&spending_block);
+                    } else {
+                        error!("Claimed UTXO's mutator set membership proof was invalid but we could not find the block in which it was spent. This is most likely a bug in the software.");
+                    }
+                }
 
                 Some(monitored_utxo)
             }
@@ -2275,14 +2300,21 @@ mod rpc_server_tests {
         #[allow(clippy::needless_return)]
         #[tokio::test]
         async fn claim_utxo_owned_before_confirmed() -> Result<()> {
-            worker::claim_utxo_owned(false).await
+            worker::claim_utxo_owned(false, false).await
         }
 
         #[traced_test]
         #[allow(clippy::needless_return)]
         #[tokio::test]
         async fn claim_utxo_owned_after_confirmed() -> Result<()> {
-            worker::claim_utxo_owned(true).await
+            worker::claim_utxo_owned(true, false).await
+        }
+
+        #[traced_test]
+        #[allow(clippy::needless_return)]
+        #[tokio::test]
+        async fn claim_utxo_owned_after_confirmed_and_after_spent() -> Result<()> {
+            worker::claim_utxo_owned(true, true).await
         }
 
         #[traced_test]
@@ -2449,7 +2481,14 @@ mod rpc_server_tests {
                 Ok(())
             }
 
-            pub(super) async fn claim_utxo_owned(claim_after_confirmed: bool) -> Result<()> {
+            pub(super) async fn claim_utxo_owned(
+                claim_after_mined: bool,
+                spent: bool,
+            ) -> Result<()> {
+                assert!(
+                    !spent || claim_after_mined,
+                    "If UTXO is spent, it must also be mined"
+                );
                 let network = Network::Main;
                 let bob_key = WalletSecret::new_random();
                 let mut bob_rpc_server =
@@ -2499,7 +2538,7 @@ mod rpc_server_tests {
                 let (tx, offchain_notifications) = bob_rpc_server
                     .clone()
                     .send_to_many_inner_invalid_proof(
-                        pay_to_self_outputs,
+                        pay_to_self_outputs.clone(),
                         UtxoNotificationMedium::OffChain,
                         UtxoNotificationMedium::OffChain,
                         fee,
@@ -2512,10 +2551,30 @@ mod rpc_server_tests {
                 let block2 = invalid_block_with_transaction(&block1, tx);
                 let block3 = invalid_empty_block(&block2);
 
-                if claim_after_confirmed {
+                if claim_after_mined {
                     // bob applies the blocks before claiming utxos.
                     bob_rpc_server.state.set_new_tip(block2.clone()).await?;
                     bob_rpc_server.state.set_new_tip(block3.clone()).await?;
+
+                    if spent {
+                        // Send entire balance somewhere else
+                        let another_address = WalletSecret::new_random()
+                            .nth_generation_spending_key(0)
+                            .to_address();
+                        let (spending_tx, _) = bob_rpc_server
+                            .clone()
+                            .send_to_many_inner_invalid_proof(
+                                vec![(another_address.into(), NeptuneCoins::new(126))],
+                                UtxoNotificationMedium::OffChain,
+                                UtxoNotificationMedium::OffChain,
+                                NeptuneCoins::zero(),
+                                in_eight_months,
+                            )
+                            .await
+                            .unwrap();
+                        let block4 = invalid_block_with_transaction(&block3, spending_tx);
+                        bob_rpc_server.state.set_new_tip(block4.clone()).await?;
+                    }
                 }
 
                 for offchain_notification in offchain_notifications {
@@ -2547,7 +2606,7 @@ mod rpc_server_tests {
                         .collect_vec()
                 );
 
-                if !claim_after_confirmed {
+                if !claim_after_mined {
                     // bob hasn't applied blocks 2,3. balance should be 128
                     assert_eq!(
                         NeptuneCoins::new(128),
@@ -2561,29 +2620,23 @@ mod rpc_server_tests {
                     bob_rpc_server.state.set_new_tip(block3).await?;
                 }
 
-                // final balance should be 126.
-                // +128  coinbase
-                // -128  coinbase spent
-                // +5 self-send via Generation
-                // +6 self-send via Symmetric
-                // +115   change (less fee == 2)
-                assert_eq!(
-                    NeptuneCoins::new(126),
-                    bob_rpc_server.synced_balance(context::current()).await,
-                );
-
-                // todo: test that claim_utxo() correctly handles case when the
-                //       claimed utxo has already been spent.
-                //
-                //       in normal wallet usage this would not happen.  However it
-                //       is possible if bob were to claim a utxo with wallet A,
-                //       spend the utxo and then restore wallet B from A's seed.
-                //       When bob performs claim_utxo() in wallet B the balance
-                //       should reflect that the utxo was already spent.
-                //
-                //       this is a bit tricky to test, as it requires using a
-                //       different data directory for wallet B and test infrastructure
-                //       isn't setup for that.
+                if spent {
+                    assert!(bob_rpc_server
+                        .synced_balance(context::current())
+                        .await
+                        .is_zero(),);
+                } else {
+                    // final balance should be 126.
+                    // +128  coinbase
+                    // -128  coinbase spent
+                    // +5 self-send via Generation
+                    // +6 self-send via Symmetric
+                    // +115   change (less fee == 2)
+                    assert_eq!(
+                        NeptuneCoins::new(126),
+                        bob_rpc_server.synced_balance(context::current()).await,
+                    );
+                }
                 Ok(())
             }
         }