bitwarden_user_crypto_management/key_rotation/

unlock.rs

//! Functionality for re-encrypting unlock (decryption) methods during user key rotation.
//! During key-rotation, a new user-key is sampled. The unlock module then creates a set of newly
//! encrypted copies, one for each decryption/unlock method.
use std::str::FromStr;

use bitwarden_api_api::models::{
    self, CommonUnlockDataRequestModel, EmergencyAccessKeyDataResponseModel,
    EmergencyAccessWithIdRequestModel, MasterPasswordUnlockAndAuthenticationDataModel,
    OrganizationPasswordResetKeyDataResponseModel, OtherDeviceKeysUpdateRequestModel,
    ResetPasswordWithOrgIdRequestModel, UnlockDataRequestModel, V2UpgradeTokenRequestModel,
    WebAuthnLoginRotateKeyRequestModel,
};
use bitwarden_core::{
    key_management::{
        KeySlotIds, MasterPasswordAuthenticationData, MasterPasswordUnlockData, SymmetricKeySlotId,
        V2UpgradeToken,
    },
    require,
};
use bitwarden_crypto::{
    Kdf, KeyStoreContext, PublicKey, SpkiPublicKeyBytes, SymmetricKeyAlgorithm, UnsignedSharedKey,
};
use bitwarden_encoding::B64;
use serde::{Deserialize, Serialize};
use tracing::{debug, debug_span, error, info};
#[cfg(feature = "wasm")]
use tsify::Tsify;

use crate::key_rotation::{
    KeyRotationDataParseError, partial_rotateable_keyset::PartialRotateableKeyset,
    rotate_user_keys::UpgradeTokenAction,
};

/// The data necessary to re-share the user-key to a V1 emergency access membership. Note: The
/// Public-key must be verified/trusted. Further, there is no sender authentication possible here.
#[derive(Serialize, Deserialize, Clone)]
#[cfg_attr(feature = "wasm", derive(Tsify), tsify(into_wasm_abi, from_wasm_abi))]
pub struct V1EmergencyAccessMembership {
    pub id: uuid::Uuid,
    pub grantee_id: uuid::Uuid,
    pub name: String,
    pub public_key: PublicKey,
}

impl TryFrom<EmergencyAccessKeyDataResponseModel> for V1EmergencyAccessMembership {
    type Error = KeyRotationDataParseError;

    fn try_from(ea: EmergencyAccessKeyDataResponseModel) -> Result<Self, Self::Error> {
        Ok(Self {
            id: require!(ea.id),
            grantee_id: require!(ea.grantee_id),
            // The name can be null if a user does not set a name; fall back to the email and
            // then to "Unknown" so we always have a non-empty display name.
            name: ea
                .grantee_name
                .or(ea.grantee_email)
                .unwrap_or_else(|| "Unknown".to_string()),
            public_key: parse_public_key(&require!(ea.public_key))?,
        })
    }
}

/// The data necessary to re-share the user-key to a V1 organization membership. Note: The
/// Public-key must be verified/trusted. Further, there is no sender authentication possible here.
#[derive(Serialize, Deserialize, Clone)]
#[cfg_attr(feature = "wasm", derive(Tsify), tsify(into_wasm_abi, from_wasm_abi))]
pub struct V1OrganizationMembership {
    pub organization_id: uuid::Uuid,
    pub name: String,
    pub public_key: PublicKey,
}

impl TryFrom<OrganizationPasswordResetKeyDataResponseModel> for V1OrganizationMembership {
    type Error = KeyRotationDataParseError;

    fn try_from(o: OrganizationPasswordResetKeyDataResponseModel) -> Result<Self, Self::Error> {
        Ok(Self {
            organization_id: require!(o.organization_id),
            name: require!(o.organization_name),
            public_key: parse_public_key(&require!(o.organization_public_key))?,
        })
    }
}

fn parse_public_key(public_key_b64: &str) -> Result<PublicKey, KeyRotationDataParseError> {
    Ok(PublicKey::from_der(&SpkiPublicKeyBytes::from(
        B64::from_str(public_key_b64)?.into_bytes(),
    ))?)
}

#[derive(Debug)]
pub(super) enum ReencryptError {
    /// Failed to update the unlock data for the master password
    MasterPasswordDerivation,
    /// Failed to update the unlock data for TDE/PRF-Passkey
    KeysetUnlockDataReencryption,
    /// Failed to update the unlock data for emergency access or organization membership
    KeySharingError,
    /// Failed to wrap the user key with the Key Connector key
    KeyConnectorWrapping,
    /// Failed to create v2 upgrade token
    UpgradeTokenCreation,
}

