Add canister method to store imported tokens (#5210)

# Motivation

We want to allow users to import custom tokens in NNS dapp by specifying
a ledger canister ID and optional index canister ID.
The imported tokens should be persisted with the account so we need to
store them in the nns-dapp backend canister.
We need to limit the number of tokens per user to have a bound on the
amount of storage used.
We arbitrarily set the limit to 10 and can increase it later if we want.
We store the tokens in the user account which is stored in stable
structures.
To keep things simple we just support setting and getting the full list
of tokens.

# Changes

Add `set_imported_tokens` and `get_imported_tokens` canister methods.

# Tests

1. Unit tests added.
2. Tested manually with dfx:
```
$ dfx canister call nns-dapp get_imported_tokens --candid rs/backend/nns-dapp.did
(variant { AccountNotFound })

$ dfx canister call nns-dapp  --candid rs/backend/nns-dapp.did add_account
("a6e3ccdf7ede8e4b95c522d990a7be4d9d970b1f7bdcedf3baec857b65f7d8bd")

$ dfx canister call nns-dapp --candid rs/backend/nns-dapp.did get_imported_tokens
(variant { Ok = record { imported_tokens = vec {} } })

$ dfx canister call nns-dapp --candid rs/backend/nns-dapp.did set_imported_tokens '(record{imported_tokens = vec{ record{ledger_canister_id=principal"xnjld-hqaaa-aaaal-qb56q-cai"}  }})'
(variant { Ok })

$ dfx canister call nns-dapp --candid rs/backend/nns-dapp.did get_imported_tokens
(
  variant {
    Ok = record {
      imported_tokens = vec {
        record {
          index_canister_id = null;
          ledger_canister_id = principal "xnjld-hqaaa-aaaal-qb56q-cai";
        };
      };
    }
  },
)

$ dfx canister call nns-dapp --candid rs/backend/nns-dapp.did set_imported_tokens '(record{imported_tokens = vec{ record{ledger_canister_id=principal"xnjld-hqaaa-aaaal-qb56q-cai"};record{ledger_canister_id=principal"xnjld-hqaaa-aaaal-qb56q-cai"};record{ledger_canister_id=principal"xnjld-hqaaa-aaaal-qb56q-cai"};record{ledger_canister_id=principal"xnjld-hqaaa-aaaal-qb56q-cai"};record{ledger_canister_id=principal"xnjld-hqaaa-aaaal-qb56q-cai"};record{ledger_canister_id=principal"xnjld-hqaaa-aaaal-qb56q-cai"};record{ledger_canister_id=principal"xnjld-hqaaa-aaaal-qb56q-cai"};record{ledger_canister_id=principal"xnjld-hqaaa-aaaal-qb56q-cai"};record{ledger_canister_id=principal"xnjld-hqaaa-aaaal-qb56q-cai"};record{ledger_canister_id=principal"xnjld-hqaaa-aaaal-qb56q-cai"};record{ledger_canister_id=principal"xnjld-hqaaa-aaaal-qb56q-cai"};  }})'
(variant { TooManyImportedTokens = record { limit = 10 : int32 } })
```
3. Max tested it in his branch with the frontend code.

# Todos

- [x] Add entry to changelog (if necessary).

CHANGELOG-Nns-Dapp-unreleased.md

@@ -18,6 +18,7 @@ proposal is successful, the changes it released will be moved from this file to
 * Add 2 new topics `ProtocolCanisterMangement` and
 * Make the WASM accessible via beta subdomains so we can deploy early versions there.
 * Enable warning dialog on beta deployment.
+* Back-end support for storing imported tokens.
 
 #### Changed
 

rs/backend/nns-dapp-exports-production.txt

@@ -7,6 +7,7 @@ canister_query get_canisters
 canister_query get_exceptional_transactions
 canister_query get_histogram
 canister_query get_stats
+canister_query get_imported_tokens
 canister_query http_request
 canister_update add_account
 canister_update add_assets_tar_xz
@@ -18,5 +19,6 @@ canister_update get_proposal_payload
 canister_update register_hardware_wallet
 canister_update rename_canister
 canister_update rename_sub_account
+canister_update set_imported_tokens
 canister_update step_migration
 main

rs/backend/nns-dapp-exports-test.txt

@@ -8,6 +8,7 @@ canister_query get_exceptional_transactions
 canister_query get_histogram
 canister_query get_stats
 canister_query get_toy_account
+canister_query get_imported_tokens
 canister_query http_request
 canister_update add_account
 canister_update add_assets_tar_xz
@@ -20,5 +21,6 @@ canister_update get_proposal_payload
 canister_update register_hardware_wallet
 canister_update rename_canister
 canister_update rename_sub_account
+canister_update set_imported_tokens
 canister_update step_migration
 main

rs/backend/nns-dapp.did

@@ -193,6 +193,30 @@ type SchemaLabel = variant {
     AccountsInStableMemory;
 };
 
+type ImportedToken =
+    record {
+        ledger_canister_id: principal;
+        index_canister_id: opt principal;
+    };
+
+type ImportedTokens =
+    record {
+        imported_tokens: vec ImportedToken;
+    };
+
+type SetImportedTokensResponse =
+    variant {
+        Ok;
+        AccountNotFound;
+        TooManyImportedTokens: record{limit: int32};
+    };
+
+type GetImportedTokensResponse =
+    variant {
+        Ok: ImportedTokens;
+        AccountNotFound;
+    };
+
 service: (opt Config) -> {
     get_account: () -> (GetAccountResponse) query;
     add_account: () -> (AccountIdentifier);
@@ -203,6 +227,8 @@ service: (opt Config) -> {
     attach_canister: (AttachCanisterRequest) -> (AttachCanisterResponse);
     rename_canister: (RenameCanisterRequest) -> (RenameCanisterResponse);
     detach_canister: (DetachCanisterRequest) -> (DetachCanisterResponse);
+    set_imported_tokens: (ImportedTokens) -> (SetImportedTokensResponse);
+    get_imported_tokens: () -> (GetImportedTokensResponse) query;
     get_proposal_payload: (nat64) -> (GetProposalPayloadResponse);
     get_stats: () -> (Stats) query;
     get_histogram: () -> (Histogram) query;

rs/backend/src/accounts_store.rs

@@ -37,6 +37,10 @@ use self::schema::SchemaLabel;
 /// accounts we avoid some complications.
 const PRE_MIGRATION_LIMIT: u64 = 300_000;
 
+// Conservatively limit the number of imported tokens to prevent using too much memory.
+// Can be revisited if users find this too restrictive.
+const MAX_IMPORTED_TOKENS: i32 = 20;
+
 /// Accounts, transactions and related data.
 ///
 /// Note: Some monitoring fields are not included in the `Eq` and `PartialEq` implementations.  Additionally, please note
@@ -154,6 +158,7 @@ pub struct Account {
     sub_accounts: HashMap<u8, NamedSubAccount>,
     hardware_wallet_accounts: Vec<NamedHardwareWalletAccount>,
     canisters: Vec<NamedCanister>,
+    imported_tokens: Option<ImportedTokens>,
     // default_account_transactions: Do not reuse this field. There are still accounts in stable memor with this unused field.
 }
 
@@ -216,6 +221,30 @@ impl PartialOrd for NamedCanister {
     }
 }
 
+#[derive(CandidType, Clone, Copy, Default, Deserialize, Debug, Eq, PartialEq)]
+pub struct ImportedToken {
+    ledger_canister_id: PrincipalId,
+    index_canister_id: Option<PrincipalId>,
+}
+
+#[derive(CandidType, Clone, Default, Deserialize, Debug, Eq, PartialEq)]
+pub struct ImportedTokens {
+    imported_tokens: Vec<ImportedToken>,
+}
+
+#[derive(CandidType, Debug, PartialEq)]
+pub enum SetImportedTokensResponse {
+    Ok,
+    AccountNotFound,
+    TooManyImportedTokens { limit: i32 },
+}
+
+#[derive(CandidType, Debug, PartialEq)]
+pub enum GetImportedTokensResponse {
+    Ok(ImportedTokens),
+    AccountNotFound,
+}
+
 #[derive(Copy, Clone, CandidType, Deserialize, Debug, Eq, PartialEq)]
 pub enum TransactionType {
     Burn,
@@ -746,6 +775,36 @@ impl AccountsStore {
         }
     }
 
+    pub fn set_imported_tokens(
+        &mut self,
+        caller: PrincipalId,
+        new_imported_tokens: ImportedTokens,
+    ) -> SetImportedTokensResponse {
+        if new_imported_tokens.imported_tokens.len() > (MAX_IMPORTED_TOKENS as usize) {
+            return SetImportedTokensResponse::TooManyImportedTokens {
+                limit: MAX_IMPORTED_TOKENS,
+            };
+        }
+        let account_identifier = AccountIdentifier::from(caller).to_vec();
+        let Some(mut account) = self.accounts_db.db_get_account(&account_identifier) else {
+            return SetImportedTokensResponse::AccountNotFound;
+        };
+
+        account.imported_tokens = Some(new_imported_tokens);
+
+        self.accounts_db.db_insert_account(&account_identifier, account);
+        SetImportedTokensResponse::Ok
+    }
+
+    pub fn get_imported_tokens(&mut self, caller: PrincipalId) -> GetImportedTokensResponse {
+        let account_identifier = AccountIdentifier::from(caller).to_vec();
+        let Some(account) = self.accounts_db.db_get_account(&account_identifier) else {
+            return GetImportedTokensResponse::AccountNotFound;
+        };
+
+        GetImportedTokensResponse::Ok(account.imported_tokens.unwrap_or_default())
+    }
+
     #[must_use]
     pub fn get_block_height_synced_up_to(&self) -> Option<BlockIndex> {
         self.block_height_synced_up_to
@@ -1116,6 +1175,7 @@ impl Account {
             sub_accounts: HashMap::new(),
             hardware_wallet_accounts: Vec::new(),
             canisters: Vec::new(),
+            imported_tokens: None,
         }
     }
 }

rs/backend/src/accounts_store/schema/tests.rs

@@ -24,6 +24,7 @@ pub fn toy_account(account_index: u64, num_canisters: u64) -> Account {
         sub_accounts: HashMap::new(),
         hardware_wallet_accounts: Vec::new(),
         canisters: Vec::new(),
+        imported_tokens: None,
     };
     // Attaches canisters to the account.
     for canister_index in 0..num_canisters {

rs/backend/src/accounts_store/tests.rs

@@ -779,6 +779,145 @@ fn detach_canister_canister_not_found() {
     assert_eq!(canister_id1, canisters[0].canister_id);
 }
 
+#[test]
+fn set_and_get_imported_tokens() {
+    let mut store = setup_test_store();
+    let principal = PrincipalId::from_str(TEST_ACCOUNT_1).unwrap();
+    let ledger_canister_id = PrincipalId::new_user_test_id(101);
+    let index_canister_id = PrincipalId::new_user_test_id(102);
+
+    assert_eq!(
+        store.get_imported_tokens(principal),
+        GetImportedTokensResponse::Ok(ImportedTokens::default())
+    );
+
+    let imported_token = ImportedToken {
+        ledger_canister_id,
+        index_canister_id: Some(index_canister_id),
+    };
+
+    assert_eq!(
+        store.set_imported_tokens(
+            principal,
+            ImportedTokens {
+                imported_tokens: vec![imported_token.clone()],
+            },
+        ),
+        SetImportedTokensResponse::Ok
+    );
+
+    assert_eq!(
+        store.get_imported_tokens(principal),
+        GetImportedTokensResponse::Ok(ImportedTokens {
+            imported_tokens: vec![imported_token],
+        })
+    );
+}
+
+#[test]
+fn set_and_get_imported_tokens_without_index_canister() {
+    let mut store = setup_test_store();
+    let principal = PrincipalId::from_str(TEST_ACCOUNT_1).unwrap();
+    let ledger_canister_id = PrincipalId::new_user_test_id(101);
+
+    assert_eq!(
+        store.get_imported_tokens(principal),
+        GetImportedTokensResponse::Ok(ImportedTokens::default())
+    );
+
+    let imported_token = ImportedToken {
+        ledger_canister_id,
+        index_canister_id: None,
+    };
+
+    assert_eq!(
+        store.set_imported_tokens(
+            principal,
+            ImportedTokens {
+                imported_tokens: vec![imported_token.clone()],
+            },
+        ),
+        SetImportedTokensResponse::Ok
+    );
+
+    assert_eq!(
+        store.get_imported_tokens(principal),
+        GetImportedTokensResponse::Ok(ImportedTokens {
+            imported_tokens: vec![imported_token],
+        })
+    );
+}
+
+fn get_unique_imported_tokens(count: u64) -> Vec<ImportedToken> {
+    (0..count)
+        .map(|i| ImportedToken {
+            ledger_canister_id: PrincipalId::new_user_test_id(i),
+            index_canister_id: Some(PrincipalId::new_user_test_id(i + 1000)),
+        })
+        .collect()
+}
+
+#[test]
+fn set_and_get_20_imported_tokens() {
+    let mut store = setup_test_store();
+    let principal = PrincipalId::from_str(TEST_ACCOUNT_1).unwrap();
+
+    assert_eq!(
+        store.get_imported_tokens(principal),
+        GetImportedTokensResponse::Ok(ImportedTokens::default())
+    );
+
+    let imported_tokens = get_unique_imported_tokens(20);
+
+    assert_eq!(
+        store.set_imported_tokens(
+            principal,
+            ImportedTokens {
+                imported_tokens: imported_tokens.clone()
+            },
+        ),
+        SetImportedTokensResponse::Ok
+    );
+
+    assert_eq!(
+        store.get_imported_tokens(principal),
+        GetImportedTokensResponse::Ok(ImportedTokens { imported_tokens })
+    );
+}
+
+#[test]
+fn set_imported_tokens_account_not_found() {
+    let mut store = setup_test_store();
+    let non_existing_principal = PrincipalId::from_str(TEST_ACCOUNT_3).unwrap();
+    assert_eq!(
+        store.set_imported_tokens(non_existing_principal, ImportedTokens::default()),
+        SetImportedTokensResponse::AccountNotFound
+    );
+}
+
+#[test]
+fn set_imported_tokens_too_many() {
+    let mut store = setup_test_store();
+    let principal = PrincipalId::from_str(TEST_ACCOUNT_1).unwrap();
+
+    let imported_tokens = get_unique_imported_tokens(21);
+
+    assert_eq!(
+        store.set_imported_tokens(principal, ImportedTokens { imported_tokens },),
+        SetImportedTokensResponse::TooManyImportedTokens { limit: 20 }
+    );
+}
+
+#[test]
+fn get_imported_tokens_account_not_found() {
+    let mut store = setup_test_store();
+    let non_existing_principal = PrincipalId::from_str(TEST_ACCOUNT_3).unwrap();
+    assert_eq!(
+        store.get_imported_tokens(non_existing_principal),
+        GetImportedTokensResponse::AccountNotFound
+    );
+}
+
 #[test]
 fn sub_account_name_too_long() {
     let mut store = setup_test_store();

rs/backend/src/accounts_store/toy_data.rs

@@ -59,6 +59,7 @@ pub fn toy_account(account_index: u64, size: ToyAccountSize) -> Account {
         sub_accounts: HashMap::new(),
         hardware_wallet_accounts: Vec::new(),
         canisters: Vec::new(),
+        imported_tokens: None,
     };
     // Creates linked sub-accounts:
     // Note: Successive accounts have 0, 1, 2 ... MAX_SUB_ACCOUNTS_PER_ACCOUNT-1 sub accounts, restarting at 0.

rs/backend/src/main.rs

@@ -1,8 +1,9 @@
 use crate::accounts_store::histogram::AccountsStoreHistogram;
 use crate::accounts_store::{
     AccountDetails, AttachCanisterRequest, AttachCanisterResponse, CreateSubAccountResponse, DetachCanisterRequest,
-    DetachCanisterResponse, NamedCanister, RegisterHardwareWalletRequest, RegisterHardwareWalletResponse,
-    RenameCanisterRequest, RenameCanisterResponse, RenameSubAccountRequest, RenameSubAccountResponse,
+    DetachCanisterResponse, GetImportedTokensResponse, ImportedTokens, NamedCanister, RegisterHardwareWalletRequest,
+    RegisterHardwareWalletResponse, RenameCanisterRequest, RenameCanisterResponse, RenameSubAccountRequest,
+    RenameSubAccountResponse, SetImportedTokensResponse,
 };
 use crate::arguments::{set_canister_arguments, CanisterArguments, CANISTER_ARGUMENTS};
 use crate::assets::{hash_bytes, insert_asset, insert_tar_xz, Asset};
@@ -238,6 +239,26 @@ fn detach_canister_impl(request: DetachCanisterRequest) -> DetachCanisterRespons
     STATE.with(|s| s.accounts_store.borrow_mut().detach_canister(principal, request))
 }
 
+#[export_name = "canister_update set_imported_tokens"]
+pub fn set_imported_tokens() {
+    over(candid_one, set_imported_tokens_impl);
+}
+
+fn set_imported_tokens_impl(settings: ImportedTokens) -> SetImportedTokensResponse {
+    let principal = dfn_core::api::caller();
+    STATE.with(|s| s.accounts_store.borrow_mut().set_imported_tokens(principal, settings))
+}
+
+#[export_name = "canister_query get_imported_tokens"]
+pub fn get_imported_tokens() {
+    over(candid_one, |()| get_imported_tokens_impl());
+}
+
+fn get_imported_tokens_impl() -> GetImportedTokensResponse {
+    let principal = dfn_core::api::caller();
+    STATE.with(|s| s.accounts_store.borrow_mut().get_imported_tokens(principal))
+}
+
 #[export_name = "canister_update get_proposal_payload"]
 pub fn get_proposal_payload() {
     over_async(candid_one, proposals::get_proposal_payload);