add gated Polkadot asset discovery

Author: cmazakas

browser/about_flags.cc

@@ -155,6 +155,16 @@ const flags_ui::FeatureEntry::FeatureVariation kZCashFeatureVariations[] = {
     {"- Shielded support disabled", kZCashShieldedTransactionsDisabled,
      nullptr},
     {"- Shielded support enabled", kZCashShieldedTransactionsEnabled, nullptr}};
+
+const flags_ui::FeatureEntry::FeatureParam kPolkadotAssetDiscoveryDisabled[] = {
+    {"polkadot_asset_discovery", "false"}};
+
+const flags_ui::FeatureEntry::FeatureParam kPolkadotAssetDiscoveryEnabled[] = {
+    {"polkadot_asset_discovery", "true"}};
+
+const flags_ui::FeatureEntry::FeatureVariation kPolkadotFeatureVariations[] = {
+    {"- Asset discovery disabled", kPolkadotAssetDiscoveryDisabled, nullptr},
+    {"- Asset discovery enabled", kPolkadotAssetDiscoveryEnabled, nullptr}};
 #endif  // BUILDFLAG(ENABLE_BRAVE_WALLET)
 
 #if defined(TOOLKIT_VIEWS)
@@ -237,8 +247,10 @@ const char* const kBraveSyncImplLink[1] = {"https://github.com/brave/go-sync"};
           "Enable Brave Wallet Polkadot support",                             \
           "Polkadot support for native Brave Wallet",                         \
           kOsDesktop,                                                         \