pub(super) struct ReencryptMasterPasswordChangeAndUnlockInput {
    /// Master password change data.
    pub(super) password: String,
    pub(super) hint: Option<String>,
    pub(super) kdf: Kdf,
    pub(super) salt: String,
    /// Common unlock data to re-encrypt
    pub(super) common_unlock_data: ReencryptCommonUnlockDataInput,
}

pub(super) struct ReencryptCommonUnlockDataInput {
    /// The trusted device keysets.
    pub(super) trusted_devices: Vec<PartialRotateableKeyset>,
    /// The webauthn credential keysets.
    pub(super) webauthn_credentials: Vec<PartialRotateableKeyset>,
    /// The V1 organization memberships.
    pub(super) trusted_organization_keys: Vec<V1OrganizationMembership>,
    /// The V1 emergency access memberships.
    pub(super) trusted_emergency_access_keys: Vec<V1EmergencyAccessMembership>,
}

pub(super) fn reencrypt_master_password_change_unlock_data(
    input: ReencryptMasterPasswordChangeAndUnlockInput,
    current_user_key_id: SymmetricKeySlotId,
    new_user_key_id: SymmetricKeySlotId,
    ctx: &mut KeyStoreContext<KeySlotIds>,
) -> Result<UnlockDataRequestModel, ReencryptError> {
    let master_password_unlock_data = reencrypt_userkey_for_masterpassword_unlock(
        input.password,
        input.hint,
        input.kdf,
        input.salt,
        new_user_key_id,
        ctx,
    )?;

    let common_unlock_data = reencrypt_common_unlock_data(
        input.common_unlock_data,
        current_user_key_id,
        new_user_key_id,
        UpgradeTokenAction::Skip,
        ctx,
    )?;

    Ok(UnlockDataRequestModel {
        master_password_unlock_data: Box::new(master_password_unlock_data),
        emergency_access_unlock_data: common_unlock_data.emergency_access_unlock_data,
        organization_account_recovery_unlock_data: common_unlock_data
            .organization_account_recovery_unlock_data,
        passkey_unlock_data: common_unlock_data.passkey_unlock_data,
        device_key_unlock_data: common_unlock_data.device_key_unlock_data,
        // Master password change + key rotation always needs to logout other sessions, so no
        // upgrade token is created.
        v2_upgrade_token: None,
    })
}

pub(super) fn reencrypt_common_unlock_data(
    input: ReencryptCommonUnlockDataInput,
    current_user_key_id: SymmetricKeySlotId,
    new_user_key_id: SymmetricKeySlotId,
    upgrade_token_action: UpgradeTokenAction,
    ctx: &mut KeyStoreContext<KeySlotIds>,
) -> Result<CommonUnlockDataRequestModel, ReencryptError> {
    let tde_device_unlock_data = reencrypt_tde_devices(
        &input.trusted_devices,
        current_user_key_id,
        new_user_key_id,
        ctx,
    )?;
    let prf_passkey_unlock_data = reencrypt_passkey_credentials(
        &input.webauthn_credentials,
        current_user_key_id,
        new_user_key_id,
        ctx,
    )?;
    let emergency_accesses =
        reencrypt_emergency_access_keys(input.trusted_emergency_access_keys, new_user_key_id, ctx)?;
    let organizations_memberships =
        reencrypt_organization_memberships(input.trusted_organization_keys, new_user_key_id, ctx)?;

    let upgrade_token = make_upgrade_token_if_needed(
        current_user_key_id,
        new_user_key_id,
        upgrade_token_action,
        ctx,
    )?;

    Ok(CommonUnlockDataRequestModel {
        emergency_access_unlock_data: Some(emergency_accesses),
        organization_account_recovery_unlock_data: Some(organizations_memberships),
        passkey_unlock_data: Some(prf_passkey_unlock_data),
        device_key_unlock_data: Some(tde_device_unlock_data),
        v2_upgrade_token: upgrade_token,
    })
}

/// Re-encrypt TDE device keys for the new user key.
fn reencrypt_tde_devices(
    trusted_devices: &[PartialRotateableKeyset],
    current_user_key_id: SymmetricKeySlotId,
    new_user_key_id: SymmetricKeySlotId,
    ctx: &mut KeyStoreContext<KeySlotIds>,
) -> Result<Vec<OtherDeviceKeysUpdateRequestModel>, ReencryptError> {
    trusted_devices
        .iter()
        .map(|device| {
            let _span = debug_span!("reencrypt_device_key", device_id = ?device.id).entered();
            device
                .rotate_userkey(current_user_key_id, new_user_key_id, ctx)
                .map_err(|_| ReencryptError::KeysetUnlockDataReencryption)
                .map(Into::into)
        })
        .collect()
}

