olsemeno · olsemeno · Jun 1, 2022 · Jun 9, 2022
diff --git a/backend-tests/backend-tests.hs b/backend-tests/backend-tests.hs
@@ -254,6 +254,7 @@ device1 = empty
     .+ #pubkey .== webauth1PK
     .+ #credential_id .== Nothing
     .+ #purpose .== enum #authentication
+    .+ #protection_type .== Nothing
     .+ #key_type .== enum #cross_platform
 
 webauth2SK :: SecretKey
@@ -269,8 +270,34 @@ device2 = empty
     .+ #pubkey .== webauth2PK
     .+ #credential_id .== Just "foobar"
     .+ #purpose .== enum #authentication
+    .+ #protection_type .== Nothing
     .+ #key_type .== enum #platform
 
+webauth3SK :: SecretKey
+webauth3SK = createSecretKeyWebAuthnRSA "foobar3"
+webauth3PK = toPublicKey webauth3SK
+webauth3PK :: PublicKey
+webauth3ID :: EntityId
+webauth3ID = EntityId $ mkSelfAuthenticatingId webauth3PK
+
+device3 :: DeviceData
+device3 = empty
+    .+ #alias .== "device3"
+    .+ #pubkey .== webauth3PK
+    .+ #credential_id .== Just "foobar"
+    .+ #purpose .== enum #authentication
+    .+ #key_type .== enum #seed_phrase
+    .+ #protection_type .== Just (enum #protected)
+
+device4 :: DeviceData
+device4 = empty
+    .+ #alias .== "device4"
+    .+ #pubkey .== webauth1PK
+    .+ #credential_id .== Just "foobar"
+    .+ #purpose .== enum #authentication
+    .+ #key_type .== enum #seed_phrase
+    .+ #protection_type .== Just (enum #protected)
+
 anonymousID :: EntityId
 anonymousID = EntityId "\x04"
 
@@ -735,6 +762,18 @@ tests wasm_file = testGroup "Tests" $ upgradeGroups $
     when (user_number == user_number2) $
       lift $ assertFailure "Identity Anchor re-used"
 
+  , withUpgrade $ \should_upgrade -> iiTest "remove protected seed phrase" $ \cid -> do
+    user_number <- register cid webauth3ID device3 >>= mustGetUserNumber
+    callII cid webauth3ID #add (user_number, device2)
+    lookupIs cid user_number [device3, device2]
+    callIIRejectWith cid webauth2ID #remove (user_number, webauth3PK) "failed to remove protected recovery phrase"
+    callIIRejectWith cid webauth3ID #add (user_number, device4) "recovery mechanism already protected"
+    when should_upgrade $ doUpgrade cid
+    callII cid webauth3ID #remove (user_number, webauth3PK)
+    lookupIs cid user_number [device2]
+    callII cid webauth2ID #add (user_number, device4)
+    lookupIs cid user_number [device2, device4]
+
   , withUpgrade $ \should_upgrade -> iiTestWithInit "init range" (100, 103) $ \cid -> do
     s <- queryII cid dummyUserId #stats ()
     lift $ s .! #assigned_user_number_range @?= (100, 103)
@@ -822,15 +861,16 @@ tests wasm_file = testGroup "Tests" $ upgradeGroups $
     -- and should be auto-corrected on upgrade
     lift $ s .! #assigned_user_number_range @?= (10_000, 10_000 + 3_774_873)
     lookupIs cid 9_999 []
