Add missing principal links to the 2 way mapping (#6562)

open-chat-labs · Oct 11, 2024 · 92c6c88 · 92c6c88
1 parent 1342737
commit 92c6c88
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 1 deletion.
diff --git a/backend/canisters/identity/CHANGELOG.md b/backend/canisters/identity/CHANGELOG.md
@@ -10,6 +10,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 - Switch from using `canister_sig_util` to `ic-canister-sig-creation` ([#6537](https://github.com/open-chat-labs/open-chat/pull/6537)) 
 
+### Fixed
+
+- Add missing principal links to the 2 way mapping ([#6562](https://github.com/open-chat-labs/open-chat/pull/6562))
+
 ## [[2.0.1373](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1373-identity)] - 2024-10-07
 
 ### Changed

diff --git a/backend/canisters/identity/impl/src/lifecycle/post_upgrade.rs b/backend/canisters/identity/impl/src/lifecycle/post_upgrade.rs
@@ -1,6 +1,6 @@
 use crate::lifecycle::{init_env, init_state};
 use crate::memory::get_upgrades_memory;
-use crate::Data;
+use crate::{mutate_state, Data};
 use canister_logger::LogEntry;
 use canister_tracing_macros::trace;
 use ic_cdk::post_upgrade;
@@ -23,5 +23,10 @@ fn post_upgrade(args: Args) {
     init_cycles_dispenser_client(data.cycles_dispenser_canister_id, data.test_mode);
     init_state(env, data, args.wasm_version);
 
+    mutate_state(|state| {
+        let missing_principal_links = state.data.user_principals.add_missing_principal_links();
+        info!("Missing principal links: {missing_principal_links}");
+    });
+
     info!(version = %args.wasm_version, "Post-upgrade complete");
 }
diff --git a/backend/canisters/identity/impl/src/model/user_principals.rs b/backend/canisters/identity/impl/src/model/user_principals.rs
@@ -1,6 +1,7 @@
 use candid::Principal;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
+use tracing::info;
 use types::{is_default, CanisterId, PushIfNotContains, UserId};
 
 #[derive(Serialize, Deserialize, Default)]
@@ -39,6 +40,21 @@ struct AuthPrincipalInternal {
 }
 
 impl UserPrincipals {
+    pub fn add_missing_principal_links(&mut self) -> u32 {
+        let mut count = 0;
+        for (principal, details) in self.auth_principals.iter() {
+            if let Some(user_principal) = self.user_principals.get_mut(details.user_principal_index as usize) {
+                if !user_principal.auth_principals.contains(principal) {
+                    user_principal.auth_principals.push(*principal);
+                    let user_id = user_principal.user_id.map(|u| u.to_string());
+                    info!(user_id, "Missing principal link added");
+                    count += 1;
+                }
+            }
+        }
+        count
+    }
+
     pub fn push(
         &mut self,
         index: u32,