/// Re-encrypt passkey (WebAuthn PRF) credentials for the new user key.
fn reencrypt_passkey_credentials(
    webauthn_credentials: &[PartialRotateableKeyset],
    current_user_key_id: SymmetricKeySlotId,
    new_user_key_id: SymmetricKeySlotId,
    ctx: &mut KeyStoreContext<KeySlotIds>,
) -> Result<Vec<WebAuthnLoginRotateKeyRequestModel>, ReencryptError> {
    webauthn_credentials
        .iter()
        .map(|cred| {
            let _span =
                debug_span!("reencrypt_webauthn_credential", credential_id = ?cred.id).entered();
            cred.rotate_userkey(current_user_key_id, new_user_key_id, ctx)
                .map_err(|_| ReencryptError::KeysetUnlockDataReencryption)
                .map(Into::into)
        })
        .collect()
}

/// Re-encrypt emergency access keys for the new user key.
fn reencrypt_emergency_access_keys(
    trusted_emergency_access_keys: Vec<V1EmergencyAccessMembership>,
    new_user_key_id: SymmetricKeySlotId,
    ctx: &mut KeyStoreContext<KeySlotIds>,
) -> Result<Vec<EmergencyAccessWithIdRequestModel>, ReencryptError> {
    trusted_emergency_access_keys
        .into_iter()
        .map(|ea| {
            let _span =
                debug_span!("reencrypt_emergency_access_key", grantee_id = ?ea.id).entered();
            // Share the key to the organization. Note: No sender authentication
            // and the passed in public-key must be verified/trusted.
            match UnsignedSharedKey::encapsulate(new_user_key_id, &ea.public_key, ctx) {
                Ok(reencrypted_key) => Ok(EmergencyAccessWithIdRequestModel {
                    // Default value that is ignored on the server
                    r#type: models::EmergencyAccessType::Takeover,
                    // Default value that is ignored on the server
                    wait_time_days: 1,
                    id: ea.id,
                    key_encrypted: reencrypted_key.to_string().into(),
                }),
                Err(_) => Err(ReencryptError::KeySharingError),
            }
        })
        .collect()
}

/// Re-encrypt organization membership keys for the new user key.
fn reencrypt_organization_memberships(
    trusted_organization_keys: Vec<V1OrganizationMembership>,
    new_user_key_id: SymmetricKeySlotId,
    ctx: &mut KeyStoreContext<KeySlotIds>,
) -> Result<Vec<ResetPasswordWithOrgIdRequestModel>, ReencryptError> {
    trusted_organization_keys
        .into_iter()
        .map(|org_membership| {
            let _span =
                debug_span!("reencrypt_organization_key", organization = ?org_membership.organization_id)
                    .entered();
            // Share the key to the organization. Note: No sender authentication
            // and the passed in public-key must be verified/trusted.
            match UnsignedSharedKey::encapsulate(new_user_key_id, &org_membership.public_key, ctx) {
                Ok(reencrypted_key) => Ok(ResetPasswordWithOrgIdRequestModel {
                    reset_password_key: Some(reencrypted_key.to_string()),
                    master_password_hash: None,
                    organization_id: org_membership.organization_id,
                }),
                Err(_) => Err(ReencryptError::KeySharingError),
            }
        })
        .collect()
}

fn reencrypt_userkey_for_masterpassword_unlock(
    password: String,
    hint: Option<String>,
    kdf: Kdf,
    salt: String,
    new_user_key_id: SymmetricKeySlotId,
    ctx: &mut KeyStoreContext<KeySlotIds>,
) -> Result<MasterPasswordUnlockAndAuthenticationDataModel, ReencryptError> {
    let _span = debug_span!("derive_master_password_unlock_data").entered();
    let unlock_data =
        MasterPasswordUnlockData::derive(&password, &kdf, &salt, new_user_key_id, ctx)
            .map_err(|_| ReencryptError::MasterPasswordDerivation)?;
    let authentication_data = MasterPasswordAuthenticationData::derive(&password, &kdf, &salt)
        .map_err(|_| ReencryptError::MasterPasswordDerivation)?;
    to_authentication_and_unlock_data(unlock_data, authentication_data, hint)
        .map_err(|_| ReencryptError::MasterPasswordDerivation)
}

#[derive(Debug)]
struct ParsingError;

