Skip to content

Commit 254e6c6

Browse files
committed
Update to the overall functionality, and added integration test
1 parent 476d693 commit 254e6c6

File tree

16 files changed

+302
-129
lines changed

16 files changed

+302
-129
lines changed

backend/canisters/identity/api/src/main.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,9 @@ fn main() {
1010
generate_ts_method!(identity, auth_principals);
1111
generate_ts_method!(identity, check_auth_principal_v2);
1212
generate_ts_method!(identity, get_delegation);
13+
generate_ts_method!(identity, get_account_linking_code);
1314
generate_ts_method!(identity, lookup_webauthn_pubkey);
15+
generate_ts_method!(identity, verify_account_linking_code);
1416

1517
generate_ts_method!(identity, approve_identity_link);
1618
generate_ts_method!(identity, create_identity);
@@ -19,6 +21,8 @@ fn main() {
1921
generate_ts_method!(identity, initiate_identity_link);
2022
generate_ts_method!(identity, prepare_delegation);
2123
generate_ts_method!(identity, remove_identity_link);
24+
generate_ts_method!(identity, create_account_linking_code);
25+
generate_ts_method!(identity, link_with_account_linking_code);
2226

2327
candid::export_service!();
2428
std::print!("{}", __export_service());

backend/canisters/identity/api/src/queries/get_account_linking_code.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use types::{AccountLinkingCode, Empty};
66
pub type Args = Empty;
77

88
#[ts_export(identity, get_account_linking_code)]
9-
#[derive(CandidType, Serialize, Deserialize, Debug)]
9+
#[derive(CandidType, Serialize, Deserialize, Debug, Eq, PartialEq)]
1010
pub enum Response {
1111
Success(AccountLinkingCode),
1212
NotFound,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,5 @@ pub type Args = Empty;
1010
pub enum Response {
1111
Success(AccountLinkingCode),
1212
UserNotFound,
13+
FailedToGenerateCode,
1314
}

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

Lines changed: 2 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use crate::WebAuthnKey;
22
use candid::{CandidType, Deserialize};
33
use serde::Serialize;
44
use ts_export::ts_export;
5-
use types::CanisterId;
5+
use types::UnitResult;
66

77
#[ts_export(identity, use_account_linking_code)]
88
#[derive(CandidType, Serialize, Deserialize, Debug)]
@@ -13,17 +13,4 @@ pub struct Args {
1313
pub webauthn_key: Option<WebAuthnKey>,
1414
}
1515

16-
#[ts_export(identity, use_account_linking_code)]
17-
#[derive(CandidType, Serialize, Deserialize, Debug)]
18-
pub enum Response {
19-
Success,
20-
AlreadyRegistered,
21-
AlreadyLinkedToPrincipal,
22-
TargetUserNotFound,
23-
PublicKeyInvalid(String),
24-
OriginatingCanisterInvalid(CanisterId),
25-
LinkedIdentitiesLimitReached(u32),
26-
LinkingCodeNotFound,
27-
LinkingCodeExpired,
28-
LinkingCodeInvalid,
29-
}
16+
pub type Response = UnitResult;

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,7 @@ impl RuntimeState {
168168
user_index: self.data.user_index_canister_id,
169169
cycles_dispenser: self.data.cycles_dispenser_canister_id,
170170
},
171+
account_linking_codes_count: self.data.account_linking_codes.len(),
171172
}
172173
}
173174
}
@@ -194,7 +195,7 @@ struct Data {
194195
challenges: Challenges,
195196
test_mode: bool,
196197
#[serde(default)]
197-
account_linking_codes: HashMap<String, (UserId, AccountLinkingCode)>,
198+
account_linking_codes: HashMap<String, AccountLinkingCode>,
198199
}
199200

200201
impl Data {
@@ -316,6 +317,7 @@ pub struct Metrics {
316317
pub originating_canisters: HashMap<CanisterId, u32>,
317318
pub stable_memory_sizes: BTreeMap<u8, u64>,
318319
pub canister_ids: CanisterIds,
320+
pub account_linking_codes_count: usize,
319321
}
320322

321323
#[derive(Serialize, Debug)]

backend/canisters/identity/impl/src/model/user_principals.rs

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -224,11 +224,17 @@ impl UserPrincipals {
224224
}
225225
}
226226

