fix: remove old ks in memory (#80)

consenlabs · Mar 25, 2024 · 4b97465 · 4b97465
1 parent 32229dd
commit 4b97465
Show file tree

Hide file tree

Showing 4 changed files with 136 additions and 2 deletions.
diff --git a/token-core/tcx/src/handler.rs b/token-core/tcx/src/handler.rs
@@ -727,7 +727,11 @@ pub(crate) fn delete_keystore(data: &[u8]) -> Result<Vec<u8>> {
         delete_keystore_file(&param.id)?;
         map.remove(&param.id);
 
-        remove_old_keystore_by_id(&param.id.clone());
+        if let Some(file_ids) = remove_old_keystore_by_id(&param.id.clone()) {
+            for file_id in file_ids {
+                map.remove(&file_id);
+            }
+        }
 
         let rsp = GeneralResult {
             is_success: true,

diff --git a/token-core/tcx/src/migration.rs b/token-core/tcx/src/migration.rs
@@ -40,7 +40,7 @@ fn read_migrated_map() -> (String, HashMap<String, Vec<String>>) {
     (migrated_file, map)
 }
 
-pub fn remove_old_keystore_by_id(id: &str) {
+pub fn remove_old_keystore_by_id(id: &str) -> Option<Vec<String>> {
     let legacy_file_dir = {
         let dir = LEGACY_WALLET_FILE_DIR.read();
         dir.to_string()
@@ -51,6 +51,7 @@ pub fn remove_old_keystore_by_id(id: &str) {
     let migrated_file = result.0;
     let mut map = result.1;
 
+    let marked_files = map.get(id).and_then(|x| Some(x.to_vec())).clone();
     if let Some(files) = map.get(id) {
         for file_id in files.iter() {
             let mut file_path = format!("{}/{}.json", legacy_file_dir, file_id);
@@ -72,6 +73,7 @@ pub fn remove_old_keystore_by_id(id: &str) {
     } else {
         fs::remove_file(&migrated_file);
     }
+    marked_files
 }
 
 fn mark_keystore_as_migrated(legacy_file_id: &str, migrated_file_id: &str) {

diff --git a/token-core/tcx/tests/migration_test.rs b/token-core/tcx/tests/migration_test.rs
@@ -854,6 +854,96 @@ fn test_migrate_duplicate_then_delete_keystore() {
     );
 }
 
+#[test]
+#[serial]
+fn test_migrate_duplicate_delete_all_keystore_migrate_again() {
+    setup_test("../test-data/migrate-duplication-fixtures");
+    let param = MigrateKeystoreParam {
+        id: "300b42bc-0948-4734-82cb-4293dfeeefd2".to_string(),
+        network: "TESTNET".to_string(),
+        key: Some(migrate_keystore_param::Key::Password(
+            TEST_PASSWORD.to_string(),
+        )),
+    };
+    call_api("migrate_keystore", param).unwrap();
+    let param = ExportMnemonicParam {
+        id: "300b42bc-0948-4734-82cb-4293dfeeefd2".to_string(),
+        key: Some(export_mnemonic_param::Key::Password(
+            TEST_PASSWORD.to_string(),
+        )),
+    };
+    let ret = call_api("export_mnemonic", param).unwrap();
+    let exported = ExportMnemonicResult::decode(ret.as_slice())
+        .unwrap()
+        .mnemonic;
+    assert_eq!(OTHER_MNEMONIC, exported);
+    // CKB imported 300b42bc-0948-4734-82cb-4293dfeeefd2
+    // 9b696367-69c1-4cfe-8325-e5530399fc3f
+    let param = MigrateKeystoreParam {
+        id: "9b696367-69c1-4cfe-8325-e5530399fc3f".to_string(),
+        network: "TESTNET".to_string(),
+        key: Some(migrate_keystore_param::Key::Password(
+            TEST_PASSWORD.to_string(),
+        )),
+    };
+    call_api("migrate_keystore", param).unwrap();
+    let param = ExportMnemonicParam {
+        id: "9b696367-69c1-4cfe-8325-e5530399fc3f".to_string(),
+        key: Some(export_mnemonic_param::Key::Password(
+            TEST_PASSWORD.to_string(),
+        )),
+    };
+    let ret = call_api("export_mnemonic", param).unwrap();
+    let exported = ExportMnemonicResult::decode(ret.as_slice())
+        .unwrap()
+        .mnemonic;
+    assert_eq!(OTHER_MNEMONIC, exported);
+
+    let param = WalletKeyParam {
+        id: "300b42bc-0948-4734-82cb-4293dfeeefd2".to_string(),
+        key: Some(wallet_key_param::Key::Password(TEST_PASSWORD.to_string())),
+    };
+    let ret = call_api("delete_keystore", param).unwrap();
+
+    let param = MigrateKeystoreParam {
+        id: "792a0051-16d7-44a7-921a-9b4a0c893b8f".to_string(),
+        network: "TESTNET".to_string(),
+        key: Some(migrate_keystore_param::Key::Password(
+            TEST_PASSWORD.to_string(),
+        )),
+    };
+    call_api("migrate_keystore", param).unwrap();
+    let param = ExportMnemonicParam {
+        id: "792a0051-16d7-44a7-921a-9b4a0c893b8f".to_string(),
+        key: Some(export_mnemonic_param::Key::Password(
+            TEST_PASSWORD.to_string(),
+        )),
+    };
+    let ret = call_api("export_mnemonic", param).unwrap();
+    let exported = ExportMnemonicResult::decode(ret.as_slice())
+        .unwrap()
+        .mnemonic;
+    assert_eq!(OTHER_MNEMONIC, exported);
+
+    assert_eq!(
+        Path::new("/tmp/token-core-x/wallets/300b42bc-0948-4734-82cb-4293dfeeefd2.json").exists(),
+        false
+    );
+    assert_eq!(
+        Path::new("/tmp/token-core-x/wallets/9b696367-69c1-4cfe-8325-e5530399fc3f").exists(),
+        false
+    );
+    assert_eq!(
+        Path::new("/tmp/token-core-x/wallets/_migrated").exists(),
+        false
+    );
+
+    assert_eq!(
+        Path::new("/tmp/token-core-x/wallets/792a0051-16d7-44a7-921a-9b4a0c893b8f").exists(),
+        true
+    );
+}
+
 #[test]
 #[serial]
 fn test_migrate_android_old_eos_private_keystore() {

diff --git a/...-core/test-data/migrate-duplication-fixtures/wallets/792a0051-16d7-44a7-921a-9b4a0c893b8f b/...-core/test-data/migrate-duplication-fixtures/wallets/792a0051-16d7-44a7-921a-9b4a0c893b8f
@@ -0,0 +1,38 @@
+{
+  "id": "792a0051-16d7-44a7-921a-9b4a0c893b8f",
+  "mnemonicPath": "m/44'/60'/0'/0/1",
+  "encMnemonic": {
+    "nonce": "27d0fbb70bd03ce0e96824c0f88bb6ad",
+    "encStr": "b8e34c4ff5ee92f80cd893244c89be1ad75614550b4fde40eb81a214c4676623ec5f92076f495dc3204ab06533cccabbb218b01fbaa47255b56d7c17543aa3462d2f246169fb3377e422"
+  },
+  "imTokenMeta": {
+    "source": "MNEMONIC",
+    "passwordHint": "",
+    "backup": [],
+    "chain": "ETHEREUM",
+    "mode": "normal",
+    "network": "TESTNET",
+    "segWit": "NONE",
+    "version": "iOS-2.14.1.1742",
+    "timestamp": "1703217176.7788029",
+    "name": "Imported 12"
+  },
+  "crypto": {
+    "cipher": "aes-128-ctr",
+    "kdf": "scrypt",
+    "ciphertext": "a3f81f417cf96399c178fd18c058cf0cdc1023116e293c3c6fe5aa8a58f972f9",
+    "mac": "42f62a4647194f4efe78a07e8517ef9f4974b4bfb71524ed8dc8282f36780ce5",
+    "kdfparams": {
+      "salt": "59663950940345f9e69183f11f74319db35646ae7cd5be2f9f9bc478c07b4f1e",
+      "n": 262144,
+      "dklen": 32,
+      "r": 8,
+      "p": 1
+    },
+    "cipherparams": {
+      "iv": "4f64c9458e04b387dad5c7d0eaab8c91"
+    }
+  },
+  "address": "7152bcad819b084d57179e293d2765ffa0109e04",
+  "version": 3
+}