fn to_authentication_and_unlock_data(
    master_password_unlock_data: MasterPasswordUnlockData,
    master_password_authentication_data: MasterPasswordAuthenticationData,
    hint: Option<String>,
) -> Result<MasterPasswordUnlockAndAuthenticationDataModel, ParsingError> {
    let (kdf_type, kdf_iterations, kdf_memory, kdf_parallelism) =
        match master_password_unlock_data.kdf {
            bitwarden_crypto::Kdf::PBKDF2 { iterations } => {
                (models::KdfType::PBKDF2_SHA256, iterations, None, None)
            }
            bitwarden_crypto::Kdf::Argon2id {
                iterations,
                memory,
                parallelism,
            } => (
                models::KdfType::Argon2id,
                iterations,
                Some(memory),
                Some(parallelism),
            ),
        };
    Ok(MasterPasswordUnlockAndAuthenticationDataModel {
        kdf_type,
        kdf_iterations: kdf_iterations.get().try_into().map_err(|_| ParsingError)?,
        kdf_memory: kdf_memory
            .map(|m| m.get().try_into().map_err(|_| ParsingError))
            .transpose()?,
        kdf_parallelism: kdf_parallelism
            .map(|p| p.get().try_into().map_err(|_| ParsingError))
            .transpose()?,
        email: Some(master_password_unlock_data.salt.clone()),
        master_key_authentication_hash: Some(
            master_password_authentication_data
                .master_password_authentication_hash
                .to_string(),
        ),
        master_key_encrypted_user_key: Some(
            master_password_unlock_data
                .master_key_wrapped_user_key
                .to_string(),
        ),
        master_password_hint: hint,
        master_password_salt: Some(master_password_unlock_data.salt.clone()),
    })
}

fn make_upgrade_token_if_needed(
    current_user_key_id: SymmetricKeySlotId,
    new_user_key_id: SymmetricKeySlotId,
    upgrade_token_action: UpgradeTokenAction,
    ctx: &mut KeyStoreContext<KeySlotIds>,
) -> Result<Option<Box<V2UpgradeTokenRequestModel>>, ReencryptError> {
    if matches!(upgrade_token_action, UpgradeTokenAction::Skip) {
        debug!("UpgradeTokenAction::Skip, skipping upgrade token creation");
        return Ok(None);
    }

    match (
        ctx.get_symmetric_key_algorithm(current_user_key_id),
        ctx.get_symmetric_key_algorithm(new_user_key_id),
    ) {
        (Ok(SymmetricKeyAlgorithm::Aes256CbcHmac), Ok(SymmetricKeyAlgorithm::XAes256Gcm)) => {
            let token =
                V2UpgradeToken::create(current_user_key_id, new_user_key_id, ctx).map_err(|e| {
                    error!("Failed to create V2 upgrade token: {e}");
                    ReencryptError::UpgradeTokenCreation
                })?;
            info!("Upgrade token created for the key rotation");
            Ok(Some(Box::new(token.into())))
        }
        _ => Ok(None),
    }
}

#[cfg(test)]
mod tests {
    use std::num::NonZeroU32;

    use bitwarden_api_api::models::KdfType;
    use bitwarden_core::key_management::KeySlotIds;
    use bitwarden_crypto::{Kdf, KeyStore, PublicKeyEncryptionAlgorithm, UnsignedSharedKey};
    use uuid::Uuid;

    use super::*;
    use crate::key_rotation::partial_rotateable_keyset::PartialRotateableKeyset;

    fn create_test_kdf_pbkdf2() -> Kdf {
        Kdf::PBKDF2 {
            iterations: NonZeroU32::new(600000).expect("valid iterations"),
        }
    }

    fn create_test_kdf_argon2id() -> Kdf {
        Kdf::Argon2id {
            iterations: NonZeroU32::new(3).expect("valid iterations"),
            memory: NonZeroU32::new(64).expect("valid memory"),
            parallelism: NonZeroU32::new(4).expect("valid parallelism"),
        }
    }

    fn assert_symmetric_keys_equal(
        key_id_1: SymmetricKeySlotId,
        key_id_2: SymmetricKeySlotId,
        ctx: &mut KeyStoreContext<KeySlotIds>,
    ) {
        #[allow(deprecated)]
        let key_1 = ctx
            .dangerous_get_symmetric_key(key_id_1)
            .expect("key 1 should exist");
        #[allow(deprecated)]
        let key_2 = ctx
            .dangerous_get_symmetric_key(key_id_2)
            .expect("key 2 should exist");
        assert_eq!(key_1, key_2, "symmetric keys should be equal");
    }

    fn empty_common_unlock_input() -> ReencryptCommonUnlockDataInput {
        ReencryptCommonUnlockDataInput {
            trusted_devices: vec![],
            webauthn_credentials: vec![],
            trusted_organization_keys: vec![],
            trusted_emergency_access_keys: vec![],
        }
    }

    fn request_model_to_token(request: V2UpgradeTokenRequestModel) -> V2UpgradeToken {
        V2UpgradeToken {
            wrapped_user_key_1: request
                .wrapped_user_key1
                .parse()
                .expect("wrapped_user_key1 should parse"),
            wrapped_user_key_2: request
                .wrapped_user_key2
                .parse()
                .expect("wrapped_user_key2 should parse"),
        }
    }

