Refactor temp_keys in preparation for dynamic captcha registration fl…

…ow (#2618)

For the dynamic captcha registration flow, temp_keys need to be tracked
before the identity number or authn method public key is known. For that
reason, this refactoring abstracts out the `TempKeyId` such that we can
introduce a different variant to track the keys.
For example a registration flow id.

src/internet_identity/src/anchor_management.rs

@@ -1,4 +1,5 @@
 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};
@@ -133,7 +134,10 @@ 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| temp_keys.remove_temp_key(anchor_number, &old_device));
+    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)
+    });
     Operation::ReplaceDevice {
         old_device,
         new_device: DeviceDataWithoutAlias::from(new_device),
@@ -151,7 +155,10 @@ 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| temp_keys.remove_temp_key(anchor_number, &device_key));
+    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)
+    });
     Operation::RemoveDevice { device: device_key }
 }
 

src/internet_identity/src/anchor_management/registration.rs

@@ -10,6 +10,7 @@ use internet_identity_interface::internet_identity::types::*;
 use rand_core::{RngCore, SeedableRng};
 use std::collections::{HashMap, HashSet};
 
+use crate::state::temp_keys::TempKeyId;
 #[cfg(not(feature = "dummy_captcha"))]
 use captcha::filters::Wave;
 use captcha::fonts::Default as DefaultFont;
@@ -267,7 +268,10 @@ pub fn register(
     // `TempKeys`
     if let Some(temp_key) = temp_key {
         state::with_temp_keys_mut(|temp_keys| {
-            temp_keys.add_temp_key(&device.pubkey, anchor_number, temp_key)
+            temp_keys.add_temp_key(
+                TempKeyId::from_identity_authn_method(anchor_number, device.pubkey.clone()),
+                temp_key,
+            )
         });
     }
 

src/internet_identity/src/authz_utils.rs

@@ -1,5 +1,6 @@
 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};
@@ -101,9 +102,9 @@ pub fn check_authorization(
     for device in anchor.devices() {
         if caller == Principal::self_authenticating(&device.pubkey)
             || state::with_temp_keys_mut(|temp_keys| {
-                temp_keys
-                    .check_temp_key(&caller, &device.pubkey, anchor_number)
-                    .is_ok()
+                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()
             })
         {
             return Ok((anchor.clone(), device.pubkey.clone()));

src/internet_identity/src/state.rs

@@ -19,7 +19,7 @@ use std::collections::HashMap;
 use std::ops::{Deref, DerefMut};
 use std::time::Duration;
 
-mod temp_keys;
+pub mod temp_keys;
 
 /// Default captcha config
 pub const DEFAULT_CAPTCHA_CONFIG: CaptchaConfig = CaptchaConfig {

src/internet_identity/src/state/temp_keys.rs

@@ -1,12 +1,36 @@
 use crate::MINUTE_NS;
 use candid::Principal;
 use ic_cdk::api::time;
-use internet_identity_interface::internet_identity::types::{AnchorNumber, DeviceKey, Timestamp};
+use internet_identity_interface::internet_identity::types::{IdentityNumber, PublicKey, 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.
@@ -17,59 +41,49 @@ pub struct TempKeys {
     ///
     /// Note: we link the temporary keys to a device so that we can make sure the temporary key is dropped
     /// if the device itself is removed. This ensures the temporary key is no more powerful than the device,
-    /// i.e. if the user decides to remove the device right after registration, then the the temporary key
+    /// i.e. if the user decides to remove the device right after registration, then the temporary key
     /// cannot be used to authenticate (similarly to how the browser session key pair cannot be used to
     /// authenticate if the actual delegated WebAuthn device is removed).
     ///
     /// In addition, the temporary key is also linked to an anchor so that even if the same device is
     /// added to multiple anchors, it can only be used in the context of the anchor it was created for.
     ///
-    /// Since temp keys can only be added during registration, the the max number of temp keys is
+    /// 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<(AnchorNumber, DeviceKey), TempKey>,
+    temp_keys: HashMap<TempKeyId, TempKey>,
 
     /// Deque to efficiently prune expired temp keys
     expirations: VecDeque<TempKeyExpiration>,
 }
 
 impl TempKeys {
-    pub fn add_temp_key(
-        &mut self,
-        device_key: &DeviceKey,
-        anchor: AnchorNumber,
-        temp_key: Principal,
-    ) {
+    pub fn add_temp_key(&mut self, key_id: TempKeyId, temp_key: Principal) {
         let tmp_key = TempKey {
             principal: temp_key,
             expiration: time() + TEMP_KEY_EXPIRATION_NS,
         };
 
         self.expirations.push_back(TempKeyExpiration {
-            key: (anchor, device_key.clone()),
+            key_id: key_id.clone(),
             expiration: tmp_key.expiration,
         });
-        self.temp_keys.insert((anchor, device_key.clone()), tmp_key);
+        self.temp_keys.insert(key_id, 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, anchor: AnchorNumber, device_key: &DeviceKey) {
+    pub fn remove_temp_key(&mut self, key_id: &TempKeyId) {
         // we can skip the removal from expirations because there it will be removed
         // during amortized clean-up operations
-        self.temp_keys.remove(&(anchor, device_key.clone()));
+        self.temp_keys.remove(key_id);
     }
 
     /// 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,
-        device_key: &DeviceKey,
-        anchor: AnchorNumber,
-    ) -> Result<(), ()> {
+    pub fn check_temp_key(&mut self, caller: &Principal, key_id: &TempKeyId) -> Result<(), ()> {
         self.prune_expired_keys();
 
-        let Some(temp_key) = self.temp_keys.get(&(anchor, device_key.clone())) else {
+        let Some(temp_key) = self.temp_keys.get(key_id) else {
             return Err(());
         };
         if temp_key.expiration < time() {
@@ -96,7 +110,7 @@ impl TempKeys {
             if expiration.expiration > now {
                 break;
             }
-            self.temp_keys.remove(&expiration.key);
+            self.temp_keys.remove(&expiration.key_id);
             self.expirations.pop_front();
         }
     }
@@ -107,17 +121,17 @@ impl TempKeys {
 }
 
 #[derive(Clone, Debug, Eq, PartialEq)]
-pub struct TempKey {
+struct TempKey {
     /// The temp key principal
     principal: Principal,
     /// The expiration timestamp of the temp key
     expiration: Timestamp,
 }
 
 #[derive(Clone, Debug, Eq, PartialEq)]
-pub struct TempKeyExpiration {
+struct TempKeyExpiration {
     /// Key with which to find the temp key in the `temp_keys` map
-    pub key: (AnchorNumber, DeviceKey),
+    pub key_id: TempKeyId,
     /// The expiration timestamp of the temp key
     pub expiration: Timestamp,
 }