-    lookupIs cid 10_000 [#alias .== "Desktop" .+ #credential_id .== Just "c\184\175\179\134\221u}\250[\169U\v\202f\147g\ETBvo9[\175\173\144R\163\132\237\196F\177\DC2(\188\185\203hI\128\187Z\129'\v1\212\185V\ETB\135)m@ M1\233l\ESC8kI\132" .+ #pubkey .== "0^0\f\ACK\n+\ACK\SOH\EOT\SOH\131\184C\SOH\SOH\ETXN\NUL\165\SOH\STX\ETX& \SOH!X \238o!-\ESC\148\252\192\DC4\240P\176\135\240j\211AW\255S\193\153\129\227\151hB\177dK\n\FS\"X \rk\197\238\a{\210\&0\v<\134\223\135\170_\223\144\210V\208\DC3\RS\251\228D$3\r\232\176\EOTq" .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown ]
-    lookupIs cid 10_002 [#alias .== "andrew-mbp" .+ #credential_id .== Just "\SOH\191#%\217u\247\178L-K\182\254\249J.m\187\179_I\ACK\137\244`\163o\SI\150qz\197Hz\214\&8\153\239\213\159\208\RS\243\138\171\138\"\139\173\170\ESC\148\205\129\149ri\\Dn,7\151\146\175\DEL" .+ #pubkey .== "0^0\f\ACK\n+\ACK\SOH\EOT\SOH\131\184C\SOH\SOH\ETXN\NUL\165\SOH\STX\ETX& \SOH!X rMm*\229BDe\SOH4\228u\170\206\216\216-ER\v\166r\217\137,\141\227M*@\230\243\"X \225\248\159\191\242\224Z>\241\163\\\GS\155\178\222\139^\136V\253q\v\SUBSJ\bA\131\\\183\147\170" .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown,#alias .== "andrew phone chrome" .+ #credential_id .== Just ",\235x\NUL\a\140`~\148\248\233C/\177\205\158\ETX0\129\167" .+ #pubkey .== "0^0\f\ACK\n+\ACK\SOH\EOT\SOH\131\184C\SOH\SOH\ETXN\NUL\165\SOH\STX\ETX& \SOH!X \140\169\203@\ETX\CAN\ETB,\177\153\179\223/|`\US\STX\252r\190s(.\188\136\171\SI\181V*\174@\"X \245<\174AbV\225\234\ENQ\146\247\129\245\ACK\200\205\217\250g\219\179)\197\252\164i\172kXh\180\205" .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown]
-    lookupIs cid 10_029 [#alias .== "Pixel" .+ #credential_id .== Just "\SOH\146\238\160b\223\132\205\231b\239\243F\170\163\167\251D\241\170\EM\216\136\174@r\149\183|LuKu[+{\144\217\ETBL\f\244\GS>\179\146\143\RS\179\DLE\227\179\164\188\NULDQy\223\SI\132\183\248\177\219" .+ #pubkey .== "0^0\f\ACK\n+\ACK\SOH\EOT\SOH\131\184C\SOH\SOH\ETXN\NUL\165\SOH\STX\ETX& \SOH!X \200B>\DEL\GS\248\220\145\245\153\221\&6\131\243uAQCAd>\145k\nw\233\&5\218\SUB~_\244\"X O]7\167=n\ESC\SUB\198\235\208\215s\158\191Gz\143\136\237i\146\203\&6\182\196\129\239\238\SOH\180b" .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown]
+    lookupIs cid 10_000 [#alias .== "Desktop" .+ #credential_id .== Just "c\184\175\179\134\221u}\250[\169U\v\202f\147g\ETBvo9[\175\173\144R\163\132\237\196F\177\DC2(\188\185\203hI\128\187Z\129'\v1\212\185V\ETB\135)m@ M1\233l\ESC8kI\132" .+ #pubkey .== "0^0\f\ACK\n+\ACK\SOH\EOT\SOH\131\184C\SOH\SOH\ETXN\NUL\165\SOH\STX\ETX& \SOH!X \238o!-\ESC\148\252\192\DC4\240P\176\135\240j\211AW\255S\193\153\129\227\151hB\177dK\n\FS\"X \rk\197\238\a{\210\&0\v<\134\223\135\170_\223\144\210V\208\DC3\RS\251\228D$3\r\232\176\EOTq" .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown .+ #protection_type .== Nothing]
+    lookupIs cid 10_000 [#alias .== "Desktop" .+ #credential_id .== Just "c\184\175\179\134\221u}\250[\169U\v\202f\147g\ETBvo9[\175\173\144R\163\132\237\196F\177\DC2(\188\185\203hI\128\187Z\129'\v1\212\185V\ETB\135)m@ M1\233l\ESC8kI\132" .+ #pubkey .== "0^0\f\ACK\n+\ACK\SOH\EOT\SOH\131\184C\SOH\SOH\ETXN\NUL\165\SOH\STX\ETX& \SOH!X \238o!-\ESC\148\252\192\DC4\240P\176\135\240j\211AW\255S\193\153\129\227\151hB\177dK\n\FS\"X \rk\197\238\a{\210\&0\v<\134\223\135\170_\223\144\210V\208\DC3\RS\251\228D$3\r\232\176\EOTq" .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown .+ #protection_type .== Nothing]
+    lookupIs cid 10_002 [#protection_type .== Nothing .+ #alias .== "andrew-mbp" .+ #credential_id .== Just "\SOH\191#%\217u\247\178L-K\182\254\249J.m\187\179_I\ACK\137\244`\163o\SI\150qz\197Hz\214\&8\153\239\213\159\208\RS\243\138\171\138\"\139\173\170\ESC\148\205\129\149ri\\Dn,7\151\146\175\DEL" .+ #pubkey .== "0^0\f\ACK\n+\ACK\SOH\EOT\SOH\131\184C\SOH\SOH\ETXN\NUL\165\SOH\STX\ETX& \SOH!X rMm*\229BDe\SOH4\228u\170\206\216\216-ER\v\166r\217\137,\141\227M*@\230\243\"X \225\248\159\191\242\224Z>\241\163\\\GS\155\178\222\139^\136V\253q\v\SUBSJ\bA\131\\\183\147\170" .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown,#alias .== "andrew phone chrome" .+ #credential_id .== Just ",\235x\NUL\a\140`~\148\248\233C/\177\205\158\ETX0\129\167" .+ #pubkey .== "0^0\f\ACK\n+\ACK\SOH\EOT\SOH\131\184C\SOH\SOH\ETXN\NUL\165\SOH\STX\ETX& \SOH!X \140\169\203@\ETX\CAN\ETB,\177\153\179\223/|`\US\STX\252r\190s(.\188\136\171\SI\181V*\174@\"X \245<\174AbV\225\234\ENQ\146\247\129\245\ACK\200\205\217\250g\219\179)\197\252\164i\172kXh\180\205" .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown .+ #protection_type .== Nothing]
+    lookupIs cid 10_029 [#alias .== "Pixel" .+ #credential_id .== Just "\SOH\146\238\160b\223\132\205\231b\239\243F\170\163\167\251D\241\170\EM\216\136\174@r\149\183|LuKu[+{\144\217\ETBL\f\244\GS>\179\146\143\RS\179\DLE\227\179\164\188\NULDQy\223\SI\132\183\248\177\219" .+ #pubkey .== "0^0\f\ACK\n+\ACK\SOH\EOT\SOH\131\184C\SOH\SOH\ETXN\NUL\165\SOH\STX\ETX& \SOH!X \200B>\DEL\GS\248\220\145\245\153\221\&6\131\243uAQCAd>\145k\nw\233\&5\218\SUB~_\244\"X O]7\167=n\ESC\SUB\198\235\208\215s\158\191Gz\143\136\237i\146\203\&6\182\196\129\239\238\SOH\180b" .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown .+ #protection_type .== Nothing]
     -- This user record has been created manullay with dfx and our test
     -- webauth1PK has been added, so that we can actually log into this now
     let dfxPK =  "0*0\ENQ\ACK\ETX+ep\ETX!\NUL\241\186;\128\206$\243\130\250\&2\253\a#<\235\142\&0]W\218\254j\211\209\192\SO@\DC3\NAKi&1"
-    lookupIs cid 10_030 [#alias .== "dfx" .+ #credential_id .== Nothing .+ #pubkey .== dfxPK .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown,#alias .== "testkey" .+ #credential_id .== Nothing .+ #pubkey .== webauth1PK .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown]
+    lookupIs cid 10_030 [#protection_type .== Nothing .+ #alias .== "dfx" .+ #credential_id .== Nothing .+ #pubkey .== dfxPK .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown,#alias .== "testkey" .+ #credential_id .== Nothing .+ #pubkey .== webauth1PK .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown .+ #protection_type .== Nothing]
     callII cid webauth1ID #remove (10_030, dfxPK)
-    lookupIs cid 10_030 [#alias .== "testkey" .+ #credential_id .== Nothing .+ #pubkey .== webauth1PK .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown]
+    lookupIs cid 10_030 [#alias .== "testkey" .+ #credential_id .== Nothing .+ #pubkey .== webauth1PK .+ #purpose .== enum #authentication .+ #key_type .== enum #unknown .+ #protection_type .== Nothing]
     let delegationArgs = (10_030, "example.com", "dummykey", Nothing)
     (userPK,_) <- callII cid webauth1ID #prepare_delegation delegationArgs
     -- Check that we get the same user key; this proves that the salt was

diff --git a/src/internet_identity/internet_identity.did b/src/internet_identity/internet_identity.did
@@ -49,6 +49,11 @@ type KeyType = variant {
     seed_phrase;
 };
 
+type ProtectionType = variant {
+    protected;
+    unknown;
+};
+
 type Challenge = record {
     png_base64: text;
     challenge_key: ChallengeKey;
@@ -60,6 +65,7 @@ type DeviceData = record {
   credential_id : opt CredentialId;
   purpose: Purpose;
   key_type: KeyType;
+  protection_type: opt ProtectionType;
 };
 
 type RegisterResponse = variant {

diff --git a/src/internet_identity/src/main.rs b/src/internet_identity/src/main.rs
@@ -16,6 +16,7 @@ use serde_bytes::ByteBuf;
 use std::cell::{Cell, RefCell};
 use std::collections::HashMap;
 use std::convert::TryInto;
+use std::ptr::eq;
 use storage::{Salt, Storage};
 
 mod assets;
@@ -27,6 +28,8 @@ const fn secs_to_nanos(secs: u64) -> u64 {
 
 #[cfg(not(feature = "dummy_captcha"))]
 use captcha::filters::Wave;
+use crate::KeyType::SeedPhrase;
+use crate::ProtectionType::Protected;
 
 // 30 mins
 const DEFAULT_EXPIRATION_PERIOD_NS: u64 = secs_to_nanos(30 * 60);
@@ -57,7 +60,8 @@ type DeviceKey = PublicKey;
 type UserKey = PublicKey;
 type SessionKey = PublicKey;
 type FrontendHostname = String;
-type Timestamp = u64; // in nanos since epoch
+type Timestamp = u64;
+// in nanos since epoch
 type Signature = ByteBuf;
 type DeviceVerificationCode = String;
 type FailedAttemptsCounter = u8;
@@ -72,7 +76,7 @@ enum Purpose {
     Authentication,
 }
 
-#[derive(Clone, Debug, CandidType, Deserialize)]
+#[derive(Clone, Debug, CandidType, Deserialize, PartialEq)]
 enum KeyType {
     #[serde(rename = "unknown")]
     Unknown,
@@ -84,13 +88,20 @@ enum KeyType {
     SeedPhrase,
 }
 
+#[derive(Clone, Debug, CandidType, Deserialize, PartialEq)]
+enum ProtectionType {
+    #[serde(rename = "protected")]
+    Protected,
+}
+
 #[derive(Clone, Debug, CandidType, Deserialize)]
 struct DeviceData {
     pubkey: DeviceKey,
     alias: String,
     credential_id: Option<CredentialId>,
     purpose: Purpose,
     key_type: KeyType,
+    protection_type: Option<ProtectionType>,
 }
 
 /// This is an internal version of `DeviceData` primarily useful to provide a
@@ -103,6 +114,7 @@ struct DeviceDataInternal {
     credential_id: Option<CredentialId>,
     purpose: Option<Purpose>,
     key_type: Option<KeyType>,
+    protection_type: Option<ProtectionType>,
 }
 
 impl From<DeviceData> for DeviceDataInternal {
@@ -113,6 +125,7 @@ impl From<DeviceData> for DeviceDataInternal {
             credential_id: device_data.credential_id,
             purpose: Some(device_data.purpose),
             key_type: Some(device_data.key_type),
+            protection_type: device_data.protection_type,
         }
     }
 }
@@ -127,6 +140,7 @@ impl From<DeviceDataInternal> for DeviceData {
                 .purpose
                 .unwrap_or(Purpose::Authentication),
             key_type: device_data_internal.key_type.unwrap_or(KeyType::Unknown),
+            protection_type: device_data_internal.protection_type,
         }
     }
 }
@@ -296,7 +310,7 @@ async fn init_salt() {
     });
 
     let res: Vec<u8> = match call(Principal::management_canister(), "raw_rand", ()).await {
-        Ok((res,)) => res,
+        Ok((res, )) => res,
         Err((_, err)) => trap(&format!("failed to get salt: {}", err)),
     };
     let salt: Salt = res[..].try_into().unwrap_or_else(|_| {
@@ -396,9 +410,9 @@ async fn add_tentative_device(
                 AddTentativeDeviceResponse::DeviceRegistrationModeOff
             }
             Some(TentativeDeviceRegistration {
-                state: DeviceTentativelyAdded { .. },
-                ..
-            }) => AnotherDeviceTentativelyAdded,
+                     state: DeviceTentativelyAdded { .. },
+                     ..
+                 }) => AnotherDeviceTentativelyAdded,
             Some(mut registration) => {
                 registration.state = DeviceTentativelyAdded {
                     tentative_device: device_data,
@@ -483,7 +497,7 @@ fn get_verified_device(
 /// Return a decimal representation of a random `u32` to be used as verification code
 async fn new_verification_code() -> DeviceVerificationCode {
     let res: Vec<u8> = match call(Principal::management_canister(), "raw_rand", ()).await {
-        Ok((res,)) => res,
+        Ok((res, )) => res,
         Err((_, err)) => trap(&format!("failed to get randomness: {}", err)),
     };
     let rand = u32::from_be_bytes(res[..4].try_into().unwrap_or_else(|_| {
@@ -562,10 +576,28 @@ async fn add(user_number: UserNumber, device_data: DeviceData) {
 
         trap_if_not_authenticated(entries.iter().map(|e| &e.pubkey));
 
+        let is_protected_recovery_device = match device_data.protection_type {
+            None => { false }
+            Some(ref protection_type) => {
+                protection_type.eq(&ProtectionType::Protected)
+                    && device_data.key_type.eq(&KeyType::SeedPhrase)
+            }
+        };
+
         for e in entries.iter_mut() {
             if e.pubkey == device_data.pubkey {
                 trap("Device already added.");
             }
+            if is_protected_recovery_device {
+                match e.protection_type {
+                    None => {}
+                    Some(_) => {
+                        if e.key_type.is_some() && e.key_type.as_ref().unwrap().eq(&KeyType::SeedPhrase) {
+                            trap("recovery mechanism already protected");
+                        }
+                    }
+                }
+            }
         }
 
         if entries.len() >= MAX_ENTRIES_PER_USER {
@@ -606,6 +638,15 @@ async fn remove(user_number: UserNumber, device_key: DeviceKey) {
         trap_if_not_authenticated(entries.iter().map(|e| &e.pubkey));
 
         if let Some(i) = entries.iter().position(|e| e.pubkey == device_key) {
+            let entry_to_remove = entries.get(i as usize).unwrap();
+
+            if entry_to_remove.key_type.is_some() && entry_to_remove.key_type.as_ref().unwrap() == &SeedPhrase
+                && entry_to_remove.protection_type.is_some() && entry_to_remove.protection_type.as_ref().unwrap() == &Protected {
+                if caller() != Principal::self_authenticating(entry_to_remove.pubkey.clone()) {
+                    trap("failed to remove protected recovery phrase");
+                }
+            }
+
             entries.swap_remove(i as usize);
         }
 
@@ -706,14 +747,14 @@ fn random_string<T: RngCore>(rng: &mut T, n: usize) -> String {
 // Get a random number generator based on 'raw_rand'
 async fn make_rng() -> rand_chacha::ChaCha20Rng {
     let raw_rand: Vec<u8> = match call(Principal::management_canister(), "raw_rand", ()).await {
-        Ok((res,)) => res,
+        Ok((res, )) => res,
         Err((_, err)) => trap(&format!("failed to get seed: {}", err)),
     };
     let seed: Salt = raw_rand[..].try_into().unwrap_or_else(|_| {
         trap(&format!(
-                "when creating seed from raw_rand output, expected raw randomness to be of length 32, got {}",
-                raw_rand.len()
-                ));
+            "when creating seed from raw_rand output, expected raw randomness to be of length 32, got {}",
+            raw_rand.len()
+        ));
     });
 
     rand_chacha::ChaCha20Rng::from_seed(seed)
@@ -803,12 +844,12 @@ fn get_anchor_info(user_number: UserNumber) -> IdentityAnchorInfo {
             .get(&user_number)
         {
             Some(TentativeDeviceRegistration {
-                expiration,
-                state:
-                    DeviceTentativelyAdded {
-                        tentative_device, ..
-                    },
-            }) if *expiration > now => IdentityAnchorInfo {
+                     expiration,
+                     state:
+                     DeviceTentativelyAdded {
+                         tentative_device, ..
+                     },
+                 }) if *expiration > now => IdentityAnchorInfo {
                 devices,
                 device_registration: Some(DeviceRegistrationInfo {
                     expiration: *expiration,
@@ -1164,7 +1205,7 @@ fn prune_expired_signatures(asset_hashes: &AssetHashes, sigs: &mut SignatureMap)
 
 // Checks if the caller is authenticated against any of the public keys provided
 // and traps if not.
-fn trap_if_not_authenticated<'a>(public_keys: impl Iterator<Item = &'a PublicKey>) {
+fn trap_if_not_authenticated<'a>(public_keys: impl Iterator<Item=&'a PublicKey>) {
     for pk in public_keys {
         if caller() == Principal::self_authenticating(pk) {
             return;