    #[test]
    fn test_to_authentication_and_unlock_data_pbkdf2() {
        let store: KeyStore<KeySlotIds> = KeyStore::default();
        let mut ctx = store.context_mut();

        let kdf = create_test_kdf_pbkdf2();
        let salt = "[email protected]";
        let password = "test_password";

        let user_key_id = ctx.generate_symmetric_key();
        let unlock_data = MasterPasswordUnlockData::derive(password, &kdf, salt, user_key_id, &ctx)
            .expect("derive should succeed");
        let auth_data = MasterPasswordAuthenticationData::derive(password, &kdf, salt)
            .expect("derive should succeed");

        let result = to_authentication_and_unlock_data(unlock_data, auth_data, None);
        assert!(result.is_ok());

        let model = result.expect("should be ok");
        assert_eq!(model.kdf_type, KdfType::PBKDF2_SHA256);
        assert_eq!(model.kdf_iterations, 600000);
        assert!(model.kdf_memory.is_none());
        assert!(model.kdf_parallelism.is_none());
        assert_eq!(model.email, Some(salt.to_string()));
        assert!(model.master_key_authentication_hash.is_some());
        assert!(model.master_key_encrypted_user_key.is_some());
        assert!(model.master_password_hint.is_none());

        // Verify the unlock data can decrypt the user key
        let master_password_unlock_data = MasterPasswordUnlockData {
            master_key_wrapped_user_key: model
                .master_key_encrypted_user_key
                .expect("should be present")
                .parse()
                .expect("should parse"),
            kdf: kdf.clone(),
            salt: salt.to_string(),
        };
        let decrypted_user_key = master_password_unlock_data
            .unwrap_to_context(password, &mut ctx)
            .expect("unwrap should succeed");
        assert_symmetric_keys_equal(user_key_id, decrypted_user_key, &mut ctx);
    }

    #[test]
    fn test_to_authentication_and_unlock_data_argon2id() {
        let store: KeyStore<KeySlotIds> = KeyStore::default();
        let mut ctx = store.context_mut();

        let kdf = create_test_kdf_argon2id();
        let salt = "[email protected]";
        let password = "test_password";

        let user_key_id = ctx.generate_symmetric_key();
        let unlock_data = MasterPasswordUnlockData::derive(password, &kdf, salt, user_key_id, &ctx)
            .expect("derive should succeed");
        let auth_data = MasterPasswordAuthenticationData::derive(password, &kdf, salt)
            .expect("derive should succeed");

        let result = to_authentication_and_unlock_data(unlock_data, auth_data, None);
        assert!(result.is_ok());

        let model = result.expect("should be ok");
        assert_eq!(model.kdf_type, KdfType::Argon2id);
        assert_eq!(model.kdf_iterations, 3);
        assert_eq!(model.kdf_memory, Some(64));
        assert_eq!(model.kdf_parallelism, Some(4));
        assert_eq!(model.email, Some(salt.to_string()));
        assert!(model.master_key_authentication_hash.is_some());
        assert!(model.master_key_encrypted_user_key.is_some());

        // Verify the unlock data can decrypt the user key
        let master_password_unlock_data = MasterPasswordUnlockData {
            master_key_wrapped_user_key: model
                .master_key_encrypted_user_key
                .expect("should be present")
                .parse()
                .expect("should parse"),
            kdf: kdf.clone(),
            salt: salt.to_string(),
        };
        let decrypted_user_key = master_password_unlock_data
            .unwrap_to_context(password, &mut ctx)
            .expect("unwrap should succeed");
        assert_symmetric_keys_equal(user_key_id, decrypted_user_key, &mut ctx);
    }

    #[test]
    fn test_reencrypt_unlock_device_key_data() {
        let store: KeyStore<KeySlotIds> = KeyStore::default();
        let mut ctx = store.context_mut();

        let current_user_key_id = ctx.generate_symmetric_key();
        let new_user_key_id = ctx.generate_symmetric_key();

        let (device_keyset, device_private_key) =
            PartialRotateableKeyset::make_test_keyset(current_user_key_id, &mut ctx);

        let result = reencrypt_common_unlock_data(
            ReencryptCommonUnlockDataInput {
                trusted_devices: vec![device_keyset],
                webauthn_credentials: vec![],
                trusted_organization_keys: vec![],
                trusted_emergency_access_keys: vec![],
            },
            current_user_key_id,
            new_user_key_id,
            UpgradeTokenAction::CreateIfNeeded,
            &mut ctx,
        );

        let unlock_data = result.expect("should be ok");

        let device_unlock = unlock_data
            .device_key_unlock_data
            .as_ref()
            .expect("should be present")
            .first()
            .expect("should have at least one");
        let decrypted_user_key = device_unlock
            .encrypted_user_key
            .parse::<UnsignedSharedKey>()
            .expect("should parse")
            .decapsulate(device_private_key, &mut ctx)
            .expect("unwrap should succeed");
        assert_symmetric_keys_equal(new_user_key_id, decrypted_user_key, &mut ctx);
    }

