Fix #8481: Race condition in NFT tab when user hide/unhide NFT before…

… balance or metadata is fetched (#8486)
brave · Dec 14, 2023 · e443fa6 · e443fa6
1 parent 5b40778
commit e443fa6
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 66 deletions.
diff --git a/Sources/BraveWallet/Crypto/Stores/NFTStore.swift b/Sources/BraveWallet/Crypto/Stores/NFTStore.swift
@@ -261,24 +261,6 @@ public class NFTStore: ObservableObject, WalletObserverStore {
       let selectedAccounts = filters.accounts.filter(\.isSelected).map(\.model)
       let selectedNetworks = filters.networks.filter(\.isSelected).map(\.model)
 
-      // user visible NFTs
-      let userVisibleNFTs = assetManager.getAllUserAssetsInNetworkAssetsByVisibility(networks: selectedNetworks, visible: true)
-        .map { networkAssets in
-          NetworkAssets(
-            network: networkAssets.network,
-            tokens: networkAssets.tokens.filter { $0.isNft || $0.isErc721 },
-            sortOrder: networkAssets.sortOrder
-          )
-        }
-      // user hidden NFTs
-      let userHiddenNFTs = assetManager.getAllUserAssetsInNetworkAssetsByVisibility(networks: selectedNetworks, visible: false)
-        .map { networkAssets in
-          NetworkAssets(
-            network: networkAssets.network,
-            tokens: networkAssets.tokens.filter { $0.isNft || $0.isErc721 },
-            sortOrder: networkAssets.sortOrder
-          )
-        }
       // all spam NFTs marked by SimpleHash
       simpleHashSpamNFTs = await walletService.simpleHashSpamNFTs(for: selectedAccounts, on: selectedNetworks)
       let unionedSpamNFTs = computeSpamNFTs(
@@ -287,29 +269,21 @@ public class NFTStore: ObservableObject, WalletObserverStore {
         simpleHashSpamNFTs: simpleHashSpamNFTs
       )
 
-      let allNetworkNFTs = generateAllNFTsInNetworks(
-        userVisibleNFTs: userVisibleNFTs,
-        userHiddenNFTs: userHiddenNFTs,
-        computedSpamNFTs: unionedSpamNFTs
-      )
-      userNFTGroups = buildNFTGroupModels(
+      let (userNFTGroups, allUserNFTs) = buildNFTGroupModels(
         groupBy: filters.groupBy,
-        allUserNFTs: allNetworkNFTs,
+        spams: unionedSpamNFTs,
         selectedAccounts: selectedAccounts,
         selectedNetworks: selectedNetworks
       )
+      self.userNFTGroups = userNFTGroups
 
-      var allNFTs: [BraveWallet.BlockchainToken] = []
-      for networkAssets in [userVisibleNFTs, userHiddenNFTs, unionedSpamNFTs] {
-        allNFTs.append(contentsOf: networkAssets.flatMap(\.tokens))
-      }
       // if we're not hiding unowned or grouping by account, balance isn't needed
       if filters.isHidingUnownedNFTs || filters.groupBy == .accounts {
         let allAccounts = filters.accounts.map(\.model)
         nftBalancesCache = await withTaskGroup(
           of: [String: [String: Int]].self,
           body: { @MainActor [nftBalancesCache, rpcService] group in
-            for nft in allNFTs { // for each NFT
+            for nft in allUserNFTs { // for each NFT
               guard let networkForNFT = allNetworks.first(where: { $0.chainId == nft.chainId }) else {
                 continue
               }
@@ -343,25 +317,27 @@ public class NFTStore: ObservableObject, WalletObserverStore {
       }
 
       guard !Task.isCancelled else { return }
-      userNFTGroups = buildNFTGroupModels(
+      let (userNFTGroupsWithBalance, _) = buildNFTGroupModels(
         groupBy: filters.groupBy,
-        allUserNFTs: allNetworkNFTs,
+        spams: unionedSpamNFTs,
         selectedAccounts: selectedAccounts,
         selectedNetworks: selectedNetworks
       )
+      self.userNFTGroups = userNFTGroupsWithBalance
 
       // fetch nft metadata for all NFTs
-      let allMetadata = await rpcService.fetchNFTMetadata(tokens: allNFTs, ipfsApi: ipfsApi)
+      let allMetadata = await rpcService.fetchNFTMetadata(tokens: allUserNFTs, ipfsApi: ipfsApi)
       for (key, value) in allMetadata { // update cached values
         metadataCache[key] = value
       }
       guard !Task.isCancelled else { return }
-      userNFTGroups = buildNFTGroupModels(
+      let (userNFTGroupsWithMetadata, _) = buildNFTGroupModels(
         groupBy: filters.groupBy,
-        allUserNFTs: allNetworkNFTs,
+        spams: unionedSpamNFTs,
         selectedAccounts: selectedAccounts,
         selectedNetworks: selectedNetworks
       )
+      self.userNFTGroups = userNFTGroupsWithMetadata
 
       isShowingNFTLoadingState = false
     }
@@ -520,29 +496,57 @@ public class NFTStore: ObservableObject, WalletObserverStore {
 
   private func buildNFTGroupModels(
     groupBy: GroupBy,
-    allUserNFTs: [NetworkAssets],
+    spams: [NetworkAssets],
     selectedAccounts: [BraveWallet.AccountInfo],
     selectedNetworks: [BraveWallet.NetworkInfo]
-  ) -> [NFTGroupViewModel] {
+  ) -> ([NFTGroupViewModel], [BraveWallet.BlockchainToken]) {
+
+    // user visible NFTs
+    let userVisibleNFTs = assetManager.getAllUserAssetsInNetworkAssetsByVisibility(networks: selectedNetworks, visible: true)
+      .map { networkAssets in
+        NetworkAssets(
+          network: networkAssets.network,
+          tokens: networkAssets.tokens.filter { $0.isNft || $0.isErc721 },
+          sortOrder: networkAssets.sortOrder
+        )
+      }
+    // user hidden NFTs
+    let userHiddenNFTs = assetManager.getAllUserAssetsInNetworkAssetsByVisibility(networks: selectedNetworks, visible: false)
+      .map { networkAssets in
+        NetworkAssets(
+          network: networkAssets.network,
+          tokens: networkAssets.tokens.filter { $0.isNft || $0.isErc721 },
+          sortOrder: networkAssets.sortOrder
+        )
+      }
+
+    let allUserNFTsInNetworks = generateAllNFTsInNetworks(
+      userVisibleNFTs: userVisibleNFTs,
+      userHiddenNFTs: userHiddenNFTs,
+      computedSpamNFTs: spams
+    )
+
+    let allUserNFTs = (userVisibleNFTs + userHiddenNFTs + spams).flatMap(\.tokens)
+
     let groups: [NFTGroupViewModel]
     switch filters.groupBy {
     case .none:
       let assets = buildNFTAssetViewModels(
         for: .none,
-        allUserNFTs: allUserNFTs
+        allUserNFTs: allUserNFTsInNetworks
       )
-      return [
+      return ([
         .init(
           groupType: .none,
           assets: assets
         )
-      ]
+      ], allUserNFTs)
     case .accounts:
       groups = selectedAccounts.map { account in
         let groupType: AssetGroupType = .account(account)
         let assets = buildNFTAssetViewModels(
           for: .account(account),
-          allUserNFTs: allUserNFTs
+          allUserNFTs: allUserNFTsInNetworks
         )
         return NFTGroupViewModel(
           groupType: groupType,
@@ -554,15 +558,15 @@ public class NFTStore: ObservableObject, WalletObserverStore {
         let groupType: AssetGroupType = .network(network)
         let assets = buildNFTAssetViewModels(
           for: .network(network),
-          allUserNFTs: allUserNFTs
+          allUserNFTs: allUserNFTsInNetworks
         )
         return NFTGroupViewModel(
           groupType: groupType,
           assets: assets
         )
       }
     }
-    return groups
+    return (groups, allUserNFTs)
   }
 
   @MainActor func isNFTDiscoveryEnabled() async -> Bool {
@@ -589,36 +593,16 @@ public class NFTStore: ObservableObject, WalletObserverStore {
       guard let self else { return }
       let selectedAccounts = self.filters.accounts.filter(\.isSelected).map(\.model)
       let selectedNetworks = self.filters.networks.filter(\.isSelected).map(\.model)
-      let userVisibleAssets = self.assetManager.getAllUserAssetsInNetworkAssetsByVisibility(networks: selectedNetworks, visible: true)
-        .map { networkAssets in
-          NetworkAssets(
-            network: networkAssets.network,
-            tokens: networkAssets.tokens.filter { $0.isNft || $0.isErc721 },
-            sortOrder: networkAssets.sortOrder
-          )
-        }
-      let userHiddenAssets = self.assetManager.getAllUserAssetsInNetworkAssetsByVisibility(networks: selectedNetworks, visible: false)
-        .map { networkAssets in
-          NetworkAssets(
-            network: networkAssets.network,
-            tokens: networkAssets.tokens.filter { $0.isNft || $0.isErc721 },
-            sortOrder: networkAssets.sortOrder
-          )
-        }
+
       let unionedSpamNFTs = computeSpamNFTs(
         selectedNetworks: selectedNetworks,
         selectedAccounts: selectedAccounts,
         simpleHashSpamNFTs: simpleHashSpamNFTs
       )
 
-      let allNetworkNFTs = generateAllNFTsInNetworks(
-        userVisibleNFTs: userVisibleAssets,
-        userHiddenNFTs: userHiddenAssets,
-        computedSpamNFTs: unionedSpamNFTs
-      )
-      userNFTGroups = buildNFTGroupModels(
+      (userNFTGroups, _) = buildNFTGroupModels(
         groupBy: filters.groupBy,
-        allUserNFTs: allNetworkNFTs,
+        spams: unionedSpamNFTs,
         selectedAccounts: selectedAccounts,
         selectedNetworks: selectedNetworks
       )

diff --git a/Sources/BraveWallet/Extensions/RpcServiceExtensions.swift b/Sources/BraveWallet/Extensions/RpcServiceExtensions.swift
@@ -103,6 +103,8 @@ extension BraveWalletJsonRpcService {
       }
     case .btc:
       completion(nil)
+    case .zec:
+      completion(nil)
     @unknown default:
       completion(nil)
     }

diff --git a/Tests/BraveWalletTests/NFTStoreTests.swift b/Tests/BraveWalletTests/NFTStoreTests.swift
@@ -73,6 +73,8 @@ class NFTStoreTests: XCTestCase {
         completion([.mockFilecoinTestnet])
       case .btc:
         XCTFail("Should not fetch btc network")
+      case .zec:
+        XCTFail("Should not fetch zec network")
       @unknown default:
         XCTFail("Should not fetch unknown network")
       }