Allow account linking via a code

ivan-jukic · ivan-jukic · commit 476d693ece9b · 2025-07-17T20:19:37.000+02:00
---
Adds a feature to create account linking codes. These will be used with
native apps to link existing accounts to the passkeys created on these
devices.
diff --git a/backend/canisters/identity/api/src/queries/get_account_linking_code.rs b/backend/canisters/identity/api/src/queries/get_account_linking_code.rs
@@ -0,0 +1,13 @@
+use candid::{CandidType, Deserialize};
+use serde::Serialize;
+use ts_export::ts_export;
+use types::{AccountLinkingCode, Empty};
+
+pub type Args = Empty;
+
+#[ts_export(identity, get_account_linking_code)]
+#[derive(CandidType, Serialize, Deserialize, Debug)]
+pub enum Response {
+    Success(AccountLinkingCode),
+    NotFound,
+}
diff --git a/backend/canisters/identity/api/src/queries/mod.rs b/backend/canisters/identity/api/src/queries/mod.rs
@@ -1,4 +1,6 @@
 pub mod auth_principals;
 pub mod check_auth_principal_v2;
+pub mod get_account_linking_code;
 pub mod get_delegation;
 pub mod lookup_webauthn_pubkey;
+pub mod verify_account_linking_code;
diff --git a/backend/canisters/identity/api/src/queries/verify_account_linking_code.rs b/backend/canisters/identity/api/src/queries/verify_account_linking_code.rs
@@ -0,0 +1,13 @@
+use candid::{CandidType, Deserialize};
+use serde::Serialize;
+use ts_export::ts_export;
+
+#[ts_export(identity, verify_account_linking_code)]
+#[derive(CandidType, Serialize, Deserialize, Debug)]
+pub struct Args {
+    pub code: String,
+}
+
+#[ts_export(notifications_index, fcm_token_exists)]
+#[derive(CandidType, Serialize, Deserialize, Debug)]
+pub struct Response(pub bool);
diff --git a/backend/canisters/identity/api/src/updates/create_account_linking_code.rs b/backend/canisters/identity/api/src/updates/create_account_linking_code.rs
@@ -0,0 +1,13 @@
+use candid::{CandidType, Deserialize};
+use serde::Serialize;
+use ts_export::ts_export;
+use types::{AccountLinkingCode, Empty};
+
+pub type Args = Empty;
+
+#[ts_export(identity, create_identity)]
+#[derive(CandidType, Serialize, Deserialize, Debug)]
+pub enum Response {
+    Success(AccountLinkingCode),
+    UserNotFound,
+}
diff --git a/backend/canisters/identity/api/src/updates/link_with_account_linking_code.rs b/backend/canisters/identity/api/src/updates/link_with_account_linking_code.rs
@@ -0,0 +1,29 @@
+use crate::WebAuthnKey;
+use candid::{CandidType, Deserialize};
+use serde::Serialize;
+use ts_export::ts_export;
+use types::CanisterId;
+
+#[ts_export(identity, use_account_linking_code)]
+#[derive(CandidType, Serialize, Deserialize, Debug)]
+pub struct Args {
+    pub code: String,
+    #[serde(with = "serde_bytes")]
+    pub public_key: Vec<u8>,
+    pub webauthn_key: Option<WebAuthnKey>,
+}
+
+#[ts_export(identity, use_account_linking_code)]
+#[derive(CandidType, Serialize, Deserialize, Debug)]
+pub enum Response {
+    Success,
+    AlreadyRegistered,
+    AlreadyLinkedToPrincipal,
+    TargetUserNotFound,
+    PublicKeyInvalid(String),
+    OriginatingCanisterInvalid(CanisterId),
+    LinkedIdentitiesLimitReached(u32),
+    LinkingCodeNotFound,
+    LinkingCodeExpired,
+    LinkingCodeInvalid,
+}
diff --git a/backend/canisters/identity/api/src/updates/mod.rs b/backend/canisters/identity/api/src/updates/mod.rs
@@ -1,11 +1,13 @@
 pub mod accept_identity_link_via_qr_code;
 pub mod approve_identity_link;
 pub mod c2c_set_user_ids;
+pub mod create_account_linking_code;
 pub mod create_identity;
 pub mod delete_user;
 pub mod generate_challenge;
 pub mod get_encryption_key;
 pub mod initiate_identity_link;
 pub mod initiate_identity_link_via_qr_code;
+pub mod link_with_account_linking_code;
 pub mod prepare_delegation;
 pub mod remove_identity_link;