    #[test]
    fn test_reencrypt_unlock_webauthn_prf_credential_data() {
        let store: KeyStore<KeySlotIds> = KeyStore::default();
        let mut ctx = store.context_mut();

        let current_user_key_id = ctx.generate_symmetric_key();
        let new_user_key_id = ctx.generate_symmetric_key();

        let (credential_keyset, credential_private_key) =
            PartialRotateableKeyset::make_test_keyset(current_user_key_id, &mut ctx);

        let result = reencrypt_common_unlock_data(
            ReencryptCommonUnlockDataInput {
                trusted_devices: vec![],
                webauthn_credentials: vec![credential_keyset],
                trusted_organization_keys: vec![],
                trusted_emergency_access_keys: vec![],
            },
            current_user_key_id,
            new_user_key_id,
            UpgradeTokenAction::CreateIfNeeded,
            &mut ctx,
        );

        let unlock_data = result.expect("should be ok");

        // Ensure it decrypts to the correct key after rotation
        let credential_unlock = unlock_data
            .passkey_unlock_data
            .as_ref()
            .expect("should be present")
            .first()
            .expect("should have at least one");
        let decrypted_user_key = credential_unlock
            .encrypted_user_key
            .parse::<UnsignedSharedKey>()
            .expect("should parse")
            .decapsulate(credential_private_key, &mut ctx)
            .expect("unwrap should succeed");
        assert_symmetric_keys_equal(new_user_key_id, decrypted_user_key, &mut ctx);
    }

    #[test]
    fn test_reencrypt_unlock_emergency_access_data() {
        let store: KeyStore<KeySlotIds> = KeyStore::default();
        let mut ctx = store.context_mut();

        let current_user_key_id = ctx.generate_symmetric_key();
        let new_user_key_id = ctx.generate_symmetric_key();

        let organization_private_key =
            ctx.make_private_key(PublicKeyEncryptionAlgorithm::RsaOaepSha1);
        let emergency_access = V1EmergencyAccessMembership {
            id: Uuid::new_v4(),
            grantee_id: Uuid::new_v4(),
            name: "Test User".to_string(),
            public_key: ctx
                .get_public_key(organization_private_key)
                .expect("key exists"),
        };

        let result = reencrypt_common_unlock_data(
            ReencryptCommonUnlockDataInput {
                trusted_devices: vec![],
                webauthn_credentials: vec![],
                trusted_organization_keys: vec![],
                trusted_emergency_access_keys: vec![emergency_access],
            },
            current_user_key_id,
            new_user_key_id,
            UpgradeTokenAction::CreateIfNeeded,
            &mut ctx,
        );

        let unlock_data = result.expect("should be ok");

        // Ensure it decrypts to the correct key after rotation
        let emergency_access_unlock = unlock_data
            .emergency_access_unlock_data
            .as_ref()
            .expect("should be present")
            .first()
            .expect("should have at least one");
        let decrypted_user_key = emergency_access_unlock
            .key_encrypted
            .as_ref()
            .map(|k| k.parse::<UnsignedSharedKey>())
            .expect("should be present")
            .expect("should parse")
            .decapsulate(organization_private_key, &mut ctx)
            .expect("unwrap should succeed");
        assert_symmetric_keys_equal(new_user_key_id, decrypted_user_key, &mut ctx);
    }

    #[test]
    fn test_reencrypt_unlock_organization_membership_data() {
        let store: KeyStore<KeySlotIds> = KeyStore::default();
        let mut ctx = store.context_mut();

        let current_user_key_id = ctx.generate_symmetric_key();
        let new_user_key_id = ctx.generate_symmetric_key();

        let org_key = ctx.make_private_key(PublicKeyEncryptionAlgorithm::RsaOaepSha1);
        let org_membership = V1OrganizationMembership {
            organization_id: Uuid::new_v4(),
            name: "Test Org".to_string(),
            public_key: ctx.get_public_key(org_key).expect("key exists"),
        };

        let result = reencrypt_common_unlock_data(
            ReencryptCommonUnlockDataInput {
                trusted_devices: vec![],
                webauthn_credentials: vec![],
                trusted_organization_keys: vec![org_membership],
                trusted_emergency_access_keys: vec![],
            },
            current_user_key_id,
            new_user_key_id,
            UpgradeTokenAction::CreateIfNeeded,
            &mut ctx,
        );

        let unlock_data = result.expect("should be ok");

        let org_membership_unlock = unlock_data
            .organization_account_recovery_unlock_data
            .as_ref()
            .expect("should be present")
            .first()
            .expect("should have at least one");
        let decrypted_user_key = org_membership_unlock
            .reset_password_key
            .as_ref()
            .map(|k| k.parse::<UnsignedSharedKey>())
            .expect("should be present")
            .expect("should parse")
            .decapsulate(org_key, &mut ctx)
            .expect("unwrap should succeed");
        assert_symmetric_keys_equal(new_user_key_id, decrypted_user_key, &mut ctx);
    }

