Revert "Refactor temp_keys in preparation for dynamic captcha registr…

…ation flow (#2618)" (#2632)

* Revert "Refactor temp_keys in preparation for dynamic captcha registration flow (#2618)"

This reverts commit 60175a1.

The only changes kept from #2618 are two typo fixes in the comments.
Reverting required to adapt the new flow module slightly.

* Remove unnecessary pub access modifiers

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 }
 }
 

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)
         });
     }
 

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 })

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()));

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 {

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<TempKeyId, TempKey>,
+    temp_keys: HashMap<(AnchorNumber, DeviceKey), TempKey>,
 
     /// Deque to efficiently prune expired temp keys
     expirations: VecDeque<TempKeyExpiration>,
 }
 
 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,
 }