Run check to ensure auth principal <-> user principal maps are in sync (

#6512)
open-chat-labs · Oct 5, 2024 · 4bb9da1 · 4bb9da1
1 parent d8cca2b
commit 4bb9da1
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 5 deletions.
diff --git a/backend/canisters/identity/CHANGELOG.md b/backend/canisters/identity/CHANGELOG.md
@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Increase max stable memory read / write buffer size ([#6440](https://github.com/open-chat-labs/open-chat/pull/6440))
 - Extract certificate handling code into `identity_utils` ([#6459](https://github.com/open-chat-labs/open-chat/pull/6459)) 
 - Add serde default attribute in preparation for skipping serialization if default ([#6475](https://github.com/open-chat-labs/open-chat/pull/6475))
+- Run check to ensure auth principal <-> user principal maps are in sync ([#6512](https://github.com/open-chat-labs/open-chat/pull/6512))
 
 ## [[2.0.1344](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1344-identity)] - 2024-09-10
 

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::{read_state, Data};
 use canister_logger::LogEntry;
 use canister_tracing_macros::trace;
 use ic_cdk::post_upgrade;
@@ -24,4 +24,9 @@ fn post_upgrade(args: Args) {
     init_state(env, data, args.wasm_version);
 
     info!(version = %args.wasm_version, "Post-upgrade complete");
+
+    read_state(|state| {
+        let missing_principal_links = state.data.user_principals.count_missing_principal_links();
+        info!("Missing principal links: {missing_principal_links}");
+    })
 }
diff --git a/backend/canisters/identity/impl/src/model/user_principals.rs b/backend/canisters/identity/impl/src/model/user_principals.rs
@@ -1,7 +1,7 @@
 use candid::Principal;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
-use types::{is_default, CanisterId, UserId};
+use types::{is_default, CanisterId, PushIfNotContains, UserId};
 
 #[derive(Serialize, Deserialize, Default)]
 pub struct UserPrincipals {
@@ -39,6 +39,18 @@ struct AuthPrincipalInternal {
 }
 
 impl UserPrincipals {
+    pub fn count_missing_principal_links(&self) -> u32 {
+        let mut count = 0;
+        for (principal, details) in self.auth_principals.iter() {
+            if let Some(user_principal) = self.user_principals.get(details.user_principal_index as usize) {
+                if !user_principal.auth_principals.contains(principal) {
+                    count += 1;
+                }
+            }
+        }
+        count
+    }
+
     pub fn push(
         &mut self,
         index: u32,
@@ -79,9 +91,7 @@ impl UserPrincipals {
         {
             false
         } else if let Some(user_principal) = self.user_principals.get_mut(user_principal_index as usize) {
-            if !user_principal.auth_principals.contains(&new_principal) {
-                user_principal.auth_principals.push(new_principal);
-            }
+            user_principal.auth_principals.push_if_not_contains(new_principal);
             self.auth_principals.insert(
                 new_principal,
                 AuthPrincipalInternal {