    #[test]
    fn test_reencrypt_common_unlock_data_v1_to_v2_creates_upgrade_token() {
        let store: KeyStore<KeySlotIds> = KeyStore::default();
        let mut ctx = store.context_mut();

        let current_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
        let new_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::XAes256Gcm);

        let result = reencrypt_common_unlock_data(
            empty_common_unlock_input(),
            current_user_key_id,
            new_user_key_id,
            UpgradeTokenAction::CreateIfNeeded,
            &mut ctx,
        );

        let unlock_data = result.expect("should be ok");
        let token_request = *unlock_data
            .v2_upgrade_token
            .expect("v2_upgrade_token should be populated for V1 -> V2 rotation");
        let token = request_model_to_token(token_request);

        let unwrapped_v2_id = token
            .unwrap_v2(current_user_key_id, &mut ctx)
            .expect("unwrap_v2 should succeed");
        assert_symmetric_keys_equal(new_user_key_id, unwrapped_v2_id, &mut ctx);

        let unwrapped_v1_id = token
            .unwrap_v1(new_user_key_id, &mut ctx)
            .expect("unwrap_v1 should succeed");
        assert_symmetric_keys_equal(current_user_key_id, unwrapped_v1_id, &mut ctx);
    }

    #[test]
    fn test_reencrypt_common_unlock_data_v1_to_v2_upgrade_token_action_skip_returns_none() {
        let store: KeyStore<KeySlotIds> = KeyStore::default();
        let mut ctx = store.context_mut();

        let current_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
        let new_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::XAes256Gcm);

        let result = reencrypt_common_unlock_data(
            empty_common_unlock_input(),
            current_user_key_id,
            new_user_key_id,
            UpgradeTokenAction::Skip,
            &mut ctx,
        );

        let unlock_data = result.expect("should be ok");
        assert!(
            unlock_data.v2_upgrade_token.is_none(),
            "UpgradeTokenAction::Skip skips the creation of the upgrade token"
        );
    }

    #[test]
    fn test_reencrypt_common_unlock_data_v2_to_v2_upgrade_token_action_create_if_needed_returns_none()
     {
        let store: KeyStore<KeySlotIds> = KeyStore::default();
        let mut ctx = store.context_mut();

        let current_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::XAes256Gcm);
        let new_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::XAes256Gcm);

        let result = reencrypt_common_unlock_data(
            empty_common_unlock_input(),
            current_user_key_id,
            new_user_key_id,
            UpgradeTokenAction::CreateIfNeeded,
            &mut ctx,
        );

        let unlock_data = result.expect("should be ok");
        assert!(
            unlock_data.v2_upgrade_token.is_none(),
            "UpgradeTokenAction::CreateIfNeeded should not create a v2_upgrade_token for V2 -> V2 rotation"
        );
    }

    fn make_valid_public_key_b64() -> String {
        let store: KeyStore<KeySlotIds> = KeyStore::default();
        let mut ctx = store.context_mut();
        let key_id = ctx.make_private_key(PublicKeyEncryptionAlgorithm::RsaOaepSha1);
        ctx.get_public_key(key_id)
            .expect("public key exists")
            .to_string()
    }

    #[test]
    fn test_v1_emergency_access_membership_try_from_uses_name_when_present() {
        let id = Uuid::new_v4();
        let grantee_id = Uuid::new_v4();
        let model = EmergencyAccessKeyDataResponseModel {
            id: Some(id),
            grantee_id: Some(grantee_id),
            grantee_name: Some("Alice".to_string()),
            grantee_email: Some("[email protected]".to_string()),
            public_key: Some(make_valid_public_key_b64()),
            ..Default::default()
        };

        let membership = V1EmergencyAccessMembership::try_from(model).expect("should be ok");
        assert_eq!(membership.id, id);
        assert_eq!(membership.grantee_id, grantee_id);
        assert_eq!(membership.name, "Alice");
    }

    #[test]
    fn test_v1_emergency_access_membership_try_from_falls_back_to_email_when_name_missing() {
        let model = EmergencyAccessKeyDataResponseModel {
            id: Some(Uuid::new_v4()),
            grantee_id: Some(Uuid::new_v4()),
            grantee_name: None,
            grantee_email: Some("[email protected]".to_string()),
            public_key: Some(make_valid_public_key_b64()),
            ..Default::default()
        };

        let membership = V1EmergencyAccessMembership::try_from(model).expect("should be ok");
        assert_eq!(membership.name, "[email protected]");
    }

    #[test]
    fn test_v1_emergency_access_membership_try_from_falls_back_to_unknown_when_name_and_email_missing()
     {
        let model = EmergencyAccessKeyDataResponseModel {
            id: Some(Uuid::new_v4()),
            grantee_id: Some(Uuid::new_v4()),
            grantee_name: None,
            grantee_email: None,
            public_key: Some(make_valid_public_key_b64()),
            ..Default::default()
        };

        let membership = V1EmergencyAccessMembership::try_from(model).expect("should be ok");
        assert_eq!(membership.name, "Unknown");
    }

    #[test]
    fn test_v1_emergency_access_membership_try_from_missing_id_returns_error() {
        let model = EmergencyAccessKeyDataResponseModel {
            id: None,
            grantee_id: Some(Uuid::new_v4()),
            grantee_name: Some("Alice".to_string()),
            grantee_email: None,
            public_key: Some(make_valid_public_key_b64()),
            ..Default::default()
        };

        let Err(err) = V1EmergencyAccessMembership::try_from(model) else {
            panic!("expected error")
        };
        assert!(matches!(err, KeyRotationDataParseError::MissingField(_)));
    }

    #[test]
    fn test_v1_emergency_access_membership_try_from_missing_grantee_id_returns_error() {
        let model = EmergencyAccessKeyDataResponseModel {
            id: Some(Uuid::new_v4()),
            grantee_id: None,
            grantee_name: Some("Alice".to_string()),
            grantee_email: None,
            public_key: Some(make_valid_public_key_b64()),
            ..Default::default()
        };

        let Err(err) = V1EmergencyAccessMembership::try_from(model) else {
            panic!("expected error")
        };
        assert!(matches!(err, KeyRotationDataParseError::MissingField(_)));
    }

    #[test]
    fn test_v1_emergency_access_membership_try_from_missing_public_key_returns_error() {
        let model = EmergencyAccessKeyDataResponseModel {
            id: Some(Uuid::new_v4()),
            grantee_id: Some(Uuid::new_v4()),
            grantee_name: Some("Alice".to_string()),
            grantee_email: None,
            public_key: None,
            ..Default::default()
        };

        let Err(err) = V1EmergencyAccessMembership::try_from(model) else {
            panic!("expected error")
        };
        assert!(matches!(err, KeyRotationDataParseError::MissingField(_)));
    }

    #[test]
    fn test_v1_emergency_access_membership_try_from_invalid_b64_public_key_returns_error() {
        let model = EmergencyAccessKeyDataResponseModel {
            id: Some(Uuid::new_v4()),
            grantee_id: Some(Uuid::new_v4()),
            grantee_name: Some("Alice".to_string()),
            grantee_email: None,
            public_key: Some("not valid base64 !!!".to_string()),
            ..Default::default()
        };

        let Err(err) = V1EmergencyAccessMembership::try_from(model) else {
            panic!("expected error")
        };
        assert!(matches!(err, KeyRotationDataParseError::B64(_)));
    }

    #[test]
    fn test_v1_emergency_access_membership_try_from_invalid_spki_public_key_returns_error() {
        let model = EmergencyAccessKeyDataResponseModel {
            id: Some(Uuid::new_v4()),
            grantee_id: Some(Uuid::new_v4()),
            grantee_name: Some("Alice".to_string()),
            grantee_email: None,
            public_key: Some(B64::from(b"not-a-real-public-key".as_slice()).to_string()),
            ..Default::default()
        };

        let Err(err) = V1EmergencyAccessMembership::try_from(model) else {
            panic!("expected error")
        };
        assert!(matches!(err, KeyRotationDataParseError::Crypto(_)));
    }

    #[test]
    fn test_reencrypt_master_password_change_unlock_data_never_returns_upgrade_token() {
        let store: KeyStore<KeySlotIds> = KeyStore::default();
        let mut ctx = store.context_mut();

        let current_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
        let new_user_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::XAes256Gcm);

        let input = ReencryptMasterPasswordChangeAndUnlockInput {
            password: "test_password".to_string(),
            hint: None,
            kdf: create_test_kdf_pbkdf2(),
            salt: "[email protected]".to_string(),
            common_unlock_data: empty_common_unlock_input(),
        };

        let unlock_data = reencrypt_master_password_change_unlock_data(
            input,
            current_user_key_id,
            new_user_key_id,
            &mut ctx,
        )
        .expect("should be ok");

        assert!(
            unlock_data.v2_upgrade_token.is_none(),
            "master password change rotation must never include a v2 upgrade token"
        );
    }
}