feat: improve network royalties identification for spends

maidsafe · Jul 22, 2024 · 5b6d160 · 5b6d160
1 parent bcf05e5
commit 5b6d160
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 31 deletions.
diff --git a/sn_client/src/audit/dag_crawling.rs b/sn_client/src/audit/dag_crawling.rs
@@ -11,8 +11,8 @@ use crate::{Client, Error, SpendDag};
 use futures::{future::join_all, StreamExt};
 use sn_networking::{GetRecordError, NetworkError};
 use sn_transfers::{
-    NanoTokens, OutputPurpose, SignedSpend, SpendAddress, SpendReason, UniquePubkey, WalletError,
-    WalletResult, DEFAULT_NETWORK_ROYALTIES_PK, GENESIS_SPEND_UNIQUE_KEY, NETWORK_ROYALTIES_PK,
+    NanoTokens, SignedSpend, SpendAddress, SpendReason, UniquePubkey, WalletError, WalletResult,
+    DEFAULT_NETWORK_ROYALTIES_PK, GENESIS_SPEND_UNIQUE_KEY, NETWORK_ROYALTIES_PK,
 };
 use std::{
     collections::{BTreeMap, BTreeSet},
@@ -174,14 +174,7 @@ impl Client {
                             let amount = spend.spend.amount();
                             let ancestors_len = spend.spend.ancestors.len();
                             let descendants_len = spend.spend.descendants.len();
-                            let roy_len = spend
-                                .spend
-                                .descendants
-                                .iter()
-                                .filter(|(_pk, (_amount, _purpose))| {
-                                    matches!(OutputPurpose::RoyaltyFee, _purpose)
-                                })
-                                .count();
+                            let roy_len = spend.spend.network_royalties().len();
                             warn!(
                                 "double spend entry {i} reason {reason:?}, amount {amount}, ancestors: {ancestors_len}, descendants: {descendants_len}, royalties: {roy_len}, {:?} - {:?}",
                                 spend.spend.ancestors, spend.spend.descendants
@@ -552,27 +545,15 @@ fn beta_track_analyze_spend(spend: &SignedSpend) -> BTreeSet<(SpendAddress, Nano
     // Filter out royalty outputs
     let royalty_pubkeys: BTreeSet<_> = spend
         .spend
-        .descendants
+        .network_royalties()
         .iter()
-        .filter_map(|(_pk, (_amount, purpose))| {
-            if let OutputPurpose::RoyaltyFee(der) = purpose {
-                Some(NETWORK_ROYALTIES_PK.new_unique_pubkey(der))
-            } else {
-                None
-            }
-        })
+        .map(|(_, _, der)| NETWORK_ROYALTIES_PK.new_unique_pubkey(der))
         .collect();
     let default_royalty_pubkeys: BTreeSet<_> = spend
         .spend
-        .descendants
+        .network_royalties()
         .iter()
-        .filter_map(|(_pk, (_amount, purpose))| {
-            if let OutputPurpose::RoyaltyFee(der) = purpose {
-                Some(DEFAULT_NETWORK_ROYALTIES_PK.new_unique_pubkey(der))
-            } else {
-                None
-            }
-        })
+        .map(|(_, _, der)| DEFAULT_NETWORK_ROYALTIES_PK.new_unique_pubkey(der))
         .collect();
 
     let new_utxos: BTreeSet<_> = spend

diff --git a/sn_client/src/audit/tests/mod.rs b/sn_client/src/audit/tests/mod.rs
@@ -46,7 +46,6 @@ fn test_spend_dag_verify_valid_simple() -> Result<()> {
     Ok(())
 }
 
-#[ignore = "Re-enabled once got DAG auditor properly updated"]
 #[test]
 fn test_spend_dag_double_spend_poisonning() -> Result<()> {
     let mut net = MockNetwork::genesis()?;
@@ -116,7 +115,7 @@ fn test_spend_dag_double_spend_poisonning() -> Result<()> {
     let expected = BTreeSet::from_iter([SpendFault::DoubleSpend(*double_spent)]);
     assert_eq!(got, expected, "spend_ko3 should be unspendable");
 
-    // make sure the effects of the rest of the DAG
+    // make sure this didn't poison the rest of the DAG
     let s4 = spend_ok4.first().expect("spend_ok4 to be unique");
     let s5 = spend_ok5.first().expect("spend_ok5 to be unique");
 
@@ -125,7 +124,6 @@ fn test_spend_dag_double_spend_poisonning() -> Result<()> {
     Ok(())
 }
 
-#[ignore = "Re-enabled once got DAG auditor properly updated"]
 #[test]
 fn test_spend_dag_double_spend_branches() -> Result<()> {
     let mut net = MockNetwork::genesis()?;
@@ -170,8 +168,7 @@ fn test_spend_dag_double_spend_branches() -> Result<()> {
         dag.insert(spend.address(), spend.clone());
     }
 
-    // TODO: re-enable the assertion once figured out the root reason of the failure.
-    // assert_eq!(dag.record_faults(&genesis), Ok(()));
+    assert_eq!(dag.record_faults(&genesis), Ok(()));
     // dag.dump_to_file("/tmp/test_spend_dag_double_spend_branches")?;
 
     // make sure double spend is detected