diff --git a/src/internet_identity/src/anchor_management.rs b/src/internet_identity/src/anchor_management.rs index 80242ac3cd..2806cf20af 100644 --- a/src/internet_identity/src/anchor_management.rs +++ b/src/internet_identity/src/anchor_management.rs @@ -1,5 +1,4 @@ use crate::archive::{archive_operation, device_diff}; -use crate::state::temp_keys::TempKeyId; use crate::state::RegistrationState::DeviceTentativelyAdded; use crate::state::TentativeDeviceRegistration; use crate::storage::anchor::{Anchor, AnchorError, Device}; @@ -134,10 +133,7 @@ pub fn replace( .add_device(new_device.clone()) .unwrap_or_else(|err| trap(&format!("failed to replace device: {err}"))); - state::with_temp_keys_mut(|temp_keys| { - let temp_key_id = TempKeyId::from_identity_authn_method(anchor_number, old_device.clone()); - temp_keys.remove_temp_key(&temp_key_id) - }); + state::with_temp_keys_mut(|temp_keys| temp_keys.remove_temp_key(anchor_number, &old_device)); Operation::ReplaceDevice { old_device, new_device: DeviceDataWithoutAlias::from(new_device), @@ -155,10 +151,7 @@ pub fn remove( .remove_device(&device_key) .unwrap_or_else(|err| trap(&format!("failed to remove device: {err}"))); - state::with_temp_keys_mut(|temp_keys| { - let temp_key_id = TempKeyId::from_identity_authn_method(anchor_number, device_key.clone()); - temp_keys.remove_temp_key(&temp_key_id) - }); + state::with_temp_keys_mut(|temp_keys| temp_keys.remove_temp_key(anchor_number, &device_key)); Operation::RemoveDevice { device: device_key } } diff --git a/src/internet_identity/src/anchor_management/registration.rs b/src/internet_identity/src/anchor_management/registration.rs index 0514f9f445..04bb5d78c6 100644 --- a/src/internet_identity/src/anchor_management/registration.rs +++ b/src/internet_identity/src/anchor_management/registration.rs @@ -1,6 +1,5 @@ use crate::anchor_management::{activity_bookkeeping, post_operation_bookkeeping}; use crate::state; -use crate::state::temp_keys::TempKeyId; use crate::state::ChallengeInfo; use crate::storage::anchor::Device; use candid::Principal; @@ -113,10 +112,7 @@ pub fn register( // `TempKeys` if let Some(temp_key) = temp_key { state::with_temp_keys_mut(|temp_keys| { - temp_keys.add_temp_key( - TempKeyId::from_identity_authn_method(anchor_number, device.pubkey.clone()), - temp_key, - ) + temp_keys.add_temp_key(&device.pubkey, anchor_number, temp_key) }); } diff --git a/src/internet_identity/src/anchor_management/registration/registration_flow_v2.rs b/src/internet_identity/src/anchor_management/registration/registration_flow_v2.rs index 37c39c167b..5d56b623ca 100644 --- a/src/internet_identity/src/anchor_management/registration/registration_flow_v2.rs +++ b/src/internet_identity/src/anchor_management/registration/registration_flow_v2.rs @@ -5,7 +5,6 @@ use crate::anchor_management::registration::rate_limit::process_rate_limit; use crate::anchor_management::{activity_bookkeeping, post_operation_bookkeeping}; use crate::state; use crate::state::flow_states::RegistrationFlowState; -use crate::state::temp_keys::TempKeyId; use crate::storage::anchor::Device; use candid::Principal; use ic_cdk::api::time; @@ -133,13 +132,7 @@ pub fn identity_registration_finish( // add temp key so the user can keep using the identity used for the registration flow state::with_temp_keys_mut(|temp_keys| { - temp_keys.add_temp_key( - TempKeyId::IdentityAuthnMethod { - identity_number, - authn_method_pubkey: arg.authn_method.public_key(), - }, - caller, - ) + temp_keys.add_temp_key(&arg.authn_method.public_key(), identity_number, caller) }); Ok(IdRegFinishResult { identity_number }) diff --git a/src/internet_identity/src/authz_utils.rs b/src/internet_identity/src/authz_utils.rs index a93ba57534..2a7b0e72f5 100644 --- a/src/internet_identity/src/authz_utils.rs +++ b/src/internet_identity/src/authz_utils.rs @@ -1,6 +1,5 @@ use crate::anchor_management::post_operation_bookkeeping; use crate::ii_domain::IIDomain; -use crate::state::temp_keys::TempKeyId; use crate::storage::anchor::Anchor; use crate::storage::StorageError; use crate::{anchor_management, state}; @@ -104,9 +103,9 @@ pub fn check_authorization( for device in anchor.devices() { if caller == Principal::self_authenticating(&device.pubkey) || state::with_temp_keys_mut(|temp_keys| { - let temp_key_id = - TempKeyId::from_identity_authn_method(anchor_number, device.pubkey.clone()); - temp_keys.check_temp_key(&caller, &temp_key_id).is_ok() + temp_keys + .check_temp_key(&caller, &device.pubkey, anchor_number) + .is_ok() }) { return Ok((anchor.clone(), device.pubkey.clone())); diff --git a/src/internet_identity/src/state.rs b/src/internet_identity/src/state.rs index 1e8d383a37..5448038991 100644 --- a/src/internet_identity/src/state.rs +++ b/src/internet_identity/src/state.rs @@ -21,7 +21,7 @@ use std::ops::{Deref, DerefMut}; use std::time::Duration; pub mod flow_states; -pub mod temp_keys; +mod temp_keys; /// Default captcha config pub const DEFAULT_CAPTCHA_CONFIG: CaptchaConfig = CaptchaConfig { diff --git a/src/internet_identity/src/state/temp_keys.rs b/src/internet_identity/src/state/temp_keys.rs index 4347b281f4..cbc36f8486 100644 --- a/src/internet_identity/src/state/temp_keys.rs +++ b/src/internet_identity/src/state/temp_keys.rs @@ -1,36 +1,12 @@ use crate::MINUTE_NS; use candid::Principal; use ic_cdk::api::time; -use internet_identity_interface::internet_identity::types::{IdentityNumber, PublicKey, Timestamp}; +use internet_identity_interface::internet_identity::types::{AnchorNumber, DeviceKey, Timestamp}; use std::collections::{HashMap, VecDeque}; // Expiration for temp keys, the same as the front-end delegation expiry const TEMP_KEY_EXPIRATION_NS: u64 = 10 * MINUTE_NS; -/// A temp key identifier -#[derive(Debug, Clone, Eq, PartialOrd, PartialEq, Hash)] -pub enum TempKeyId { - /// A temp key that is tied to an existing identity and authn_method. - /// We need to track the authn_method because the associated temp_key - /// must become invalid if the authn_method is removed from the identity. - IdentityAuthnMethod { - identity_number: IdentityNumber, - authn_method_pubkey: PublicKey, - }, -} - -impl TempKeyId { - pub fn from_identity_authn_method( - identity_number: IdentityNumber, - authn_method_pubkey: PublicKey, - ) -> Self { - TempKeyId::IdentityAuthnMethod { - identity_number, - authn_method_pubkey, - } - } -} - #[derive(Default, Debug)] pub struct TempKeys { /// A map of "temporary keys" attached to specific anchor and device combination. @@ -50,40 +26,50 @@ pub struct TempKeys { /// /// Since temp keys can only be added during registration, the max number of temp keys is /// bounded by the registration rate limit. - temp_keys: HashMap, + temp_keys: HashMap<(AnchorNumber, DeviceKey), TempKey>, /// Deque to efficiently prune expired temp keys expirations: VecDeque, } impl TempKeys { - pub fn add_temp_key(&mut self, key_id: TempKeyId, temp_key: Principal) { + pub fn add_temp_key( + &mut self, + device_key: &DeviceKey, + anchor: AnchorNumber, + temp_key: Principal, + ) { let tmp_key = TempKey { principal: temp_key, expiration: time() + TEMP_KEY_EXPIRATION_NS, }; self.expirations.push_back(TempKeyExpiration { - key_id: key_id.clone(), + key: (anchor, device_key.clone()), expiration: tmp_key.expiration, }); - self.temp_keys.insert(key_id, tmp_key); + self.temp_keys.insert((anchor, device_key.clone()), tmp_key); } /// Removes the temporary key for the given device if it exists and is linked to the provided anchor. - pub fn remove_temp_key(&mut self, key_id: &TempKeyId) { + pub fn remove_temp_key(&mut self, anchor: AnchorNumber, device_key: &DeviceKey) { // we can skip the removal from expirations because there it will be removed // during amortized clean-up operations - self.temp_keys.remove(key_id); + self.temp_keys.remove(&(anchor, device_key.clone())); } /// Checks that the temporary key is valid for the given device and anchor. /// /// Requires a mutable reference because it does amortized clean-up of expired temp keys. - pub fn check_temp_key(&mut self, caller: &Principal, key_id: &TempKeyId) -> Result<(), ()> { + pub fn check_temp_key( + &mut self, + caller: &Principal, + device_key: &DeviceKey, + anchor: AnchorNumber, + ) -> Result<(), ()> { self.prune_expired_keys(); - let Some(temp_key) = self.temp_keys.get(key_id) else { + let Some(temp_key) = self.temp_keys.get(&(anchor, device_key.clone())) else { return Err(()); }; if temp_key.expiration < time() { @@ -110,7 +96,7 @@ impl TempKeys { if expiration.expiration > now { break; } - self.temp_keys.remove(&expiration.key_id); + self.temp_keys.remove(&expiration.key); self.expirations.pop_front(); } } @@ -131,7 +117,7 @@ struct TempKey { #[derive(Clone, Debug, Eq, PartialEq)] struct TempKeyExpiration { /// Key with which to find the temp key in the `temp_keys` map - pub key_id: TempKeyId, + key: (AnchorNumber, DeviceKey), /// The expiration timestamp of the temp key - pub expiration: Timestamp, + expiration: Timestamp, }