diff --git a/backend/canisters/identity/impl/src/lib.rs b/backend/canisters/identity/impl/src/lib.rs
@@ -20,7 +20,8 @@ use serde_bytes::ByteBuf;
 use sha256::sha256;
 use std::cell::RefCell;
 use std::collections::{BTreeMap, HashMap, HashSet};
-use types::{BuildVersion, CanisterId, Cycles, Milliseconds, OCResult, TimestampMillis, Timestamped};
+use types::AccountLinkingCode;
+use types::{BuildVersion, CanisterId, Cycles, Milliseconds, OCResult, TimestampMillis, Timestamped, UserId};
 use utils::env::Environment;
 use x509_parser::prelude::{FromDer, SubjectPublicKeyInfo};
 
@@ -142,6 +143,14 @@ impl RuntimeState {
         }
     }
 
+    pub fn get_user_id_by_caller(&self) -> Option<UserId> {
+        let caller = self.env.caller();
+        self.data
+            .user_principals
+            .get_by_auth_principal(&caller)
+            .and_then(|u| u.user_id)
+    }
+
     pub fn metrics(&self) -> Metrics {
         Metrics {
             heap_memory_used: utils::memory::heap(),
@@ -184,6 +193,8 @@ struct Data {
     rng_seed: [u8; 32],
     challenges: Challenges,
     test_mode: bool,
+    #[serde(default)]
+    account_linking_codes: HashMap<String, (UserId, AccountLinkingCode)>,
 }
 
 impl Data {
@@ -213,6 +224,7 @@ impl Data {
             rng_seed: [0; 32],
             challenges: Challenges::default(),
             test_mode,
+            account_linking_codes: HashMap::new(),
         }
     }
 
diff --git a/backend/canisters/identity/impl/src/model/user_principals.rs b/backend/canisters/identity/impl/src/model/user_principals.rs
@@ -224,6 +224,13 @@ impl UserPrincipals {
         }
     }
 
+    pub fn find_user_principal_by_user_id(&self, user_id: UserId) -> Option<(Principal, Vec<Principal>)> {
+        self.user_principals
+            .iter()
+            .find(|u| u.user_id == Some(user_id))
+            .map(|u| (u.principal, u.auth_principals.clone()))
+    }
+
     // This is O(number of users) so only use in rare cases
     pub fn get_originating_canisters_by_user_id_slow(&self, user_id: UserId) -> Vec<(CanisterId, bool)> {
         if let Some(user_principal) = self.user_principals.iter().find(|u| u.user_id == Some(user_id)) {
diff --git a/backend/canisters/identity/impl/src/queries/get_account_linking_code.rs b/backend/canisters/identity/impl/src/queries/get_account_linking_code.rs
@@ -0,0 +1,33 @@
+use crate::{RuntimeState, read_state};
+use canister_api_macros::query;
+use identity_canister::get_account_linking_code::{Response::*, *};
+use types::{AccountLinkingCode, UserId};
+
+#[query(msgpack = true, candid = true)]
+fn get_account_linking_code(_args: Args) -> Response {
+    read_state(get_account_linking_code_impl)
+}
+
+// Fetches existing account linking code for the caller's user ID.
+fn get_account_linking_code_impl(state: &RuntimeState) -> Response {
+    let now = state.env.now();
+
+    if let Some(user_id) = state.get_user_id_by_caller() {
+        if let Some(alc) = find_account_linking_code_for_user_id(state, &user_id) {
+            if alc.is_valid(now) { Success(alc.clone()) } else { NotFound }
+        } else {
+            NotFound
+        }
+    } else {
+        NotFound
+    }
+}
+
+fn find_account_linking_code_for_user_id(state: &RuntimeState, user_id: &UserId) -> Option<AccountLinkingCode> {
+    state
+        .data
+        .account_linking_codes
+        .iter()
+        .find(|(_, (alc_user_id, _))| alc_user_id == user_id)
+        .map(|(_, (_, alc))| alc.clone())
+}
diff --git a/backend/canisters/identity/impl/src/queries/mod.rs b/backend/canisters/identity/impl/src/queries/mod.rs
@@ -1,5 +1,7 @@
 mod auth_principals;
 mod check_auth_principal_v2;
+mod get_account_linking_code;
 mod get_delegation;
 mod http_request;
 mod lookup_webauthn_pubkey;
+mod verify_account_linking_code;
diff --git a/backend/canisters/identity/impl/src/queries/verify_account_linking_code.rs b/backend/canisters/identity/impl/src/queries/verify_account_linking_code.rs
@@ -0,0 +1,24 @@
+use crate::{RuntimeState, read_state};
+use canister_api_macros::query;
+use identity_canister::verify_account_linking_code::{Args, Response};
+
+#[query(msgpack = true, candid = true)]
+fn verify_account_linking_code(args: Args) -> Response {
+    read_state(|state| verify_account_linking_code_impl(args, state))
+}
+
+// Verify that the account linking code is correct & valid for the caller's user ID.
+fn verify_account_linking_code_impl(args: Args, state: &RuntimeState) -> Response {
+    let now = state.env.now();
+
+    if let Some(user_id) = state.get_user_id_by_caller() {
+        if let Some((alc_user_id, alc)) = state.data.account_linking_codes.get(&args.code) {
+            // Check if the account linking code is valid and matches the user ID.
+            Response(alc.is_valid(now) && alc.value == args.code && user_id == *alc_user_id)
+        } else {
+            Response(false)
+        }
+    } else {
+        Response(false)
+    }
+}
diff --git a/backend/canisters/identity/impl/src/updates/create_account_linking_code.rs b/backend/canisters/identity/impl/src/updates/create_account_linking_code.rs
@@ -0,0 +1,44 @@
+use crate::{RuntimeState, mutate_state};
+use canister_api_macros::update;
+use canister_tracing_macros::trace;
+use identity_canister::create_account_linking_code::{Response::*, *};
+use types::AccountLinkingCode;
+
+#[update(msgpack = true, candid = true)]
+#[trace]
+fn create_account_linking_code(_args: Args) -> Response {
+    mutate_state(create_account_linking_code_impl)
+}
+
+fn create_account_linking_code_impl(state: &mut RuntimeState) -> Response {
+    let now = state.env.now();
+    let caller = state.env.caller();
+
+    if let Some(user_id) = state
+        .data
+        .user_principals
+        .get_by_auth_principal(&caller)
+        .and_then(|u| u.user_id)
+    {
+        // Clean up expired codes - keeps the memory footprint smaller in
+        // exchange for a bit of extra CPU time.
+        state.data.account_linking_codes.retain(|_, (_, alc)| alc.is_valid(now));
+
+        let mut alc = AccountLinkingCode::new(now);
+
+        // Verify that the generated code is unique
+        while state.data.account_linking_codes.contains_key(&alc.value) {
+            alc = AccountLinkingCode::new(now);
+        }
+
+        // Store the code in the state (this part is assumed, as the original code does not specify how to store it)
+        state
+            .data
+            .account_linking_codes
+            .insert(alc.value.clone(), (user_id, alc.clone()));
+
+        Success(alc)
+    } else {
+        UserNotFound
+    }
+}
diff --git a/backend/canisters/identity/impl/src/updates/link_with_account_linking_code.rs b/backend/canisters/identity/impl/src/updates/link_with_account_linking_code.rs
@@ -0,0 +1,88 @@
+use crate::{RuntimeState, VerifyNewIdentityArgs, VerifyNewIdentityError, VerifyNewIdentitySuccess, mutate_state};
+use canister_api_macros::update;
+use canister_tracing_macros::trace;
+use identity_canister::link_with_account_linking_code::{Response::*, *};
+
+// TODO Moved to a shared location, also consider what this means now in
+// context of users using passkeys.
+const MAX_LINKED_IDENTITIES: usize = 10;
+
+#[update(msgpack = true, candid = true)]
+#[trace]
+fn link_with_account_linking_code(args: Args) -> Response {
+    mutate_state(|state| link_with_account_linking_code_impl(args, state))
+}
+
+// TODO making this more secure:
+// - rate limit the number of attempts to link,
+// - lockout after too many attempts
+// - log and audit failed attempts (?)
+// - consider reducing the time window for the linking code to be valid
+fn link_with_account_linking_code_impl(args: Args, state: &mut RuntimeState) -> Response {
+    // Get user ID and account linking code from the state. `args.code` is the
+    // linking code provided by the user.
+    let Some((user_id, linking_code)) = state.data.account_linking_codes.get(&args.code) else {
+        return LinkingCodeNotFound;
+    };
+
+    // Returns user principal, and the list of auth principals linked to that user principal.
+    if let Some((link_to_principal, auth_principals)) = state.data.user_principals.find_user_principal_by_user_id(*user_id) {
+        // Verify the new identity using the provided public key and webauthn key.
+        // This will also check if the caller is already registered.
+        let VerifyNewIdentitySuccess {
+            caller: _,
+            auth_principal,
+            originating_canister,
+            webauthn_key,
+        } = match state.verify_new_identity(VerifyNewIdentityArgs {
+            public_key: args.public_key,
+            webauthn_key: args.webauthn_key,
+            allow_existing_provided_not_linked_to_oc_account: true,
+        }) {
+            Ok(ok) => ok,
+            Err(error) => {
+                return match error {
+                    VerifyNewIdentityError::AlreadyRegistered => AlreadyRegistered,
+                    VerifyNewIdentityError::PublicKeyInvalid(e) => PublicKeyInvalid(e),
+                    VerifyNewIdentityError::OriginatingCanisterInvalid(c) => OriginatingCanisterInvalid(c),
+                };
+            }
+        };
+
+        // Check if the linking code is expired.
+        if !linking_code.is_valid(state.env.now()) {
+            return LinkingCodeExpired;
+        }
+
+        // Check if the linking code is valid.
+        if linking_code.value != args.code {
+            return LinkingCodeInvalid;
+        }
+
+        // Check that the target user principal has not reached the maximum number of linked identities
+        if auth_principals.len() >= MAX_LINKED_IDENTITIES {
+            return LinkedIdentitiesLimitReached(MAX_LINKED_IDENTITIES as u32);
+        }
+
+        // Check if the auth principal is already linked to the target user principal
+        if auth_principals.contains(&auth_principal) {
+            return AlreadyLinkedToPrincipal;
+        }
+
+        state.data.identity_link_requests.push(
+            auth_principal,
+            webauthn_key,
+            originating_canister,
+            false,
+            link_to_principal,
+            state.env.now(),
+        );
+
+        // Remove the linking code from the state, as it has been used.
+        state.data.account_linking_codes.remove(&args.code);
+
+        Success
+    } else {
+        TargetUserNotFound
+    }
+}
diff --git a/backend/canisters/identity/impl/src/updates/mod.rs b/backend/canisters/identity/impl/src/updates/mod.rs
@@ -1,11 +1,13 @@
 pub mod accept_identity_link_via_qr_code;
 pub mod approve_identity_link;
 pub mod c2c_set_user_ids;
+pub mod create_account_linking_code;
 pub mod create_identity;
 pub mod delete_user;
 pub mod generate_challenge;
 pub mod get_encryption_key;
 pub mod initiate_identity_link;
 pub mod initiate_identity_link_via_qr_code;
+pub mod link_with_account_linking_code;
 pub mod prepare_delegation;
 pub mod remove_identity_link;
diff --git a/backend/libraries/types/src/account_linking_code.rs b/backend/libraries/types/src/account_linking_code.rs
@@ -0,0 +1,34 @@
+use candid::CandidType;
+use rand::{Rng, distributions::Alphanumeric};
+use serde::{Deserialize, Serialize};
+use ts_export::ts_export;
+
+// For [a-zA-Z0-9] characters this gives us 62^6 = 56,800,235,584 possible combinations
+pub const OTP_LENGTH: usize = 6;
+pub const OTP_DURATION: u64 = 5 * 60 * 1000; // 5 minutes in milliseconds
+
+#[derive(CandidType, Serialize, Deserialize, Clone, Debug)]
+#[ts_export]
+pub struct AccountLinkingCode {
+    pub value: String,
+    pub expires_at: u64, // timestamp in milliseconds
+}
+
+impl AccountLinkingCode {
+    pub fn new(now: u64) -> Self {
+        let value = rand::thread_rng()
+            .sample_iter(&Alphanumeric)
+            .take(OTP_LENGTH)
+            .map(char::from)
+            .collect::<String>();
+
+        Self {
+            value,
+            expires_at: now + OTP_DURATION,
+        }
+    }
+
+    pub fn is_valid(&self, now: u64) -> bool {
+        now < self.expires_at
+    }
+}
diff --git a/backend/libraries/types/src/lib.rs b/backend/libraries/types/src/lib.rs

Original file line number	Diff line number	Diff line change
`@@ -224,6 +224,13 @@ impl UserPrincipals {`
`224`	`224`	`}`
`225`	`225`	`}`
`226`	`226`
	`227`	`+ pub fn find_user_principal_by_user_id(&self, user_id: UserId) -> Option<(Principal, Vec<Principal>)> {`
	`228`	`+ self.user_principals`
	`229`	`+ .iter()`
	`230`	`+ .find(\|u\| u.user_id == Some(user_id))`
	`231`	`+ .map(\|u\| (u.principal, u.auth_principals.clone()))`
	`232`	`+ }`
	`233`	`+`
`227`	`234`	`// This is O(number of users) so only use in rare cases`
`228`	`235`	`pub fn get_originating_canisters_by_user_id_slow(&self, user_id: UserId) -> Vec<(CanisterId, bool)> {`
`229`	`236`	`if let Some(user_principal) = self.user_principals.iter().find(\|u\| u.user_id == Some(user_id)) {`