Sync userIds to Identity canister (#6027)

Cargo.toml

@@ -34,6 +34,7 @@ members = [
     "backend/canisters/group_index/client",
     "backend/canisters/group_index/impl",
     "backend/canisters/identity/api",
+    "backend/canisters/identity/c2c_client",
     "backend/canisters/identity/impl",
     "backend/canisters/local_group_index/api",
     "backend/canisters/local_group_index/c2c_client",

backend/canisters/identity/CHANGELOG.md

@@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## [unreleased]
 
+### Added
+
+- Sync userIds to Identity canister ([#6027](https://github.com/open-chat-labs/open-chat/pull/6027))
+
 ## [[2.0.1209](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1209-identity)] - 2024-06-20
 
 ### Changed

backend/canisters/identity/api/src/updates/c2c_sync_legacy_user_principals.rs → backend/canisters/identity/api/src/updates/c2c_set_user_ids.rs

@@ -1,9 +1,10 @@
 use candid::{CandidType, Principal};
 use serde::{Deserialize, Serialize};
+use types::UserId;
 
 #[derive(CandidType, Serialize, Deserialize, Debug)]
 pub struct Args {
-    pub principals: Vec<Principal>,
+    pub users: Vec<(Principal, Option<UserId>)>,
 }
 
 #[derive(CandidType, Serialize, Deserialize, Debug)]

backend/canisters/identity/api/src/updates/mod.rs

@@ -1,4 +1,4 @@
-pub mod c2c_sync_legacy_user_principals;
+pub mod c2c_set_user_ids;
 pub mod create_identity;
 pub mod generate_challenge;
 pub mod prepare_delegation;

backend/canisters/identity/c2c_client/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "identity_canister_c2c_client"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+canister_client = { path = "../../../libraries/canister_client" }
+identity_canister = { path = "../api" }
+ic-cdk = { workspace = true }
+msgpack = { path = "../../../libraries/msgpack" }
+tracing = { workspace = true }
+types = { path = "../../../libraries/types" }

backend/canisters/identity/c2c_client/src/lib.rs

@@ -0,0 +1,7 @@
+use canister_client::generate_c2c_call;
+use identity_canister::*;
+
+// Queries
+
+// Updates
+generate_c2c_call!(c2c_set_user_ids);

backend/canisters/identity/impl/Cargo.toml

@@ -23,6 +23,7 @@ ic-cdk-timers = { workspace = true }
 ic-certification = { workspace = true }
 ic-stable-structures = { workspace = true }
 identity_canister = { path = "../api" }
+msgpack = { path = "../../../libraries/msgpack" }
 rand = { workspace = true }
 serde = { workspace = true }
 serde_bytes = { workspace = true }

backend/canisters/identity/impl/src/guards.rs

@@ -0,0 +1,9 @@
+use crate::read_state;
+
+pub fn caller_is_user_index_canister() -> Result<(), String> {
+    if read_state(|state| state.is_caller_user_index_canister()) {
+        Ok(())
+    } else {
+        Err("Caller is not the user_index canister".to_string())
+    }
+}

backend/canisters/identity/impl/src/lib.rs

@@ -14,6 +14,7 @@ use std::collections::{HashMap, HashSet};
 use types::{BuildVersion, CanisterId, Cycles, Hash, TimestampMillis, Timestamped};
 use utils::env::Environment;
 
+mod guards;
 mod hash;
 mod lifecycle;
 mod memory;
@@ -37,6 +38,11 @@ impl RuntimeState {
         RuntimeState { env, data }
     }
 
+    pub fn is_caller_user_index_canister(&self) -> bool {
+        let caller = self.env.caller();
+        self.data.user_index_canister_id == caller
+    }
+
     pub fn der_encode_canister_sig_key(&self, seed: [u8; 32]) -> Vec<u8> {
         let canister_id = self.env.canister_id();
         CanisterSigPublicKey::new(canister_id, seed.to_vec()).to_der()

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::CanisterId;
+use types::{CanisterId, UserId};
 
 #[derive(Serialize, Deserialize, Default)]
 pub struct UserPrincipals {
@@ -23,6 +23,8 @@ struct UserPrincipalInternal {
     principal: Principal,
     #[serde(rename = "a")]
     auth_principals: Vec<Principal>,
+    #[serde(rename = "u", default, skip_serializing_if = "Option::is_none")]
+    user_id: Option<UserId>,
 }
 
 #[derive(Serialize, Deserialize)]
@@ -41,6 +43,7 @@ impl UserPrincipals {
         self.user_principals.push(UserPrincipalInternal {
             principal,
             auth_principals: vec![auth_principal],
+            user_id: None,
         });
         self.auth_principals.insert(
             auth_principal,
@@ -74,6 +77,17 @@ impl UserPrincipals {
         &self.originating_canisters
     }
 
+    // This is O(number of users) so we may need to revisit this in the future, but it is only
+    // called once per user so is fine for now.
+    pub fn set_user_id(&mut self, principal: Principal, user_id: Option<UserId>) -> bool {
+        if let Some(user) = self.user_principals.iter_mut().find(|u| u.principal == principal) {
+            user.user_id = user_id;
+            true
+        } else {
+            false
+        }
+    }
+
     fn user_principal_by_index(&self, user_principal_index: u32) -> Option<UserPrincipal> {
         self.user_principals
             .get(usize::try_from(user_principal_index).unwrap())

backend/canisters/identity/impl/src/updates/c2c_set_user_ids.rs

@@ -0,0 +1,23 @@
+use crate::guards::caller_is_user_index_canister;
+use crate::{mutate_state, RuntimeState};
+use canister_api_macros::update_msgpack;
+use canister_tracing_macros::trace;
+use identity_canister::c2c_set_user_ids::{Response::*, *};
+
+#[update_msgpack(guard = "caller_is_user_index_canister")]
+#[trace]
+fn c2c_set_user_ids(args: Args) -> Response {
+    // This function runs in O(number of users registered x batch size),
+    // so we need to ensure each batch is fairly small
+    assert!(args.users.len() <= 100);
+
+    mutate_state(|state| c2c_set_user_ids_impl(args, state))
+}
+
+fn c2c_set_user_ids_impl(args: Args, state: &mut RuntimeState) -> Response {
+    for (principal, user_id) in args.users {
+        state.data.user_principals.set_user_id(principal, user_id);
+    }
+
+    Success
+}

backend/canisters/identity/impl/src/updates/mod.rs

@@ -1,3 +1,4 @@
+pub mod c2c_set_user_ids;
 pub mod create_identity;
 pub mod generate_challenge;
 pub mod prepare_delegation;

backend/canisters/user_index/CHANGELOG.md

@@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## [unreleased]
 
+### Added
+
+- Sync userIds to Identity canister ([#6027](https://github.com/open-chat-labs/open-chat/pull/6027))
+
 ## [[2.0.1235](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1235-user_index)] - 2024-07-11
 
 ### Added

backend/canisters/user_index/impl/Cargo.toml

@@ -35,6 +35,8 @@ ic-stable-structures = { workspace = true }
 ic-verifiable-credentials = { workspace = true }
 icrc_ledger_canister_c2c_client = { path = "../../../external_canisters/icrc_ledger/c2c_client" }
 icrc-ledger-types = { workspace = true }
+identity_canister = { path = "../../identity/api" }
+identity_canister_c2c_client = { path = "../../identity/c2c_client" }
 itertools = { workspace = true }
 local_user_index_canister = { path = "../../local_user_index/api" }
 local_user_index_canister_c2c_client = { path = "../../local_user_index/c2c_client" }

backend/canisters/user_index/impl/src/jobs/mod.rs

@@ -2,6 +2,7 @@ use crate::RuntimeState;
 pub mod make_pending_payments;
 pub mod submit_message_to_modclub;
 pub mod sync_events_to_local_user_index_canisters;
+pub mod sync_users_to_identity_canister;
 pub mod sync_users_to_storage_index;
 pub mod upgrade_canisters;
 
@@ -10,5 +11,6 @@ pub(crate) fn start(state: &RuntimeState) {
     submit_message_to_modclub::start_job_if_required(state);
     sync_events_to_local_user_index_canisters::start_job_if_required(state);
     sync_users_to_storage_index::start_job_if_required(state);
+    sync_users_to_identity_canister::start_job_if_required(state);
     upgrade_canisters::start_job_if_required(state);
 }

backend/canisters/user_index/impl/src/jobs/sync_users_to_identity_canister.rs

@@ -0,0 +1,71 @@
+use crate::{mutate_state, RuntimeState};
+use candid::Principal;
+use ic_cdk_timers::TimerId;
+use std::cell::Cell;
+use std::cmp::min;
+use std::time::Duration;
+use tracing::trace;
+use types::{CanisterId, UserId};
+
+const BATCH_SIZE: usize = 100;
+
+thread_local! {
+    static TIMER_ID: Cell<Option<TimerId>> = Cell::default();
+}
+
+pub(crate) fn start_job_if_required(state: &RuntimeState) -> bool {
+    if TIMER_ID.get().is_none() && !state.data.identity_canister_user_sync_queue.is_empty() {
+        let timer_id = ic_cdk_timers::set_timer(Duration::ZERO, run);
+        TIMER_ID.set(Some(timer_id));
+        true
+    } else {
+        false
+    }
+}
+
+pub(crate) fn try_run_now(state: &mut RuntimeState) -> bool {
+    if let Some((canister_id, users)) = next_batch(state) {
+        if let Some(timer_id) = TIMER_ID.take() {
+            ic_cdk_timers::clear_timer(timer_id);
+        }
+        ic_cdk::spawn(sync_users(canister_id, users));
+        true
+    } else {
+        false
+    }
+}
+
+fn run() {
+    trace!("'sync_users_to_identity_canister' job running");
+    TIMER_ID.set(None);
+
+    if let Some((canister_id, users)) = mutate_state(next_batch) {
+        ic_cdk::spawn(sync_users(canister_id, users));
+    }
+}
+
+#[allow(clippy::type_complexity)]
+fn next_batch(state: &mut RuntimeState) -> Option<(CanisterId, Vec<(Principal, Option<UserId>)>)> {
+    let count = min(state.data.identity_canister_user_sync_queue.len(), BATCH_SIZE);
+    if count == 0 {
+        return None;
+    }
+
+    let batch: Vec<_> = state.data.identity_canister_user_sync_queue.drain(..count).collect();
+
+    Some((state.data.identity_canister_id, batch))
+}
+
+async fn sync_users(identity_canister_id: CanisterId, users: Vec<(Principal, Option<UserId>)>) {
+    let args = identity_canister::c2c_set_user_ids::Args { users: users.clone() };
+    let success = identity_canister_c2c_client::c2c_set_user_ids(identity_canister_id, &args)
+        .await
+        .is_ok();
+
+    mutate_state(|state| {
+        if !success {
+            state.data.identity_canister_user_sync_queue.extend(users);
+        }
+        start_job_if_required(state);
+    });
+}