Add canister method to store imported tokens

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,27 @@ 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;
+    };
+
+type GetImportedTokensResponse =
+    variant {
+        Ok: ImportedTokens;
+        AccountNotFound;
+    };
+
 service: (opt Config) -> {
     get_account: () -> (GetAccountResponse) query;
     add_account: () -> (AccountIdentifier);
@@ -216,4 +237,7 @@ service: (opt Config) -> {
 
     // Methods available in the test build only:
     get_toy_account: (nat64) -> (GetAccountResponse) query;
+
+    set_imported_tokens: (ImportedTokens) -> (SetImportedTokensResponse);
+    get_imported_tokens: () -> (GetImportedTokensResponse) query;
 }

rs/backend/src/accounts_store.rs

@@ -154,6 +154,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 +217,29 @@ 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,
+}
+
+#[derive(CandidType, Debug, PartialEq)]
+pub enum GetImportedTokensResponse {
+    Ok(ImportedTokens),
+    AccountNotFound,
+}
+
 #[derive(Copy, Clone, CandidType, Deserialize, Debug, Eq, PartialEq)]
 pub enum TransactionType {
     Burn,
@@ -746,6 +770,31 @@ impl AccountsStore {
         }
     }
 
+    pub fn set_imported_tokens(
+        &mut self,
+        caller: PrincipalId,
+        new_imported_tokens: ImportedTokens,
+    ) -> SetImportedTokensResponse {
+        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 +1165,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,62 @@ 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::from_str(TEST_ACCOUNT_2).unwrap();
+    let index_canister_id = PrincipalId::from_str(TEST_ACCOUNT_3).unwrap();
+
+    assert_eq!(
+        store.get_imported_tokens(principal),
+        GetImportedTokensResponse::Ok(ImportedTokens::default())
+    );
+
+    assert_eq!(
+        store.set_imported_tokens(
+            principal,
+            ImportedTokens {
+                imported_tokens: vec![ImportedToken {
+                    ledger_canister_id: ledger_canister_id,
+                    index_canister_id: Some(index_canister_id),
+                }],
+            },
+        ),
+        SetImportedTokensResponse::Ok
+    );
+
+    assert_eq!(
+        store.get_imported_tokens(principal),
+        GetImportedTokensResponse::Ok(ImportedTokens {
+            imported_tokens: vec![ImportedToken {
+                ledger_canister_id: ledger_canister_id,
+                index_canister_id: Some(index_canister_id),
+            }],
+        })
+    );
+}
+
+#[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 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);