227-
pub fn find_user_principal_by_user_id(&self, user_id: UserId) -> Option<(Principal, Vec<Principal>)> {
227+
pub fn find_user_principal_by_user_id(&self, user_id: UserId) -> Option<UserPrincipal> {
228228
self.user_principals
229229
.iter()
230-
.find(|u| u.user_id == Some(user_id))
231-
.map(|u| (u.principal, u.auth_principals.clone()))
230+
.enumerate()
231+
.find(|(_, u)| u.user_id == Some(user_id))
232+
.map(|(i, u)| UserPrincipal {
233+
principal: u.principal,
234+
auth_principals: u.auth_principals.clone(),
235+
index: i as u32,
236+
user_id: Some(user_id),
237+
})
232238
}
233239

234240
// This is O(number of users) so only use in rare cases

backend/canisters/identity/impl/src/queries/get_account_linking_code.rs

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@ fn get_account_linking_code_impl(state: &RuntimeState) -> Response {
1313
let now = state.env.now();
1414

1515
if let Some(user_id) = state.get_user_id_by_caller() {
16-
if let Some(alc) = find_account_linking_code_for_user_id(state, &user_id) {
17-
if alc.is_valid(now) { Success(alc.clone()) } else { NotFound }
16+
if let Some(linking_code) = find_account_linking_code_for_user_id(state, user_id) {
17+
if linking_code.is_valid(now) { Success(linking_code.clone()) } else { NotFound }
1818
} else {
1919
NotFound
2020
}
@@ -23,11 +23,11 @@ fn get_account_linking_code_impl(state: &RuntimeState) -> Response {
2323
}
2424
}
2525

26-
fn find_account_linking_code_for_user_id(state: &RuntimeState, user_id: &UserId) -> Option<AccountLinkingCode> {
26+
fn find_account_linking_code_for_user_id(state: &RuntimeState, user_id: UserId) -> Option<AccountLinkingCode> {
2727
state
2828
.data
2929
.account_linking_codes
3030
.iter()
31-
.find(|(_, (alc_user_id, _))| alc_user_id == user_id)
32-
.map(|(_, (_, alc))| alc.clone())
31+
.find(|(_, linking_code)| linking_code.user_id == user_id)
32+
.map(|(_, linking_code)| linking_code.clone())
3333
}

backend/canisters/identity/impl/src/queries/verify_account_linking_code.rs

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,17 +7,17 @@ fn verify_account_linking_code(args: Args) -> Response {
77
read_state(|state| verify_account_linking_code_impl(args, state))
88
}
99

10-
// Verify that the account linking code is correct & valid for the caller's user ID.
10+
// Verify that the account linking code exists and is valid & correct.
11+
//
12+
// We're assuming this will be called by the user on a new device, trying to
13+
// link to their existing account.
14+
// TODO Make this a bit more secure: rate limit number of request.
1115
fn verify_account_linking_code_impl(args: Args, state: &RuntimeState) -> Response {
1216
let now = state.env.now();
1317

14-
if let Some(user_id) = state.get_user_id_by_caller() {
15-
if let Some((alc_user_id, alc)) = state.data.account_linking_codes.get(&args.code) {
16-
// Check if the account linking code is valid and matches the user ID.
17-
Response(alc.is_valid(now) && alc.value == args.code && user_id == *alc_user_id)
18-
} else {
19-
Response(false)
20-
}
18+
if let Some(linking_code) = state.data.account_linking_codes.get(&args.code) {
19+
// Check if the account linking code is valid!
20+
Response(linking_code.is_valid(now) && linking_code.value == args.code)
2121
} else {
2222
Response(false)
2323
}

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

Lines changed: 29 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -4,40 +4,49 @@ use canister_tracing_macros::trace;
44
use identity_canister::create_account_linking_code::{Response::*, *};
55
use types::AccountLinkingCode;
66

7+
const MAX_CODE_GEN_ATTEMPTS: usize = 20;
8+
79
#[update(msgpack = true, candid = true)]
810
#[trace]
911
fn create_account_linking_code(_args: Args) -> Response {
1012
mutate_state(create_account_linking_code_impl)
1113
}
1214

1315
fn create_account_linking_code_impl(state: &mut RuntimeState) -> Response {
16+
let caller_user_id = state.get_user_id_by_caller();
1417
let now = state.env.now();
15-
let caller = state.env.caller();
16-
17-
if let Some(user_id) = state
18-
.data
19-
.user_principals
20-
.get_by_auth_principal(&caller)
21-
.and_then(|u| u.user_id)
22-
{
18+
let rng = state.env.rng();
19+
20+
if let Some(user_id) = caller_user_id {
2321
// Clean up expired codes - keeps the memory footprint smaller in
2422
// exchange for a bit of extra CPU time.
25-
state.data.account_linking_codes.retain(|_, (_, alc)| alc.is_valid(now));
26-
27-
let mut alc = AccountLinkingCode::new(now);
28-
29-
// Verify that the generated code is unique
30-
while state.data.account_linking_codes.contains_key(&alc.value) {
31-
alc = AccountLinkingCode::new(now);
32-
}
33-
34-
// Store the code in the state (this part is assumed, as the original code does not specify how to store it)
3523
state
3624
.data
3725
.account_linking_codes
38-
.insert(alc.value.clone(), (user_id, alc.clone()));
26+
.retain(|_, linking_code| linking_code.is_valid(now));
27+
28+
// Attempt to generate linking code, up to max number of attempts
29+
let mut current_attempt = 0;
30+
while current_attempt < MAX_CODE_GEN_ATTEMPTS {
31+
let Ok(linking_code) = AccountLinkingCode::new(user_id, rng, now) else {
32+
continue;
33+
};
34+
35+
// Check if the code already exists in the state, if it does, we
36+
// try again; if not, we insert it and return success.
37+
if !state.data.account_linking_codes.contains_key(&linking_code.value) {
38+
state
39+
.data
40+
.account_linking_codes
41+
.insert(linking_code.value.clone(), linking_code.clone());
42+
43+
return Success(linking_code);
44+
}
45+
46+
current_attempt += 1;
47+
}
3948

40-
Success(alc)
49+
FailedToGenerateCode
4150
} else {
4251
UserNotFound
4352
}
Lines changed: 76 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
use crate::{RuntimeState, VerifyNewIdentityArgs, VerifyNewIdentityError, VerifyNewIdentitySuccess, mutate_state};
22
use canister_api_macros::update;
33
use canister_tracing_macros::trace;
4-
use identity_canister::link_with_account_linking_code::{Response::*, *};
4+
use identity_canister::link_with_account_linking_code::{Args, Response};
5+
use oc_error_codes::OCErrorCode;
6+
use types::OCResult;
57

68
// TODO Moved to a shared location, also consider what this means now in
79
// context of users using passkeys.
@@ -10,79 +12,93 @@ const MAX_LINKED_IDENTITIES: usize = 10;
1012
#[update(msgpack = true, candid = true)]
1113
#[trace]
1214
fn link_with_account_linking_code(args: Args) -> Response {
13-
mutate_state(|state| link_with_account_linking_code_impl(args, state))
15+
mutate_state(|state| link_with_account_linking_code_impl(args, state)).into()
1416
}
1517

16-
// TODO making this more secure:
17-
// - rate limit the number of attempts to link,
18-
// - lockout after too many attempts
19-
// - log and audit failed attempts (?)
20-
// - consider reducing the time window for the linking code to be valid
21-
fn link_with_account_linking_code_impl(args: Args, state: &mut RuntimeState) -> Response {
22-
// Get user ID and account linking code from the state. `args.code` is the
23-
// linking code provided by the user.
24-
let Some((user_id, linking_code)) = state.data.account_linking_codes.get(&args.code) else {
25-
return LinkingCodeNotFound;
18+
// Link accounts!
19+
//
20+
// At this point, no user is actually logged in, which means that the `caller`
21+
// value cannot be used to identify any user; but we can figure out principals
22+
// from the user id attached to the linking code.
23+
//
24+
// TODO Consider ways to make this functionality more secure: set limit on the
25+
// number of attempts to link; lock-out after too many failed attempts; log
26+
// and audit failed attempts; consider reducing validity time for linking codes.
27+
fn link_with_account_linking_code_impl(args: Args, state: &mut RuntimeState) -> OCResult {
28+
let now = state.env.now();
29+
30+
// Basically checks if the code provided by the user match any of the codes
31+
// that are saved.
32+
let Some(linking_code) = state.data.account_linking_codes.get(&args.code) else {
33+
return Err(OCErrorCode::LinkingCodeNotFound.into());
2634
};
2735

28-
// Returns user principal, and the list of auth principals linked to that user principal.
29-
if let Some((link_to_principal, auth_principals)) = state.data.user_principals.find_user_principal_by_user_id(*user_id) {
30-
// Verify the new identity using the provided public key and webauthn key.
31-
// This will also check if the caller is already registered.
32-
let VerifyNewIdentitySuccess {
33-
caller: _,
34-
auth_principal,
35-
originating_canister,
36-
webauthn_key,
37-
} = match state.verify_new_identity(VerifyNewIdentityArgs {
38-
public_key: args.public_key,
39-
webauthn_key: args.webauthn_key,
40-
allow_existing_provided_not_linked_to_oc_account: true,
41-
}) {
42-
Ok(ok) => ok,
43-
Err(error) => {
44-
return match error {
45-
VerifyNewIdentityError::AlreadyRegistered => AlreadyRegistered,
46-
VerifyNewIdentityError::PublicKeyInvalid(e) => PublicKeyInvalid(e),
47-
VerifyNewIdentityError::OriginatingCanisterInvalid(c) => OriginatingCanisterInvalid(c),
48-
};
49-
}
50-
};
36+
// Check if the linking code is still valid (i.e. not expired).
37+
if !linking_code.is_valid(state.env.now()) {
38+
return Err(OCErrorCode::LinkingCodeExpired.into());
39+
}
5140

52-
// Check if the linking code is expired.
53-
if !linking_code.is_valid(state.env.now()) {
54-
return LinkingCodeExpired;
41+
// Verify the new identity using the provided public key and webauthn key.
42+
// This will also check if the caller is already registered.
43+
let VerifyNewIdentitySuccess {
44+
caller: _,
45+
auth_principal: new_auth_principal,
46+
originating_canister,
47+
webauthn_key,
48+
} = match state.verify_new_identity(VerifyNewIdentityArgs {
49+
public_key: args.public_key,
50+
webauthn_key: args.webauthn_key,
51+
allow_existing_provided_not_linked_to_oc_account: true,
52+
}) {
53+
Ok(ok) => ok,
54+
Err(error) => {
55+
return Err(match error {
56+
VerifyNewIdentityError::AlreadyRegistered => OCErrorCode::AlreadyRegistered.into(),
57+
VerifyNewIdentityError::PublicKeyInvalid(_) => OCErrorCode::InvalidPublicKey.into(),
58+
VerifyNewIdentityError::OriginatingCanisterInvalid(_) => OCErrorCode::InvalidOriginatingCanister.into(),
59+
});
5560
}
61+
};
5662

57-
// Check if the linking code is valid.
58-
if linking_code.value != args.code {
59-
return LinkingCodeInvalid;
60-
}
63+
// Returns the list of auth principals for a given user id.
64+
let Some(user_principal) = state
65+
.data
66+
.user_principals
67+
.find_user_principal_by_user_id(linking_code.user_id)
68+
else {
69+
return Err(OCErrorCode::InitiatorNotFound.into());
70+
};
6171

62-
// Check that the target user principal has not reached the maximum number of linked identities
63-
if auth_principals.len() >= MAX_LINKED_IDENTITIES {
64-
return LinkedIdentitiesLimitReached(MAX_LINKED_IDENTITIES as u32);
65-
}
72+
// Check that the target user principal has not reached the maximum number of linked identities
73+
if user_principal.auth_principals.len() >= MAX_LINKED_IDENTITIES {
74+
// return LinkedIdentitiesLimitReached(MAX_LINKED_IDENTITIES as u32);
75+
return Err(OCErrorCode::UserLimitReached.into());
76+
}
6677

67-
// Check if the auth principal is already linked to the target user principal
68-
if auth_principals.contains(&auth_principal) {
69-
return AlreadyLinkedToPrincipal;
70-
}
78+
// Check if the auth principal is already linked to the target user principal
79+
if user_principal.auth_principals.contains(&new_auth_principal) {
80+
return Err(OCErrorCode::PrincipalAlreadyUsed.into());
81+
}
7182

72-
state.data.identity_link_requests.push(
73-
auth_principal,
74-
webauthn_key,
75-
originating_canister,
76-
false,
77-
link_to_principal,
78-
state.env.now(),
79-
);
83+
// Link the user with the new auth principal!
84+
if state.data.user_principals.link_auth_principal_with_existing_user(
85+
new_auth_principal,
86+
originating_canister,
87+
webauthn_key.as_ref().map(|k| k.credential_id.clone().into()),
88+
false,
89+
user_principal.index,
90+
now,
91+
) {
92+
if let Some(webauthn_key) = webauthn_key {
93+
state.data.webauthn_keys.add(webauthn_key, now);
94+
}
8095

8196
// Remove the linking code from the state, as it has been used.
8297
state.data.account_linking_codes.remove(&args.code);
8398

84-
Success
99+
// Linking done!
100+
Ok(())
85101
} else {
86-
TargetUserNotFound
102+
Err(OCErrorCode::PrincipalAlreadyUsed.into())
87103
}
88104
}

0 commit comments

Comments
 (0)