Limit number of imported tokens per account

dfinity · Jul 25, 2024 · e75e1e7 · e75e1e7
1 parent 144f503
commit e75e1e7
Show file tree

Hide file tree

Showing 3 changed files with 108 additions and 10 deletions.
diff --git a/rs/backend/nns-dapp.did b/rs/backend/nns-dapp.did
@@ -206,6 +206,7 @@ type SetImportedTokensResponse =
     variant {
         Ok;
         AccountNotFound;
+        TooManyImportedTokens: record{limit: int32};
     };
 
 type GetImportedTokensResponse =

diff --git a/rs/backend/src/accounts_store.rs b/rs/backend/src/accounts_store.rs
@@ -37,6 +37,8 @@ use self::schema::SchemaLabel;
 /// accounts we avoid some complications.
 const PRE_MIGRATION_LIMIT: u64 = 300_000;
 
+const MAX_IMPORTED_TOKENS: i32 = 10;
+
 /// Accounts, transactions and related data.
 ///
 /// Note: Some monitoring fields are not included in the `Eq` and `PartialEq` implementations.  Additionally, please note
@@ -232,6 +234,7 @@ pub struct ImportedTokens {
 pub enum SetImportedTokensResponse {
     Ok,
     AccountNotFound,
+    TooManyImportedTokens { limit: i32 },
 }
 
 #[derive(CandidType, Debug, PartialEq)]
@@ -775,6 +778,11 @@ impl AccountsStore {
         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;

diff --git a/rs/backend/src/accounts_store/tests.rs b/rs/backend/src/accounts_store/tests.rs
@@ -783,22 +783,24 @@ fn detach_canister_canister_not_found() {
 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();
+    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![ImportedToken {
-                    ledger_canister_id: ledger_canister_id,
-                    index_canister_id: Some(index_canister_id),
-                }],
+                imported_tokens: vec![imported_token.clone()],
             },
         ),
         SetImportedTokensResponse::Ok
@@ -807,11 +809,75 @@ fn set_and_get_imported_tokens() {
     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),
-            }],
+            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],
+        })
+    );
+}
+
+#[test]
+fn set_and_get_10_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: Vec<ImportedToken> = (0..10)
+        .map(|i| ImportedToken {
+            ledger_canister_id: PrincipalId::new_user_test_id(i as u64),
+            index_canister_id: Some(PrincipalId::new_user_test_id(i as u64 + 100)),
         })
+        .collect();
+
+    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 })
     );
 }
 
@@ -825,6 +891,29 @@ fn set_imported_tokens_account_not_found() {
     );
 }
 
+#[test]
+fn set_imported_tokens_too_many() {
+    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);
+
+    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(); 11],
+            },
+        ),
+        SetImportedTokensResponse::TooManyImportedTokens { limit: 10 }
+    );
+}
+
 #[test]
 fn get_imported_tokens_account_not_found() {
     let mut store = setup_test_store();