do the same in the upgrade what Motoko does -- overwrite the key name

Author: altkdf

backend/rs/canisters/ic_vetkeys_encrypted_maps_canister/src/lib.rs

@@ -7,7 +7,7 @@ use ic_stable_structures::memory_manager::{MemoryId, MemoryManager, VirtualMemor
 use ic_stable_structures::storable::Blob;
 use ic_stable_structures::DefaultMemoryImpl;
 use ic_vetkeys::encrypted_maps::{EncryptedMapData, EncryptedMaps, VetKey, VetKeyVerificationKey};
-use ic_vetkeys::types::{AccessRights, ByteBuf, EncryptedMapValue, TransportKey};
+use ic_vetkeys::types::{AccessRights, ByteBuf, EncryptedMapValue, KeyManagerConfig, TransportKey};
 
 type Memory = VirtualMemory<DefaultMemoryImpl>;
 type MapId = (Principal, ByteBuf);
@@ -38,21 +38,44 @@ fn init(key_name: String) {
 }
 
 #[post_upgrade]
-fn post_upgrade() {
-    // Note that the value of the key_id is not important, as it is used only if the memory not yet initialized.
-    let key_id = VetKDKeyId {
+fn post_upgrade(key_name: String) {
+    // Note that the value of the `dummy_key_id` and `domain_separator` is not important, as it is used only if the memory not yet initialized, which is not the case here.
+    let dummy_key_id = VetKDKeyId {
         curve: VetKDCurve::Bls12_381_G2,
         name: "dummy_key_name".to_string(),
     };
+    let domain_separator = "dummy_domain_separator";
     ENCRYPTED_MAPS.with_borrow_mut(|encrypted_maps| {
         encrypted_maps.replace(EncryptedMaps::init(
-            "dummy_domain_separator",
-            key_id,
+            domain_separator,
+            dummy_key_id.clone(),
             id_to_memory(0),
             id_to_memory(1),
             id_to_memory(2),
             id_to_memory(3),
-        ))
+        ));
+
+        let config: &KeyManagerConfig = encrypted_maps.as_mut().unwrap().key_manager.config.get();
+        assert_ne!(config.key_id, dummy_key_id);
+        assert_ne!(config.domain_separator, domain_separator);
+
+        let cloned_domain_separator = config.domain_separator.clone();
+
+        let new_key_id = VetKDKeyId {
+            curve: VetKDCurve::Bls12_381_G2,
+            name: key_name,
+        };
+
+        encrypted_maps
+            .as_mut()
+            .unwrap()
+            .key_manager
+            .config
+            .set(KeyManagerConfig {
+                domain_separator: cloned_domain_separator,
+                key_id: new_key_id,
+            })
+            .unwrap();
     });
 }
 

backend/rs/canisters/ic_vetkeys_encrypted_maps_canister/tests/tests.rs

@@ -901,7 +901,7 @@ fn should_survive_canister_upgrade() {
         .upgrade_canister(
             env.example_canister_id,
             wasm_bytes,
-            encode_one("").unwrap(),
+            encode_one("dfx_test_key").unwrap(),
             None,
         )
         .unwrap();

backend/rs/canisters/ic_vetkeys_manager_canister/src/lib.rs

@@ -7,7 +7,7 @@ use ic_stable_structures::memory_manager::{MemoryId, MemoryManager, VirtualMemor
 use ic_stable_structures::storable::Blob;
 use ic_stable_structures::DefaultMemoryImpl;
 use ic_vetkeys::key_manager::{KeyManager, VetKey, VetKeyVerificationKey};
-use ic_vetkeys::types::{AccessRights, ByteBuf, TransportKey};
+use ic_vetkeys::types::{AccessRights, ByteBuf, KeyManagerConfig, TransportKey};
 
 type Memory = VirtualMemory<DefaultMemoryImpl>;
 
@@ -36,21 +36,41 @@ fn init(key_name: String) {
 }
 
 #[post_upgrade]
-fn post_upgrade() {
-    // Note that the value of the key_id is not important, as it is used only if the memory not yet initialized.
+fn post_upgrade(key_name: String) {
+    // Note that the value of the `dummy_key_id` and `domain_separator` is not important, as it is used only if the memory not yet initialized, which is not the case here.
     let key_id = VetKDKeyId {
         curve: VetKDCurve::Bls12_381_G2,
         name: "dummy_key_name".to_string(),
     };
+    let domain_separator = "dummy_domain_separator";
     KEY_MANAGER.with_borrow_mut(|km| {
         km.replace(KeyManager::init(
-            "dummy_domain_separator",
-            key_id,
+            domain_separator,
+            key_id.clone(),
             id_to_memory(0),
             id_to_memory(1),
             id_to_memory(2),
-        ))
-    });
+        ));
+
+        let config: &KeyManagerConfig = km.as_mut().unwrap().config.get();
+        assert_ne!(config.key_id, key_id);
+        assert_ne!(config.domain_separator, domain_separator);
+
+        let cloned_domain_separator = config.domain_separator.clone();
+
+        let new_key_id = VetKDKeyId {
+            curve: VetKDCurve::Bls12_381_G2,
+            name: key_name,
+        };
+        km.as_mut()
+            .unwrap()
+            .config
+            .set(KeyManagerConfig {
+                domain_separator: cloned_domain_separator,
+                key_id: new_key_id,
+            })
+            .unwrap();
+    })
 }
 
 #[query]

backend/rs/canisters/ic_vetkeys_manager_canister/tests/tests.rs

@@ -532,7 +532,7 @@ fn should_survive_canister_upgrade() {
         .upgrade_canister(
             env.example_canister_id,
             wasm_bytes,
-            encode_one("").unwrap(),
+            encode_one("dfx_test_key").unwrap(),
             None,
         )
         .unwrap();