-          FEATURE_VALUE_TYPE(                                                 \
-              brave_wallet::features::kBraveWalletPolkadotFeature),           \
+          FEATURE_WITH_PARAMS_VALUE_TYPE(                                     \
+              brave_wallet::features::kBraveWalletPolkadotFeature,            \
+              kPolkadotFeatureVariations,                                     \
+              "BraveWalletPolkadot"),                                         \
       },                                                                      \
       {                                                                       \
           "brave-wallet-bitcoin",                                             \

browser/brave_wallet/asset_discovery_task_unittest.cc

@@ -7,6 +7,7 @@
 
 #include <algorithm>
 
+#include "base/containers/extend.h"
 #include "base/no_destructor.h"
 #include "brave/components/brave_wallet/browser/brave_wallet_constants.h"
 #include "brave/components/brave_wallet/browser/brave_wallet_service.h"
@@ -108,6 +109,12 @@ AssetDiscoveryManager::GetFungibleSupportedChains() {
   supported_chains.push_back(
       mojom::ChainId::New(mojom::CoinType::SOL, mojom::kSolanaMainnet));
 
+  if (IsPolkadotAssetDiscoveryEnabled()) {
+    base::Extend(supported_chains, std::array{mojom::kPolkadotMainnetAssetHub,
+                                              mojom::kPolkadotTestnetAssetHub,
+                                              mojom::kPolkadotPaseoAssetHub});
+  }
+
   return supported_chains;
 }
 
@@ -165,11 +172,17 @@ void AssetDiscoveryManager::FinishTask() {
 
 void AssetDiscoveryManager::AccountsAdded(
     std::vector<mojom::AccountInfoPtr> added_accounts) {
+  if (!auto_discovery_enabled_) {
+    return;
+  }
+
   std::vector<mojom::AccountIdPtr> accounts;
   for (const auto& account : added_accounts) {
     auto& account_id = account->account_id;
     if (account->account_id->coin == mojom::CoinType::ETH ||
-        account->account_id->coin == mojom::CoinType::SOL) {
+        account->account_id->coin == mojom::CoinType::SOL ||
+        (account->account_id->coin == mojom::CoinType::DOT &&
+         IsPolkadotAssetDiscoveryEnabled())) {
       accounts.push_back(account_id.Clone());
     }
   }

components/brave_wallet/browser/asset_discovery_manager.cc

@@ -63,6 +63,12 @@ class AssetDiscoveryManager : public KeyringServiceObserverBase {
 
   size_t GetQueueSizeForTesting() { return queue_.size(); }
 
+  // Stops account creation from kicking off a discovery run. Explicit
+  // DiscoverAssetsOnAllSupportedChains() calls are unaffected.
+  void SetAutoDiscoveryEnabledForTesting(bool enabled) {
+    auto_discovery_enabled_ = enabled;
+  }
+
  private:
   friend class AssetDiscoveryManagerUnitTest;
   FRIEND_TEST_ALL_PREFIXES(AssetDiscoveryManagerUnitTest,
@@ -83,6 +89,7 @@ class AssetDiscoveryManager : public KeyringServiceObserverBase {
   raw_ref<KeyringService> keyring_service_;
   raw_ref<SimpleHashClient> simple_hash_client_;
   raw_ptr<PrefService> prefs_;
+  bool auto_discovery_enabled_ = true;
   mojo::Receiver<brave_wallet::mojom::KeyringServiceObserver>
       keyring_service_observer_receiver_{this};
   base::WeakPtrFactory<AssetDiscoveryManager> weak_ptr_factory_;

components/brave_wallet/browser/asset_discovery_manager.h

@@ -5,6 +5,7 @@
 
 #include "brave/components/brave_wallet/browser/asset_discovery_task.h"
 
+#include <algorithm>
 #include <map>
 #include <optional>
 #include <string_view>
@@ -14,11 +15,15 @@
 #include "base/base64.h"
 #include "base/check.h"
 #include "base/containers/span_reader.h"
+#include "base/strings/string_number_conversions.h"
 #include "brave/components/brave_wallet/browser/blockchain_registry.h"
 #include "brave/components/brave_wallet/browser/brave_wallet_constants.h"
 #include "brave/components/brave_wallet/browser/brave_wallet_service.h"
 #include "brave/components/brave_wallet/browser/brave_wallet_utils.h"
 #include "brave/components/brave_wallet/browser/json_rpc_service.h"
+#include "brave/components/brave_wallet/browser/network_manager.h"
+#include "brave/components/brave_wallet/browser/polkadot/polkadot_utils.h"
+#include "brave/components/brave_wallet/browser/polkadot/polkadot_wallet_service.h"
 #include "brave/components/brave_wallet/browser/pref_names.h"
 #include "brave/components/brave_wallet/common/brave_wallet_constants.h"
 #include "brave/components/brave_wallet/common/common_utils.h"
@@ -45,6 +50,18 @@ std::pair<std::vector<T>, std::vector<T>> SplitByCoin(
   return result;
 }
 
+template <class T>
+std::vector<T> FilterByCoin(const std::vector<T>& items, mojom::CoinType coin) {
+  std::vector<T> result;
+  for (const auto& item : items) {
+    if (item->coin == coin) {
+      result.push_back(item.Clone());
+    }
+  }
+
+  return result;
+}
+
 }  // namespace
 
 AssetDiscoveryTask::AssetDiscoveryTask(APIRequestHelper& api_request_helper,
@@ -102,12 +119,21 @@ void AssetDiscoveryTask::DiscoverAssets(
 
   bool use_ankr_discovery =
       IsAnkrBalancesEnabled() && !ankr_evm_chain_ids.empty();
+  bool discover_polkadot_assets = IsPolkadotAssetDiscoveryEnabled();
 
-  // Concurrently discover ETH ERC20s on our registry, Solana tokens on our
-  // Registry and NFTs on both platforms, then merge the results
+  // Concurrently discover ETH ERC20s on our registry, Solana tokens and
+  // Polkadot assets on our registry, and NFTs on both platforms, then merge
+  // the results
+  size_t discovery_count = 3;  // SPL tokens, EVM registry tokens and NFTs.
+  if (use_ankr_discovery) {
+    ++discovery_count;
+  }
+  if (discover_polkadot_assets) {
+    ++discovery_count;
+  }
   const auto barrier_callback =
       base::BarrierCallback<std::vector<mojom::BlockchainTokenPtr>>(
-          use_ankr_discovery ? 4 : 3,
+          discovery_count,
           base::BindOnce(&AssetDiscoveryTask::MergeDiscoveredAssets,
                          weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
   // Currently SPL tokens are only discovered on Solana Mainnet.
@@ -122,6 +148,13 @@ void AssetDiscoveryTask::DiscoverAssets(
     DiscoverERC20sFromRegistry(eth_accounts, evm_chain_ids, barrier_callback);
   }
 
+  if (discover_polkadot_assets) {
+    DiscoverPolkadotAssetsFromRegistry(
+        FilterByCoin(accounts, mojom::CoinType::DOT),
+        FilterByCoin(fungible_chain_ids, mojom::CoinType::DOT),
+        barrier_callback);
+  }
+
   DiscoverNFTs(accounts, non_fungible_chain_ids, barrier_callback);
 }
 
@@ -461,6 +494,146 @@ void AssetDiscoveryTask::OnGetSolanaTokenRegistry(
   std::move(callback).Run(std::move(discovered_tokens));
 }
 
+void AssetDiscoveryTask::DiscoverPolkadotAssetsFromRegistry(
+    base::span<const mojom::AccountIdPtr> accounts,
+    base::span<const mojom::ChainIdPtr> chain_ids,
+    DiscoverAssetsCompletedCallback callback) {
+  if (accounts.empty() || chain_ids.empty() ||
+      !wallet_service_->GetPolkadotWalletService()) {
+    std::move(callback).Run({});
+    return;
+  }
+
+  // Polkadot keyrings are chain scoped, so a testnet account has nothing to
+  // find on mainnet Asset Hub and vice versa.
+  size_t balance_query_count = 0;
+  std::vector<std::pair<std::string, std::vector<mojom::AccountIdPtr>>>
+      accounts_per_chain;
+  for (const auto& chain_id : chain_ids) {
+    auto network = wallet_service_->network_manager()->GetChain(
+        chain_id->chain_id, mojom::CoinType::DOT);
+    if (!network) {
+      continue;
+    }
+
+    std::vector<mojom::AccountIdPtr> chain_accounts;
+    for (const auto& account_id : accounts) {
+      if (std::ranges::contains(network->supported_keyrings,
+                                account_id->keyring_id)) {
+        chain_accounts.push_back(account_id.Clone());
+      }
+    }
+
+    if (chain_accounts.empty()) {
+      continue;
+    }
+
+    // We performa a balance query per account.
+    balance_query_count += chain_accounts.size();
+    accounts_per_chain.emplace_back(chain_id->chain_id,
+                                    std::move(chain_accounts));
+  }
+
+  if (balance_query_count == 0) {
+    std::move(callback).Run({});
+    return;
+  }
+
+  // Each chain fans out into one balance query per account on that chain.
+  const auto barrier_callback =
+      base::BarrierCallback<std::vector<mojom::BlockchainTokenPtr>>(
+          balance_query_count,
+          base::BindOnce(&AssetDiscoveryTask::MergeDiscoveredPolkadotAssets,
+                         weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
+
+  for (auto& [chain_id, chain_accounts] : accounts_per_chain) {
+    BlockchainRegistry::GetInstance()->GetAllTokens(
+        chain_id, mojom::CoinType::DOT,
+        base::BindOnce(&AssetDiscoveryTask::OnGetPolkadotTokenRegistry,
+                       weak_ptr_factory_.GetWeakPtr(), barrier_callback,
+                       std::move(chain_accounts), chain_id));
+  }
+}
+
+void AssetDiscoveryTask::OnGetPolkadotTokenRegistry(
+    base::RepeatingCallback<void(std::vector<mojom::BlockchainTokenPtr>)>
+        barrier_callback,
+    std::vector<mojom::AccountIdPtr> accounts,
+    const std::string& chain_id,
+    std::vector<mojom::BlockchainTokenPtr> registry_tokens) {
+  // The registry carries the pallet_assets asset id as the contract address.
+  // `candidates` is kept positionally aligned with `asset_ids` so the balances
+  // can be zipped back onto the tokens they belong to.
+  std::vector<mojom::BlockchainTokenPtr> candidates;
+  std::vector<uint32_t> asset_ids;
+  for (auto& token : registry_tokens) {
+    uint32_t asset_id = 0;
+    if (!base::StringToUint(token->contract_address, &asset_id)) {
+      continue;
+    }
+    asset_ids.push_back(asset_id);
+    candidates.push_back(std::move(token));
+  }
+
+  auto* polkadot_wallet_service = wallet_service_->GetPolkadotWalletService();
+  if (asset_ids.empty() || !polkadot_wallet_service) {
+    for (size_t i = 0; i < accounts.size(); ++i) {
+      barrier_callback.Run({});
+    }
+    return;
+  }
+
+  for (const auto& account_id : accounts) {
+    polkadot_wallet_service->GetAssetAccountBalances(
+        account_id.Clone(), asset_ids, chain_id,
+        base::BindOnce(&AssetDiscoveryTask::OnGetPolkadotAssetAccountBalances,
+                       weak_ptr_factory_.GetWeakPtr(), barrier_callback,
+                       CloneVector(candidates)));
+  }
+}
+
+void AssetDiscoveryTask::OnGetPolkadotAssetAccountBalances(
+    base::OnceCallback<void(std::vector<mojom::BlockchainTokenPtr>)>
+        barrier_callback,
+    std::vector<mojom::BlockchainTokenPtr> candidates,
+    std::vector<mojom::PolkadotAssetAccountInfoPtr> asset_accounts,
+    const std::optional<std::string>& error_message) {
+  // Balances are returned positionally zipped with the asset ids we asked for,
+  // so a size mismatch leaves us unable to tell which asset a balance is for.
+  if (error_message || asset_accounts.size() != candidates.size()) {
+    std::move(barrier_callback).Run({});
+    return;
+  }
+
+  std::vector<mojom::BlockchainTokenPtr> discovered_tokens;
+  for (size_t i = 0; i < asset_accounts.size(); ++i) {
+    if (MojomToUint128(asset_accounts[i]->balance) == 0) {
+      continue;
+    }
+    discovered_tokens.push_back(std::move(candidates[i]));
+  }
+
+  std::move(barrier_callback).Run(std::move(discovered_tokens));
+}
+
+void AssetDiscoveryTask::MergeDiscoveredPolkadotAssets(
+    DiscoverAssetsCompletedCallback callback,
+    const std::vector<std::vector<mojom::BlockchainTokenPtr>>&
+        discovered_assets) {
+  std::vector<mojom::BlockchainTokenPtr> discovered_tokens;
+  for (const auto& tokens : discovered_assets) {
+    for (const auto& token : tokens) {
+      // Several accounts can hold the same asset; AddUserAsset deduplicates.
+      if (!AddUserAsset(prefs_, token.Clone())) {
+        continue;
+      }
+      discovered_tokens.push_back(token.Clone());
+    }
+  }
+
+  std::move(callback).Run(std::move(discovered_tokens));
+}
+
 void AssetDiscoveryTask::DiscoverNFTs(
     const std::vector<mojom::AccountIdPtr>& accounts,
     const std::vector<mojom::ChainIdPtr>& chain_ids,

components/brave_wallet/browser/asset_discovery_task.cc

@@ -120,6 +120,30 @@ class AssetDiscoveryTask {
       const base::flat_set<std::string>& discovered_contract_addresses,
       std::vector<mojom::BlockchainTokenPtr> sol_token_registry);
 
+  // Polkadot assets live in pallet_assets on the Asset Hub parachains and are
+  // keyed by a numeric asset id, which the registry carries as the token's
+  // contract address.
+  void DiscoverPolkadotAssetsFromRegistry(
+      base::span<const mojom::AccountIdPtr> accounts,
+      base::span<const mojom::ChainIdPtr> chain_ids,
+      DiscoverAssetsCompletedCallback callback);
+  void OnGetPolkadotTokenRegistry(
+      base::RepeatingCallback<void(std::vector<mojom::BlockchainTokenPtr>)>
+          barrier_callback,
+      std::vector<mojom::AccountIdPtr> accounts,
+      const std::string& chain_id,
+      std::vector<mojom::BlockchainTokenPtr> registry_tokens);
+  void OnGetPolkadotAssetAccountBalances(
+      base::OnceCallback<void(std::vector<mojom::BlockchainTokenPtr>)>
+          barrier_callback,
+      std::vector<mojom::BlockchainTokenPtr> candidates,
+      std::vector<mojom::PolkadotAssetAccountInfoPtr> asset_accounts,
+      const std::optional<std::string>& error_message);
+  void MergeDiscoveredPolkadotAssets(
+      DiscoverAssetsCompletedCallback callback,
+      const std::vector<std::vector<mojom::BlockchainTokenPtr>>&
+          discovered_assets);
+
   void DiscoverNFTs(const std::vector<mojom::AccountIdPtr>& accounts,
                     const std::vector<mojom::ChainIdPtr>& chain_ids,
                     DiscoverAssetsCompletedCallback callback);

components/brave_wallet/browser/asset_discovery_task.h

@@ -257,9 +257,12 @@ bool ParseTokenList(const std::string& json, TokenListMap* token_list_map) {
     }
 
     // Determine coin type based on chain_id
-    mojom::CoinType coin = chain_id == mojom::kSolanaMainnet
-                               ? mojom::CoinType::SOL
-                               : mojom::CoinType::ETH;
+    mojom::CoinType coin = mojom::CoinType::ETH;
+    if (chain_id == mojom::kSolanaMainnet) {
+      coin = mojom::CoinType::SOL;
+    } else if (chain_id == mojom::kPolkadotMainnetAssetHub) {
+      coin = mojom::CoinType::DOT;
+    }
 
     // Iterate through tokens in this chain
     for (const auto token_pair : *tokens_dict) {

components/brave_wallet/browser/blockchain_list_parser.cc

@@ -370,6 +370,9 @@ class BraveWalletService : public KeyedService,
   AccountDiscoveryManager* account_discovery_manager() {
     return account_discovery_manager_.get();
   }
+  AssetDiscoveryManager* asset_discovery_manager() {
+    return asset_discovery_manager_.get();
+  }
 
   // Might return nullptr.
   BitcoinWalletService* GetBitcoinWalletService();

components/brave_wallet/browser/brave_wallet_service.h

@@ -105,6 +105,12 @@ std::optional<std::string> GetUserAssetAddress(const std::string& address,
     }
   }
 
+  if (coin == mojom::CoinType::DOT) {
+    // Polkadot assets are keyed by a numeric asset id, which acts as the
+    // contract address.
+    return address;
+  }
+
   return std::nullopt;
 }