Skip to main content

bitwarden_user_crypto_management/key_rotation/
unlock.rs

1//! Functionality for re-encrypting unlock (decryption) methods during user key rotation.
2//! During key-rotation, a new user-key is sampled. The unlock module then creates a set of newly
3//! encrypted copies, one for each decryption/unlock method.
4use bitwarden_api_api::models::{
5    self, CommonUnlockDataRequestModel, EmergencyAccessWithIdRequestModel,
6    MasterPasswordUnlockAndAuthenticationDataModel, OtherDeviceKeysUpdateRequestModel,
7    ResetPasswordWithOrgIdRequestModel, UnlockDataRequestModel, V2UpgradeTokenRequestModel,
8    WebAuthnLoginRotateKeyRequestModel,
9};
10use bitwarden_core::key_management::{
11    KeySlotIds, MasterPasswordAuthenticationData, MasterPasswordUnlockData, SymmetricKeySlotId,
12    V2UpgradeToken,
13};
14use bitwarden_crypto::{Kdf, KeyStoreContext, PublicKey, SymmetricKeyAlgorithm, UnsignedSharedKey};
15use serde::{Deserialize, Serialize};
16use tracing::{debug, debug_span, error, info};
17#[cfg(feature = "wasm")]
18use tsify::Tsify;
19
20use crate::key_rotation::{
21    partial_rotateable_keyset::PartialRotateableKeyset, rotate_user_keys::UpgradeTokenAction,
22};
23
24/// The data necessary to re-share the user-key to a V1 emergency access membership. Note: The
25/// Public-key must be verified/trusted. Further, there is no sender authentication possible here.
26#[derive(Serialize, Deserialize, Clone)]
27#[cfg_attr(feature = "wasm", derive(Tsify), tsify(into_wasm_abi, from_wasm_abi))]
28pub struct V1EmergencyAccessMembership {
29    pub id: uuid::Uuid,
30    pub grantee_id: uuid::Uuid,
31    pub name: String,
32    pub public_key: PublicKey,
33}
34
35/// The data necessary to re-share the user-key to a V1 organization membership. Note: The
36/// Public-key must be verified/trusted. Further, there is no sender authentication possible here.
37#[derive(Serialize, Deserialize, Clone)]
38#[cfg_attr(feature = "wasm", derive(Tsify), tsify(into_wasm_abi, from_wasm_abi))]
39pub struct V1OrganizationMembership {
40    pub organization_id: uuid::Uuid,
41    pub name: String,
42    pub public_key: PublicKey,
43}
44
45#[derive(Debug)]
46pub(super) enum ReencryptError {
47    /// Failed to update the unlock data for the master password
48    MasterPasswordDerivation,
49    /// Failed to update the unlock data for TDE/PRF-Passkey
50    KeysetUnlockDataReencryption,
51    /// Failed to update the unlock data for emergency access or organization membership
52    KeySharingError,
53    /// Failed to wrap the user key with the Key Connector key
54    KeyConnectorWrapping,
55    /// Failed to create v2 upgrade token
56    UpgradeTokenCreation,
57}
58
59pub(super) struct ReencryptMasterPasswordChangeAndUnlockInput {
60    /// Master password change data.
61    pub(super) password: String,
62    pub(super) hint: Option<String>,
63    pub(super) kdf: Kdf,
64    pub(super) salt: String,
65    /// Common unlock data to re-encrypt
66    pub(super) common_unlock_data: ReencryptCommonUnlockDataInput,
67}
68
69pub(super) struct ReencryptCommonUnlockDataInput {
70    /// The trusted device keysets.
71    pub(super) trusted_devices: Vec<PartialRotateableKeyset>,
72    /// The webauthn credential keysets.
73    pub(super) webauthn_credentials: Vec<PartialRotateableKeyset>,
74    /// The V1 organization memberships.
75    pub(super) trusted_organization_keys: Vec<V1OrganizationMembership>,
76    /// The V1 emergency access memberships.
77    pub(super) trusted_emergency_access_keys: Vec<V1EmergencyAccessMembership>,
78}
79
80pub(super) fn reencrypt_master_password_change_unlock_data(
81    input: ReencryptMasterPasswordChangeAndUnlockInput,
82    current_user_key_id: SymmetricKeySlotId,
83    new_user_key_id: SymmetricKeySlotId,
84    ctx: &mut KeyStoreContext<KeySlotIds>,
85) -> Result<UnlockDataRequestModel, ReencryptError> {
86    let master_password_unlock_data = reencrypt_userkey_for_masterpassword_unlock(
87        input.password,
88        input.hint,
89        input.kdf,
90        input.salt,
91        new_user_key_id,
92        ctx,
93    )?;
94
95    let common_unlock_data = reencrypt_common_unlock_data(
96        input.common_unlock_data,
97        current_user_key_id,
98        new_user_key_id,
99        UpgradeTokenAction::Skip,
100        ctx,
101    )?;
102
103    Ok(UnlockDataRequestModel {
104        master_password_unlock_data: Box::new(master_password_unlock_data),
105        emergency_access_unlock_data: common_unlock_data.emergency_access_unlock_data,
106        organization_account_recovery_unlock_data: common_unlock_data
107            .organization_account_recovery_unlock_data,
108        passkey_unlock_data: common_unlock_data.passkey_unlock_data,
109        device_key_unlock_data: common_unlock_data.device_key_unlock_data,
110        // Master password change + key rotation always needs to logout other sessions, so no
111        // upgrade token is created.
112        v2_upgrade_token: None,
113    })
114}
115
116pub(super) fn reencrypt_common_unlock_data(
117    input: ReencryptCommonUnlockDataInput,
118    current_user_key_id: SymmetricKeySlotId,
119    new_user_key_id: SymmetricKeySlotId,
120    upgrade_token_action: UpgradeTokenAction,
121    ctx: &mut KeyStoreContext<KeySlotIds>,
122) -> Result<CommonUnlockDataRequestModel, ReencryptError> {
123    let tde_device_unlock_data = reencrypt_tde_devices(
124        &input.trusted_devices,
125        current_user_key_id,
126        new_user_key_id,
127        ctx,
128    )?;
129    let prf_passkey_unlock_data = reencrypt_passkey_credentials(
130        &input.webauthn_credentials,
131        current_user_key_id,
132        new_user_key_id,
133        ctx,
134    )?;
135    let emergency_accesses =
136        reencrypt_emergency_access_keys(input.trusted_emergency_access_keys, new_user_key_id, ctx)?;
137    let organizations_memberships =
138        reencrypt_organization_memberships(input.trusted_organization_keys, new_user_key_id, ctx)?;
139
140    let upgrade_token = make_upgrade_token_if_needed(
141        current_user_key_id,
142        new_user_key_id,
143        upgrade_token_action,
144        ctx,
145    )?;
146
147    Ok(CommonUnlockDataRequestModel {
148        emergency_access_unlock_data: Some(emergency_accesses),
149        organization_account_recovery_unlock_data: Some(organizations_memberships),
150        passkey_unlock_data: Some(prf_passkey_unlock_data),
151        device_key_unlock_data: Some(tde_device_unlock_data),
152        v2_upgrade_token: upgrade_token,
153    })
154}
155
156/// Re-encrypt TDE device keys for the new user key.
157fn reencrypt_tde_devices(
158    trusted_devices: &[PartialRotateableKeyset],
159    current_user_key_id: SymmetricKeySlotId,
160    new_user_key_id: SymmetricKeySlotId,
161    ctx: &mut KeyStoreContext<KeySlotIds>,
162) -> Result<Vec<OtherDeviceKeysUpdateRequestModel>, ReencryptError> {
163    trusted_devices
164        .iter()
165        .map(|device| {
166            let _span = debug_span!("reencrypt_device_key", device_id = ?device.id).entered();
167            device
168                .rotate_userkey(current_user_key_id, new_user_key_id, ctx)
169                .map_err(|_| ReencryptError::KeysetUnlockDataReencryption)
170                .map(Into::into)
171        })
172        .collect()
173}
174
175/// Re-encrypt passkey (WebAuthn PRF) credentials for the new user key.
176fn reencrypt_passkey_credentials(
177    webauthn_credentials: &[PartialRotateableKeyset],
178    current_user_key_id: SymmetricKeySlotId,
179    new_user_key_id: SymmetricKeySlotId,
180    ctx: &mut KeyStoreContext<KeySlotIds>,
181) -> Result<Vec<WebAuthnLoginRotateKeyRequestModel>, ReencryptError> {
182    webauthn_credentials
183        .iter()
184        .map(|cred| {
185            let _span =
186                debug_span!("reencrypt_webauthn_credential", credential_id = ?cred.id).entered();
187            cred.rotate_userkey(current_user_key_id, new_user_key_id, ctx)
188                .map_err(|_| ReencryptError::KeysetUnlockDataReencryption)
189                .map(Into::into)
190        })
191        .collect()
192}
193
194/// Re-encrypt emergency access keys for the new user key.
195fn reencrypt_emergency_access_keys(
196    trusted_emergency_access_keys: Vec<V1EmergencyAccessMembership>,
197    new_user_key_id: SymmetricKeySlotId,
198    ctx: &mut KeyStoreContext<KeySlotIds>,
199) -> Result<Vec<EmergencyAccessWithIdRequestModel>, ReencryptError> {
200    trusted_emergency_access_keys
201        .into_iter()
202        .map(|ea| {
203            let _span =
204                debug_span!("reencrypt_emergency_access_key", grantee_id = ?ea.id).entered();
205            // Share the key to the organization. Note: No sender authentication
206            // and the passed in public-key must be verified/trusted.
207            match UnsignedSharedKey::encapsulate(new_user_key_id, &ea.public_key, ctx) {
208                Ok(reencrypted_key) => Ok(EmergencyAccessWithIdRequestModel {
209                    // Default value that is ignored on the server
210                    r#type: models::EmergencyAccessType::Takeover,
211                    // Default value that is ignored on the server
212                    wait_time_days: 1,
213                    id: ea.id,
214                    key_encrypted: reencrypted_key.to_string().into(),
215                }),
216                Err(_) => Err(ReencryptError::KeySharingError),
217            }
218        })
219        .collect()
220}
221
222/// Re-encrypt organization membership keys for the new user key.
223fn reencrypt_organization_memberships(
224    trusted_organization_keys: Vec<V1OrganizationMembership>,
225    new_user_key_id: SymmetricKeySlotId,
226    ctx: &mut KeyStoreContext<KeySlotIds>,
227) -> Result<Vec<ResetPasswordWithOrgIdRequestModel>, ReencryptError> {
228    trusted_organization_keys
229        .into_iter()
230        .map(|org_membership| {
231            let _span =
232                debug_span!("reencrypt_organization_key", organization = ?org_membership.organization_id)
233                    .entered();
234            // Share the key to the organization. Note: No sender authentication
235            // and the passed in public-key must be verified/trusted.
236            match UnsignedSharedKey::encapsulate(new_user_key_id, &org_membership.public_key, ctx) {
237                Ok(reencrypted_key) => Ok(ResetPasswordWithOrgIdRequestModel {
238                    reset_password_key: Some(reencrypted_key.to_string()),
239                    master_password_hash: None,
240                    organization_id: org_membership.organization_id,
241                }),
242                Err(_) => Err(ReencryptError::KeySharingError),
243            }
244        })
245        .collect()
246}
247
248fn reencrypt_userkey_for_masterpassword_unlock(
249    password: String,
250    hint: Option<String>,
251    kdf: Kdf,
252    salt: String,
253    new_user_key_id: SymmetricKeySlotId,
254    ctx: &mut KeyStoreContext<KeySlotIds>,
255) -> Result<MasterPasswordUnlockAndAuthenticationDataModel, ReencryptError> {
256    let _span = debug_span!("derive_master_password_unlock_data").entered();
257    let unlock_data =
258        MasterPasswordUnlockData::derive(&password, &kdf, &salt, new_user_key_id, ctx)
259            .map_err(|_| ReencryptError::MasterPasswordDerivation)?;
260    let authentication_data = MasterPasswordAuthenticationData::derive(&password, &kdf, &salt)
261        .map_err(|_| ReencryptError::MasterPasswordDerivation)?;
262    to_authentication_and_unlock_data(unlock_data, authentication_data, hint)
263        .map_err(|_| ReencryptError::MasterPasswordDerivation)
264}
265
266#[derive(Debug)]
267struct ParsingError;
268
269fn to_authentication_and_unlock_data(
270    master_password_unlock_data: MasterPasswordUnlockData,
271    master_password_authentication_data: MasterPasswordAuthenticationData,
272    hint: Option<String>,
273) -> Result<MasterPasswordUnlockAndAuthenticationDataModel, ParsingError> {
274    let (kdf_type, kdf_iterations, kdf_memory, kdf_parallelism) =
275        match master_password_unlock_data.kdf {
276            bitwarden_crypto::Kdf::PBKDF2 { iterations } => {
277                (models::KdfType::PBKDF2_SHA256, iterations, None, None)
278            }
279            bitwarden_crypto::Kdf::Argon2id {
280                iterations,
281                memory,
282                parallelism,
283            } => (
284                models::KdfType::Argon2id,
285                iterations,
286                Some(memory),
287                Some(parallelism),
288            ),
289        };
290    Ok(MasterPasswordUnlockAndAuthenticationDataModel {
291        kdf_type,
292        kdf_iterations: kdf_iterations.get().try_into().map_err(|_| ParsingError)?,
293        kdf_memory: kdf_memory
294            .map(|m| m.get().try_into().map_err(|_| ParsingError))
295            .transpose()?,
296        kdf_parallelism: kdf_parallelism
297            .map(|p| p.get().try_into().map_err(|_| ParsingError))
298            .transpose()?,
299        email: Some(master_password_unlock_data.salt.clone()),
300        master_key_authentication_hash: Some(
301            master_password_authentication_data
302                .master_password_authentication_hash
303                .to_string(),
304        ),
305        master_key_encrypted_user_key: Some(
306            master_password_unlock_data
307                .master_key_wrapped_user_key
308                .to_string(),
309        ),
310        master_password_hint: hint,
311        master_password_salt: Some(master_password_unlock_data.salt.clone()),
312    })
313}
314
315fn make_upgrade_token_if_needed(
316    current_user_key_id: SymmetricKeySlotId,
317    new_user_key_id: SymmetricKeySlotId,
318    upgrade_token_action: UpgradeTokenAction,
319    ctx: &mut KeyStoreContext<KeySlotIds>,
320) -> Result<Option<Box<V2UpgradeTokenRequestModel>>, ReencryptError> {
321    if matches!(upgrade_token_action, UpgradeTokenAction::Skip) {
322        debug!("UpgradeTokenAction::Skip, skipping upgrade token creation");
323        return Ok(None);
324    }
325
326    match (
327        ctx.get_symmetric_key_algorithm(current_user_key_id),
328        ctx.get_symmetric_key_algorithm(new_user_key_id),
329    ) {
330        (
331            Ok(SymmetricKeyAlgorithm::Aes256CbcHmac),
332            Ok(SymmetricKeyAlgorithm::XChaCha20Poly1305),
333        ) => {
334            let token =
335                V2UpgradeToken::create(current_user_key_id, new_user_key_id, ctx).map_err(|e| {
336                    error!("Failed to create V2 upgrade token: {e}");
337                    ReencryptError::UpgradeTokenCreation
338                })?;
339            info!("Upgrade token created for the key rotation");
340            Ok(Some(Box::new(token.into())))
341        }
342        _ => Ok(None),
343    }
344}
345
346#[cfg(test)]
347mod tests {
348    use std::num::NonZeroU32;
349
350    use bitwarden_api_api::models::KdfType;
351    use bitwarden_core::key_management::KeySlotIds;
352    use bitwarden_crypto::{Kdf, KeyStore, PublicKeyEncryptionAlgorithm, UnsignedSharedKey};
353    use uuid::Uuid;
354
355    use super::*;
356    use crate::key_rotation::partial_rotateable_keyset::PartialRotateableKeyset;
357
358    fn create_test_kdf_pbkdf2() -> Kdf {
359        Kdf::PBKDF2 {
360            iterations: NonZeroU32::new(600000).expect("valid iterations"),
361        }
362    }
363
364    fn create_test_kdf_argon2id() -> Kdf {
365        Kdf::Argon2id {
366            iterations: NonZeroU32::new(3).expect("valid iterations"),
367            memory: NonZeroU32::new(64).expect("valid memory"),
368            parallelism: NonZeroU32::new(4).expect("valid parallelism"),
369        }
370    }
371
372    fn assert_symmetric_keys_equal(
373        key_id_1: SymmetricKeySlotId,
374        key_id_2: SymmetricKeySlotId,
375        ctx: &mut KeyStoreContext<KeySlotIds>,
376    ) {
377        #[allow(deprecated)]
378        let key_1 = ctx
379            .dangerous_get_symmetric_key(key_id_1)
380            .expect("key 1 should exist");
381        #[allow(deprecated)]
382        let key_2 = ctx
383            .dangerous_get_symmetric_key(key_id_2)
384            .expect("key 2 should exist");
385        assert_eq!(key_1, key_2, "symmetric keys should be equal");
386    }
387
388    fn empty_common_unlock_input() -> ReencryptCommonUnlockDataInput {
389        ReencryptCommonUnlockDataInput {
390            trusted_devices: vec![],
391            webauthn_credentials: vec![],
392            trusted_organization_keys: vec![],
393            trusted_emergency_access_keys: vec![],
394        }
395    }
396
397    fn request_model_to_token(request: V2UpgradeTokenRequestModel) -> V2UpgradeToken {
398        V2UpgradeToken {
399            wrapped_user_key_1: request
400                .wrapped_user_key1
401                .parse()
402                .expect("wrapped_user_key1 should parse"),
403            wrapped_user_key_2: request
404                .wrapped_user_key2
405                .parse()
406                .expect("wrapped_user_key2 should parse"),
407        }
408    }
409
410    #[test]
411    fn test_to_authentication_and_unlock_data_pbkdf2() {
412        let store: KeyStore<KeySlotIds> = KeyStore::default();
413        let mut ctx = store.context_mut();
414
415        let kdf = create_test_kdf_pbkdf2();
416        let salt = "[email protected]";
417        let password = "test_password";
418
419        let user_key_id = ctx.generate_symmetric_key();
420        let unlock_data = MasterPasswordUnlockData::derive(password, &kdf, salt, user_key_id, &ctx)
421            .expect("derive should succeed");
422        let auth_data = MasterPasswordAuthenticationData::derive(password, &kdf, salt)
423            .expect("derive should succeed");
424
425        let result = to_authentication_and_unlock_data(unlock_data, auth_data, None);
426        assert!(result.is_ok());
427
428        let model = result.expect("should be ok");
429        assert_eq!(model.kdf_type, KdfType::PBKDF2_SHA256);
430        assert_eq!(model.kdf_iterations, 600000);
431        assert!(model.kdf_memory.is_none());
432        assert!(model.kdf_parallelism.is_none());
433        assert_eq!(model.email, Some(salt.to_string()));
434        assert!(model.master_key_authentication_hash.is_some());
435        assert!(model.master_key_encrypted_user_key.is_some());
436        assert!(model.master_password_hint.is_none());
437
438        // Verify the unlock data can decrypt the user key
439        let master_password_unlock_data = MasterPasswordUnlockData {
440            master_key_wrapped_user_key: model
441                .master_key_encrypted_user_key
442                .expect("should be present")
443                .parse()
444                .expect("should parse"),
445            kdf: kdf.clone(),
446            salt: salt.to_string(),
447        };
448        let decrypted_user_key = master_password_unlock_data
449            .unwrap_to_context(password, &mut ctx)
450            .expect("unwrap should succeed");
451        assert_symmetric_keys_equal(user_key_id, decrypted_user_key, &mut ctx);
452    }
453
454    #[test]
455    fn test_to_authentication_and_unlock_data_argon2id() {
456        let store: KeyStore<KeySlotIds> = KeyStore::default();
457        let mut ctx = store.context_mut();
458
459        let kdf = create_test_kdf_argon2id();
460        let salt = "[email protected]";
461        let password = "test_password";
462
463        let user_key_id = ctx.generate_symmetric_key();
464        let unlock_data = MasterPasswordUnlockData::derive(password, &kdf, salt, user_key_id, &ctx)
465            .expect("derive should succeed");
466        let auth_data = MasterPasswordAuthenticationData::derive(password, &kdf, salt)
467            .expect("derive should succeed");
468
469        let result = to_authentication_and_unlock_data(unlock_data, auth_data, None);
470        assert!(result.is_ok());
471
472        let model = result.expect("should be ok");
473        assert_eq!(model.kdf_type, KdfType::Argon2id);
474        assert_eq!(model.kdf_iterations, 3);
475        assert_eq!(model.kdf_memory, Some(64));
476        assert_eq!(model.kdf_parallelism, Some(4));
477        assert_eq!(model.email, Some(salt.to_string()));
478        assert!(model.master_key_authentication_hash.is_some());
479        assert!(model.master_key_encrypted_user_key.is_some());
480
481        // Verify the unlock data can decrypt the user key
482        let master_password_unlock_data = MasterPasswordUnlockData {
483            master_key_wrapped_user_key: model
484                .master_key_encrypted_user_key
485                .expect("should be present")
486                .parse()
487                .expect("should parse"),
488            kdf: kdf.clone(),
489            salt: salt.to_string(),
490        };
491        let decrypted_user_key = master_password_unlock_data
492            .unwrap_to_context(password, &mut ctx)
493            .expect("unwrap should succeed");
494        assert_symmetric_keys_equal(user_key_id, decrypted_user_key, &mut ctx);
495    }
496
497    #[test]
498    fn test_reencrypt_unlock_device_key_data() {
499        let store: KeyStore<KeySlotIds> = KeyStore::default();
500        let mut ctx = store.context_mut();
501
502        let current_user_key_id = ctx.generate_symmetric_key();
503        let new_user_key_id = ctx.generate_symmetric_key();
504
505        let (device_keyset, device_private_key) =
506            PartialRotateableKeyset::make_test_keyset(current_user_key_id, &mut ctx);
507
508        let result = reencrypt_common_unlock_data(
509            ReencryptCommonUnlockDataInput {
510                trusted_devices: vec![device_keyset],
511                webauthn_credentials: vec![],
512                trusted_organization_keys: vec![],
513                trusted_emergency_access_keys: vec![],
514            },
515            current_user_key_id,
516            new_user_key_id,
517            UpgradeTokenAction::CreateIfNeeded,
518            &mut ctx,
519        );
520
521        let unlock_data = result.expect("should be ok");
522
523        let device_unlock = unlock_data
524            .device_key_unlock_data
525            .as_ref()
526            .expect("should be present")
527            .first()
528            .expect("should have at least one");
529        let decrypted_user_key = device_unlock
530            .encrypted_user_key
531            .parse::<UnsignedSharedKey>()
532            .expect("should parse")
533            .decapsulate(device_private_key, &mut ctx)
534            .expect("unwrap should succeed");
535        assert_symmetric_keys_equal(new_user_key_id, decrypted_user_key, &mut ctx);
536    }
537
538    #[test]
539    fn test_reencrypt_unlock_webauthn_prf_credential_data() {
540        let store: KeyStore<KeySlotIds> = KeyStore::default();
541        let mut ctx = store.context_mut();
542
543        let current_user_key_id = ctx.generate_symmetric_key();
544        let new_user_key_id = ctx.generate_symmetric_key();
545
546        let (credential_keyset, credential_private_key) =
547            PartialRotateableKeyset::make_test_keyset(current_user_key_id, &mut ctx);
548
549        let result = reencrypt_common_unlock_data(
550            ReencryptCommonUnlockDataInput {
551                trusted_devices: vec![],
552                webauthn_credentials: vec![credential_keyset],
553                trusted_organization_keys: vec![],
554                trusted_emergency_access_keys: vec![],
555            },
556            current_user_key_id,
557            new_user_key_id,
558            UpgradeTokenAction::CreateIfNeeded,
559            &mut ctx,
560        );
561
562        let unlock_data = result.expect("should be ok");
563
564        // Ensure it decrypts to the correct key after rotation
565        let credential_unlock = unlock_data
566            .passkey_unlock_data
567            .as_ref()
568            .expect("should be present")
569            .first()
570            .expect("should have at least one");
571        let decrypted_user_key = credential_unlock
572            .encrypted_user_key
573            .parse::<UnsignedSharedKey>()
574            .expect("should parse")
575            .decapsulate(credential_private_key, &mut ctx)
576            .expect("unwrap should succeed");
577        assert_symmetric_keys_equal(new_user_key_id, decrypted_user_key, &mut ctx);
578    }
579
580    #[test]
581    fn test_reencrypt_unlock_emergency_access_data() {
582        let store: KeyStore<KeySlotIds> = KeyStore::default();
583        let mut ctx = store.context_mut();
584
585        let current_user_key_id = ctx.generate_symmetric_key();
586        let new_user_key_id = ctx.generate_symmetric_key();
587
588        let organization_private_key =
589            ctx.make_private_key(PublicKeyEncryptionAlgorithm::RsaOaepSha1);
590        let emergency_access = V1EmergencyAccessMembership {
591            id: Uuid::new_v4(),
592            grantee_id: Uuid::new_v4(),
593            name: "Test User".to_string(),
594            public_key: ctx
595                .get_public_key(organization_private_key)
596                .expect("key exists"),
597        };
598
599        let result = reencrypt_common_unlock_data(
600            ReencryptCommonUnlockDataInput {
601                trusted_devices: vec![],
602                webauthn_credentials: vec![],
603                trusted_organization_keys: vec![],
604                trusted_emergency_access_keys: vec![emergency_access],
605            },
606            current_user_key_id,
607            new_user_key_id,
608            UpgradeTokenAction::CreateIfNeeded,
609            &mut ctx,
610        );
611
612        let unlock_data = result.expect("should be ok");
613
614        // Ensure it decrypts to the correct key after rotation
615        let emergency_access_unlock = unlock_data
616            .emergency_access_unlock_data
617            .as_ref()
618            .expect("should be present")
619            .first()
620            .expect("should have at least one");
621        let decrypted_user_key = emergency_access_unlock
622            .key_encrypted
623            .as_ref()
624            .map(|k| k.parse::<UnsignedSharedKey>())
625            .expect("should be present")
626            .expect("should parse")
627            .decapsulate(organization_private_key, &mut ctx)
628            .expect("unwrap should succeed");
629        assert_symmetric_keys_equal(new_user_key_id, decrypted_user_key, &mut ctx);
630    }
631
632    #[test]
633    fn test_reencrypt_unlock_organization_membership_data() {
634        let store: KeyStore<KeySlotIds> = KeyStore::default();
635        let mut ctx = store.context_mut();
636
637        let current_user_key_id = ctx.generate_symmetric_key();
638        let new_user_key_id = ctx.generate_symmetric_key();
639
640        let org_key = ctx.make_private_key(PublicKeyEncryptionAlgorithm::RsaOaepSha1);
641        let org_membership = V1OrganizationMembership {
642            organization_id: Uuid::new_v4(),
643            name: "Test Org".to_string(),
644            public_key: ctx.get_public_key(org_key).expect("key exists"),
645        };
646
647        let result = reencrypt_common_unlock_data(
648            ReencryptCommonUnlockDataInput {
649                trusted_devices: vec![],
650                webauthn_credentials: vec![],
651                trusted_organization_keys: vec![org_membership],
652                trusted_emergency_access_keys: vec![],
653            },
654            current_user_key_id,
655            new_user_key_id,
656            UpgradeTokenAction::CreateIfNeeded,
657            &mut ctx,
658        );
659
660        let unlock_data = result.expect("should be ok");
661
662        let org_membership_unlock = unlock_data
663            .organization_account_recovery_unlock_data
664            .as_ref()
665            .expect("should be present")
666            .first()
667            .expect("should have at least one");
668        let decrypted_user_key = org_membership_unlock
669            .reset_password_key
670            .as_ref()
671            .map(|k| k.parse::<UnsignedSharedKey>())
672            .expect("should be present")
673            .expect("should parse")
674            .decapsulate(org_key, &mut ctx)
675            .expect("unwrap should succeed");
676        assert_symmetric_keys_equal(new_user_key_id, decrypted_user_key, &mut ctx);
677    }
678
679    #[test]
680    fn test_reencrypt_common_unlock_data_v1_to_v2_creates_upgrade_token() {
681        let store: KeyStore<KeySlotIds> = KeyStore::default();
682        let mut ctx = store.context_mut();
683
684        let current_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
685        let new_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::XChaCha20Poly1305);
686
687        let result = reencrypt_common_unlock_data(
688            empty_common_unlock_input(),
689            current_user_key_id,
690            new_user_key_id,
691            UpgradeTokenAction::CreateIfNeeded,
692            &mut ctx,
693        );
694
695        let unlock_data = result.expect("should be ok");
696        let token_request = *unlock_data
697            .v2_upgrade_token
698            .expect("v2_upgrade_token should be populated for V1 -> V2 rotation");
699        let token = request_model_to_token(token_request);
700
701        let unwrapped_v2_id = token
702            .unwrap_v2(current_user_key_id, &mut ctx)
703            .expect("unwrap_v2 should succeed");
704        assert_symmetric_keys_equal(new_user_key_id, unwrapped_v2_id, &mut ctx);
705
706        let unwrapped_v1_id = token
707            .unwrap_v1(new_user_key_id, &mut ctx)
708            .expect("unwrap_v1 should succeed");
709        assert_symmetric_keys_equal(current_user_key_id, unwrapped_v1_id, &mut ctx);
710    }
711
712    #[test]
713    fn test_reencrypt_common_unlock_data_v1_to_v2_upgrade_token_action_skip_returns_none() {
714        let store: KeyStore<KeySlotIds> = KeyStore::default();
715        let mut ctx = store.context_mut();
716
717        let current_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
718        let new_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::XChaCha20Poly1305);
719
720        let result = reencrypt_common_unlock_data(
721            empty_common_unlock_input(),
722            current_user_key_id,
723            new_user_key_id,
724            UpgradeTokenAction::Skip,
725            &mut ctx,
726        );
727
728        let unlock_data = result.expect("should be ok");
729        assert!(
730            unlock_data.v2_upgrade_token.is_none(),
731            "UpgradeTokenAction::Skip skips the creation of the upgrade token"
732        );
733    }
734
735    #[test]
736    fn test_reencrypt_common_unlock_data_v2_to_v2_upgrade_token_action_create_if_needed_returns_none()
737     {
738        let store: KeyStore<KeySlotIds> = KeyStore::default();
739        let mut ctx = store.context_mut();
740
741        let current_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::XChaCha20Poly1305);
742        let new_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::XChaCha20Poly1305);
743
744        let result = reencrypt_common_unlock_data(
745            empty_common_unlock_input(),
746            current_user_key_id,
747            new_user_key_id,
748            UpgradeTokenAction::CreateIfNeeded,
749            &mut ctx,
750        );
751
752        let unlock_data = result.expect("should be ok");
753        assert!(
754            unlock_data.v2_upgrade_token.is_none(),
755            "UpgradeTokenAction::CreateIfNeeded should not create a v2_upgrade_token for V2 -> V2 rotation"
756        );
757    }
758
759    #[test]
760    fn test_reencrypt_master_password_change_unlock_data_never_returns_upgrade_token() {
761        let store: KeyStore<KeySlotIds> = KeyStore::default();
762        let mut ctx = store.context_mut();
763
764        let current_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
765        let new_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::XChaCha20Poly1305);
766
767        let input = ReencryptMasterPasswordChangeAndUnlockInput {
768            password: "test_password".to_string(),
769            hint: None,
770            kdf: create_test_kdf_pbkdf2(),
771            salt: "[email protected]".to_string(),
772            common_unlock_data: empty_common_unlock_input(),
773        };
774
775        let unlock_data = reencrypt_master_password_change_unlock_data(
776            input,
777            current_user_key_id,
778            new_user_key_id,
779            &mut ctx,
780        )
781        .expect("should be ok");
782
783        assert!(
784            unlock_data.v2_upgrade_token.is_none(),
785            "master password change rotation must never include a v2 upgrade token"
786        );
787    }
788}