bitwarden_crypto/store/
context.rs

1use std::{
2    cell::Cell,
3    sync::{RwLockReadGuard, RwLockWriteGuard},
4};
5
6use serde::Serialize;
7use zeroize::Zeroizing;
8
9use super::KeyStoreInner;
10use crate::{
11    derive_shareable_key, error::UnsupportedOperation, signing, store::backend::StoreBackend,
12    AsymmetricCryptoKey, BitwardenLegacyKeyBytes, ContentFormat, CryptoError, EncString, KeyId,
13    KeyIds, PublicKeyEncryptionAlgorithm, Result, RotatedUserKeys, Signature, SignatureAlgorithm,
14    SignedObject, SignedPublicKey, SignedPublicKeyMessage, SigningKey, SymmetricCryptoKey,
15    UnsignedSharedKey,
16};
17
18/// The context of a crypto operation using [super::KeyStore]
19///
20/// This will usually be accessed from an implementation of [crate::Decryptable] or
21/// [crate::CompositeEncryptable], [crate::PrimitiveEncryptable],
22/// but can also be obtained
23/// through [super::KeyStore::context]
24///
25/// This context contains access to the user keys stored in the [super::KeyStore] (sometimes
26/// referred to as `global keys`) and it also contains it's own individual secure backend for key
27/// storage. Keys stored in this individual backend are usually referred to as `local keys`, they
28/// will be cleared when this context goes out of scope and is dropped and they do not affect either
29/// the global [super::KeyStore] or other instances of contexts.
30///
31/// This context-local storage is recommended for ephemeral and temporary keys that are decrypted
32/// during the course of a decrypt/encrypt operation, but won't be used after the operation itself
33/// is complete.
34///
35/// ```rust
36/// # use bitwarden_crypto::*;
37/// # key_ids! {
38/// #     #[symmetric]
39/// #     pub enum SymmKeyId {
40/// #         User,
41/// #         Local(&'static str),
42/// #     }
43/// #     #[asymmetric]
44/// #     pub enum AsymmKeyId {
45/// #         UserPrivate,
46/// #     }
47/// #     #[signing]
48/// #     pub enum SigningKeyId {
49/// #         UserSigning,
50/// #     }
51/// #     pub Ids => SymmKeyId, AsymmKeyId, SigningKeyId;
52/// # }
53/// struct Data {
54///     key: EncString,
55///     name: String,
56/// }
57/// # impl IdentifyKey<SymmKeyId> for Data {
58/// #    fn key_identifier(&self) -> SymmKeyId {
59/// #        SymmKeyId::User
60/// #    }
61/// # }
62///
63/// const LOCAL_KEY: SymmKeyId = SymmKeyId::Local("local_key_id");
64///
65/// impl CompositeEncryptable<Ids, SymmKeyId, EncString> for Data {
66///     fn encrypt_composite(&self, ctx: &mut KeyStoreContext<Ids>, key: SymmKeyId) -> Result<EncString, CryptoError> {
67///         let local_key_id = ctx.unwrap_symmetric_key(key, LOCAL_KEY, &self.key)?;
68///         self.name.encrypt(ctx, local_key_id)
69///     }
70/// }
71/// ```
72#[must_use]
73pub struct KeyStoreContext<'a, Ids: KeyIds> {
74    pub(super) global_keys: GlobalKeys<'a, Ids>,
75
76    pub(super) local_symmetric_keys: Box<dyn StoreBackend<Ids::Symmetric>>,
77    pub(super) local_asymmetric_keys: Box<dyn StoreBackend<Ids::Asymmetric>>,
78    pub(super) local_signing_keys: Box<dyn StoreBackend<Ids::Signing>>,
79
80    // Make sure the context is !Send & !Sync
81    pub(super) _phantom: std::marker::PhantomData<(Cell<()>, RwLockReadGuard<'static, ()>)>,
82}
83
84/// A KeyStoreContext is usually limited to a read only access to the global keys,
85/// which allows us to have multiple read only contexts at the same time and do multitheaded
86/// encryption/decryption. We also have the option to create a read/write context, which allows us
87/// to modify the global keys, but only allows one context at a time. This is controlled by a
88/// [std::sync::RwLock] on the global keys, and this struct stores both types of guards.
89pub(crate) enum GlobalKeys<'a, Ids: KeyIds> {
90    ReadOnly(RwLockReadGuard<'a, KeyStoreInner<Ids>>),
91    ReadWrite(RwLockWriteGuard<'a, KeyStoreInner<Ids>>),
92}
93
94impl<Ids: KeyIds> GlobalKeys<'_, Ids> {
95    pub fn get(&self) -> &KeyStoreInner<Ids> {
96        match self {
97            GlobalKeys::ReadOnly(keys) => keys,
98            GlobalKeys::ReadWrite(keys) => keys,
99        }
100    }
101
102    pub fn get_mut(&mut self) -> Result<&mut KeyStoreInner<Ids>> {
103        match self {
104            GlobalKeys::ReadOnly(_) => Err(CryptoError::ReadOnlyKeyStore),
105            GlobalKeys::ReadWrite(keys) => Ok(keys),
106        }
107    }
108}
109
110impl<Ids: KeyIds> KeyStoreContext<'_, Ids> {
111    /// Clears all the local keys stored in this context
112    /// This will not affect the global keys even if this context has write access.
113    /// To clear the global keys, you need to use [super::KeyStore::clear] instead.
114    pub fn clear_local(&mut self) {
115        self.local_symmetric_keys.clear();
116        self.local_asymmetric_keys.clear();
117        self.local_signing_keys.clear();
118    }
119
120    /// Remove all symmetric keys from the context for which the predicate returns false
121    /// This will also remove the keys from the global store if this context has write access
122    pub fn retain_symmetric_keys(&mut self, f: fn(Ids::Symmetric) -> bool) {
123        if let Ok(keys) = self.global_keys.get_mut() {
124            keys.symmetric_keys.retain(f);
125        }
126        self.local_symmetric_keys.retain(f);
127    }
128
129    /// Remove all asymmetric keys from the context for which the predicate returns false
130    /// This will also remove the keys from the global store if this context has write access
131    pub fn retain_asymmetric_keys(&mut self, f: fn(Ids::Asymmetric) -> bool) {
132        if let Ok(keys) = self.global_keys.get_mut() {
133            keys.asymmetric_keys.retain(f);
134        }
135        self.local_asymmetric_keys.retain(f);
136    }
137
138    // TODO: All these encrypt x key with x key look like they need to be made generic,
139    // but I haven't found the best way to do that yet.
140
141    /// Decrypt a symmetric key into the context by using an already existing symmetric key
142    ///
143    /// # Arguments
144    ///
145    /// * `wrapping_key` - The key id used to decrypt the `wrapped_key`. It must already exist in
146    ///   the context
147    /// * `new_key_id` - The key id where the decrypted key will be stored. If it already exists, it
148    ///   will be overwritten
149    /// * `wrapped_key` - The key to decrypt
150    pub fn unwrap_symmetric_key(
151        &mut self,
152        wrapping_key: Ids::Symmetric,
153        new_key_id: Ids::Symmetric,
154        wrapped_key: &EncString,
155    ) -> Result<Ids::Symmetric> {
156        let wrapping_key = self.get_symmetric_key(wrapping_key)?;
157
158        let key = match (wrapped_key, wrapping_key) {
159            (EncString::Aes256Cbc_B64 { iv, data }, SymmetricCryptoKey::Aes256CbcKey(key)) => {
160                SymmetricCryptoKey::try_from(&BitwardenLegacyKeyBytes::from(
161                    crate::aes::decrypt_aes256(iv, data.clone(), &key.enc_key)?,
162                ))?
163            }
164            (
165                EncString::Aes256Cbc_HmacSha256_B64 { iv, mac, data },
166                SymmetricCryptoKey::Aes256CbcHmacKey(key),
167            ) => SymmetricCryptoKey::try_from(&BitwardenLegacyKeyBytes::from(
168                crate::aes::decrypt_aes256_hmac(iv, mac, data.clone(), &key.mac_key, &key.enc_key)?,
169            ))?,
170            (
171                EncString::Cose_Encrypt0_B64 { data },
172                SymmetricCryptoKey::XChaCha20Poly1305Key(key),
173            ) => {
174                let (content_bytes, content_format) =
175                    crate::cose::decrypt_xchacha20_poly1305(data, key)?;
176                match content_format {
177                    ContentFormat::BitwardenLegacyKey => {
178                        SymmetricCryptoKey::try_from(&BitwardenLegacyKeyBytes::from(content_bytes))?
179                    }
180                    ContentFormat::CoseKey => SymmetricCryptoKey::try_from_cose(&content_bytes)?,
181                    _ => return Err(CryptoError::InvalidKey),
182                }
183            }
184            _ => return Err(CryptoError::InvalidKey),
185        };
186
187        #[allow(deprecated)]
188        self.set_symmetric_key(new_key_id, key)?;
189
190        // Returning the new key identifier for convenience
191        Ok(new_key_id)
192    }
193
194    /// Encrypt and return a symmetric key from the context by using an already existing symmetric
195    /// key
196    ///
197    /// # Arguments
198    ///
199    /// * `wrapping_key` - The key id used to wrap (encrypt) the `key_to_wrap`. It must already
200    ///   exist in the context
201    /// * `key_to_wrap` - The key id to wrap. It must already exist in the context
202    pub fn wrap_symmetric_key(
203        &self,
204        wrapping_key: Ids::Symmetric,
205        key_to_wrap: Ids::Symmetric,
206    ) -> Result<EncString> {
207        use SymmetricCryptoKey::*;
208
209        let wrapping_key_instance = self.get_symmetric_key(wrapping_key)?;
210        let key_to_wrap_instance = self.get_symmetric_key(key_to_wrap)?;
211        // `Aes256CbcHmacKey` can wrap keys by encrypting their byte serialization obtained using
212        // `SymmetricCryptoKey::to_encoded()`. `XChaCha20Poly1305Key` need to specify the
213        // content format to be either octet stream, in case the wrapped key is a Aes256CbcHmacKey
214        // or `Aes256CbcKey`, or by specifying the content format to be CoseKey, in case the
215        // wrapped key is a `XChaCha20Poly1305Key`.
216        match (wrapping_key_instance, key_to_wrap_instance) {
217            (
218                Aes256CbcHmacKey(_),
219                Aes256CbcHmacKey(_) | Aes256CbcKey(_) | XChaCha20Poly1305Key(_),
220            ) => self.encrypt_data_with_symmetric_key(
221                wrapping_key,
222                key_to_wrap_instance
223                    .to_encoded()
224                    .as_ref()
225                    .to_vec()
226                    .as_slice(),
227                ContentFormat::BitwardenLegacyKey,
228            ),
229            (XChaCha20Poly1305Key(_), _) => {
230                let encoded = key_to_wrap_instance.to_encoded_raw();
231                let content_format = encoded.content_format();
232                self.encrypt_data_with_symmetric_key(
233                    wrapping_key,
234                    Into::<Vec<u8>>::into(encoded).as_slice(),
235                    content_format,
236                )
237            }
238            _ => Err(CryptoError::OperationNotSupported(
239                UnsupportedOperation::EncryptionNotImplementedForKey,
240            )),
241        }
242    }
243
244    /// Decapsulate a symmetric key into the context by using an already existing asymmetric key
245    ///
246    /// # Arguments
247    ///
248    /// * `decapsulation_key` - The key id used to decrypt the `encrypted_key`. It must already
249    ///   exist in the context
250    /// * `new_key_id` - The key id where the decrypted key will be stored. If it already exists, it
251    ///   will be overwritten
252    /// * `encapsulated_shared_key` - The symmetric key to decrypt
253    pub fn decapsulate_key_unsigned(
254        &mut self,
255        decapsulation_key: Ids::Asymmetric,
256        new_key_id: Ids::Symmetric,
257        encapsulated_shared_key: &UnsignedSharedKey,
258    ) -> Result<Ids::Symmetric> {
259        let decapsulation_key = self.get_asymmetric_key(decapsulation_key)?;
260        let decapsulated_key =
261            encapsulated_shared_key.decapsulate_key_unsigned(decapsulation_key)?;
262
263        #[allow(deprecated)]
264        self.set_symmetric_key(new_key_id, decapsulated_key)?;
265
266        // Returning the new key identifier for convenience
267        Ok(new_key_id)
268    }
269
270    /// Encapsulate and return a symmetric key from the context by using an already existing
271    /// asymmetric key
272    ///
273    /// # Arguments
274    ///
275    /// * `encapsulation_key` - The key id used to encrypt the `encapsulated_key`. It must already
276    ///   exist in the context
277    /// * `shared_key` - The key id to encrypt. It must already exist in the context
278    pub fn encapsulate_key_unsigned(
279        &self,
280        encapsulation_key: Ids::Asymmetric,
281        shared_key: Ids::Symmetric,
282    ) -> Result<UnsignedSharedKey> {
283        UnsignedSharedKey::encapsulate_key_unsigned(
284            self.get_symmetric_key(shared_key)?,
285            &self.get_asymmetric_key(encapsulation_key)?.to_public_key(),
286        )
287    }
288
289    /// Returns `true` if the context has a symmetric key with the given identifier
290    pub fn has_symmetric_key(&self, key_id: Ids::Symmetric) -> bool {
291        self.get_symmetric_key(key_id).is_ok()
292    }
293
294    /// Returns `true` if the context has an asymmetric key with the given identifier
295    pub fn has_asymmetric_key(&self, key_id: Ids::Asymmetric) -> bool {
296        self.get_asymmetric_key(key_id).is_ok()
297    }
298
299    /// Returns `true` if the context has a signing key with the given identifier
300    pub fn has_signing_key(&self, key_id: Ids::Signing) -> bool {
301        self.get_signing_key(key_id).is_ok()
302    }
303
304    /// Generate a new random symmetric key and store it in the context
305    pub fn generate_symmetric_key(&mut self, key_id: Ids::Symmetric) -> Result<Ids::Symmetric> {
306        let key = SymmetricCryptoKey::make_aes256_cbc_hmac_key();
307        #[allow(deprecated)]
308        self.set_symmetric_key(key_id, key)?;
309        Ok(key_id)
310    }
311
312    /// Makes a new asymmetric encryption key using the current default algorithm, and stores it in
313    /// the context
314    pub fn make_asymmetric_key(&mut self, key_id: Ids::Asymmetric) -> Result<Ids::Asymmetric> {
315        let key = AsymmetricCryptoKey::make(PublicKeyEncryptionAlgorithm::RsaOaepSha1);
316        #[allow(deprecated)]
317        self.set_asymmetric_key(key_id, key)?;
318        Ok(key_id)
319    }
320
321    /// Generate a new signature key using the current default algorithm, and store it in the
322    /// context
323    pub fn make_signing_key(&mut self, key_id: Ids::Signing) -> Result<Ids::Signing> {
324        let key = SigningKey::make(SignatureAlgorithm::default_algorithm());
325        #[allow(deprecated)]
326        self.set_signing_key(key_id, key)?;
327        Ok(key_id)
328    }
329
330    /// Derive a shareable key using hkdf from secret and name and store it in the context.
331    ///
332    /// A specialized variant of this function was called `CryptoService.makeSendKey` in the
333    /// Bitwarden `clients` repository.
334    pub fn derive_shareable_key(
335        &mut self,
336        key_id: Ids::Symmetric,
337        secret: Zeroizing<[u8; 16]>,
338        name: &str,
339        info: Option<&str>,
340    ) -> Result<Ids::Symmetric> {
341        #[allow(deprecated)]
342        self.set_symmetric_key(
343            key_id,
344            SymmetricCryptoKey::Aes256CbcHmacKey(derive_shareable_key(secret, name, info)),
345        )?;
346        Ok(key_id)
347    }
348
349    #[deprecated(note = "This function should ideally never be used outside this crate")]
350    #[allow(missing_docs)]
351    pub fn dangerous_get_symmetric_key(
352        &self,
353        key_id: Ids::Symmetric,
354    ) -> Result<&SymmetricCryptoKey> {
355        self.get_symmetric_key(key_id)
356    }
357
358    #[deprecated(note = "This function should ideally never be used outside this crate")]
359    #[allow(missing_docs)]
360    pub fn dangerous_get_asymmetric_key(
361        &self,
362        key_id: Ids::Asymmetric,
363    ) -> Result<&AsymmetricCryptoKey> {
364        self.get_asymmetric_key(key_id)
365    }
366
367    /// Makes a signed public key from an asymmetric private key and signing key stored in context.
368    /// Signing a public key asserts ownership, and makes the claim to other users that if they want
369    /// to share with you, they can use this public key.
370    pub fn make_signed_public_key(
371        &self,
372        private_key_id: Ids::Asymmetric,
373        signing_key_id: Ids::Signing,
374    ) -> Result<SignedPublicKey> {
375        let public_key = self.get_asymmetric_key(private_key_id)?.to_public_key();
376        let signing_key = self.get_signing_key(signing_key_id)?;
377        let signed_public_key =
378            SignedPublicKeyMessage::from_public_key(&public_key)?.sign(signing_key)?;
379        Ok(signed_public_key)
380    }
381
382    fn get_symmetric_key(&self, key_id: Ids::Symmetric) -> Result<&SymmetricCryptoKey> {
383        if key_id.is_local() {
384            self.local_symmetric_keys.get(key_id)
385        } else {
386            self.global_keys.get().symmetric_keys.get(key_id)
387        }
388        .ok_or_else(|| crate::CryptoError::MissingKeyId(format!("{key_id:?}")))
389    }
390
391    pub(super) fn get_asymmetric_key(
392        &self,
393        key_id: Ids::Asymmetric,
394    ) -> Result<&AsymmetricCryptoKey> {
395        if key_id.is_local() {
396            self.local_asymmetric_keys.get(key_id)
397        } else {
398            self.global_keys.get().asymmetric_keys.get(key_id)
399        }
400        .ok_or_else(|| crate::CryptoError::MissingKeyId(format!("{key_id:?}")))
401    }
402
403    pub(super) fn get_signing_key(&self, key_id: Ids::Signing) -> Result<&SigningKey> {
404        if key_id.is_local() {
405            self.local_signing_keys.get(key_id)
406        } else {
407            self.global_keys.get().signing_keys.get(key_id)
408        }
409        .ok_or_else(|| crate::CryptoError::MissingKeyId(format!("{key_id:?}")))
410    }
411
412    #[deprecated(note = "This function should ideally never be used outside this crate")]
413    #[allow(missing_docs)]
414    pub fn set_symmetric_key(
415        &mut self,
416        key_id: Ids::Symmetric,
417        key: SymmetricCryptoKey,
418    ) -> Result<()> {
419        if key_id.is_local() {
420            self.local_symmetric_keys.upsert(key_id, key);
421        } else {
422            self.global_keys
423                .get_mut()?
424                .symmetric_keys
425                .upsert(key_id, key);
426        }
427        Ok(())
428    }
429
430    #[deprecated(note = "This function should ideally never be used outside this crate")]
431    #[allow(missing_docs)]
432    pub fn set_asymmetric_key(
433        &mut self,
434        key_id: Ids::Asymmetric,
435        key: AsymmetricCryptoKey,
436    ) -> Result<()> {
437        if key_id.is_local() {
438            self.local_asymmetric_keys.upsert(key_id, key);
439        } else {
440            self.global_keys
441                .get_mut()?
442                .asymmetric_keys
443                .upsert(key_id, key);
444        }
445        Ok(())
446    }
447
448    /// Sets a signing key in the context
449    #[deprecated(note = "This function should ideally never be used outside this crate")]
450    pub fn set_signing_key(&mut self, key_id: Ids::Signing, key: SigningKey) -> Result<()> {
451        if key_id.is_local() {
452            self.local_signing_keys.upsert(key_id, key);
453        } else {
454            self.global_keys.get_mut()?.signing_keys.upsert(key_id, key);
455        }
456        Ok(())
457    }
458
459    pub(crate) fn decrypt_data_with_symmetric_key(
460        &self,
461        key: Ids::Symmetric,
462        data: &EncString,
463    ) -> Result<Vec<u8>> {
464        let key = self.get_symmetric_key(key)?;
465
466        match (data, key) {
467            (EncString::Aes256Cbc_B64 { iv, data }, SymmetricCryptoKey::Aes256CbcKey(key)) => {
468                crate::aes::decrypt_aes256(iv, data.clone(), &key.enc_key)
469            }
470            (
471                EncString::Aes256Cbc_HmacSha256_B64 { iv, mac, data },
472                SymmetricCryptoKey::Aes256CbcHmacKey(key),
473            ) => crate::aes::decrypt_aes256_hmac(iv, mac, data.clone(), &key.mac_key, &key.enc_key),
474            (
475                EncString::Cose_Encrypt0_B64 { data },
476                SymmetricCryptoKey::XChaCha20Poly1305Key(key),
477            ) => {
478                let (data, _) = crate::cose::decrypt_xchacha20_poly1305(data, key)?;
479                Ok(data)
480            }
481            _ => Err(CryptoError::InvalidKey),
482        }
483    }
484
485    pub(crate) fn encrypt_data_with_symmetric_key(
486        &self,
487        key: Ids::Symmetric,
488        data: &[u8],
489        content_format: ContentFormat,
490    ) -> Result<EncString> {
491        let key = self.get_symmetric_key(key)?;
492        match key {
493            SymmetricCryptoKey::Aes256CbcKey(_) => Err(CryptoError::OperationNotSupported(
494                UnsupportedOperation::EncryptionNotImplementedForKey,
495            )),
496            SymmetricCryptoKey::Aes256CbcHmacKey(key) => EncString::encrypt_aes256_hmac(data, key),
497            SymmetricCryptoKey::XChaCha20Poly1305Key(key) => {
498                EncString::encrypt_xchacha20_poly1305(data, key, content_format)
499            }
500        }
501    }
502
503    /// Signs the given data using the specified signing key, for the given
504    /// [crate::SigningNamespace] and returns the signature and the serialized message. See
505    /// [crate::SigningKey::sign]
506    pub fn sign<Message: Serialize>(
507        &self,
508        key: Ids::Signing,
509        message: &Message,
510        namespace: &crate::SigningNamespace,
511    ) -> Result<SignedObject> {
512        self.get_signing_key(key)?.sign(message, namespace)
513    }
514
515    /// Signs the given data using the specified signing key, for the given
516    /// [crate::SigningNamespace] and returns the signature and the serialized message. See
517    /// [crate::SigningKey::sign_detached]
518    #[allow(unused)]
519    pub(crate) fn sign_detached<Message: Serialize>(
520        &self,
521        key: Ids::Signing,
522        message: &Message,
523        namespace: &crate::SigningNamespace,
524    ) -> Result<(Signature, signing::SerializedMessage)> {
525        self.get_signing_key(key)?.sign_detached(message, namespace)
526    }
527
528    /// Re-encrypts the user's keys with the provided symmetric key for a v2 user.
529    pub fn dangerous_get_v2_rotated_account_keys(
530        &self,
531        current_user_private_key_id: Ids::Asymmetric,
532        current_user_signing_key_id: Ids::Signing,
533    ) -> Result<RotatedUserKeys> {
534        crate::dangerous_get_v2_rotated_account_keys(
535            current_user_private_key_id,
536            current_user_signing_key_id,
537            self,
538        )
539    }
540}
541
542#[cfg(test)]
543#[allow(deprecated)]
544mod tests {
545    use serde::{Deserialize, Serialize};
546
547    use crate::{
548        store::{
549            tests::{Data, DataView},
550            KeyStore,
551        },
552        traits::tests::{TestAsymmKey, TestIds, TestSigningKey, TestSymmKey},
553        AsymmetricCryptoKey, AsymmetricPublicCryptoKey, CompositeEncryptable, CoseKeyBytes,
554        CoseSerializable, CryptoError, Decryptable, KeyDecryptable, Pkcs8PrivateKeyBytes,
555        PublicKeyEncryptionAlgorithm, SignatureAlgorithm, SigningKey, SigningNamespace,
556        SymmetricCryptoKey,
557    };
558
559    #[test]
560    fn test_set_signing_key() {
561        let store: KeyStore<TestIds> = KeyStore::default();
562
563        // Generate and insert a key
564        let key_a0_id = TestSigningKey::A(0);
565        let key_a0 = SigningKey::make(SignatureAlgorithm::Ed25519);
566        store
567            .context_mut()
568            .set_signing_key(key_a0_id, key_a0)
569            .unwrap();
570    }
571
572    #[test]
573    fn test_set_keys_for_encryption() {
574        let store: KeyStore<TestIds> = KeyStore::default();
575
576        // Generate and insert a key
577        let key_a0_id = TestSymmKey::A(0);
578        let key_a0 = SymmetricCryptoKey::make_aes256_cbc_hmac_key();
579
580        store
581            .context_mut()
582            .set_symmetric_key(TestSymmKey::A(0), key_a0.clone())
583            .unwrap();
584
585        assert!(store.context().has_symmetric_key(key_a0_id));
586
587        // Encrypt some data with the key
588        let data = DataView("Hello, World!".to_string(), key_a0_id);
589        let _encrypted: Data = data
590            .encrypt_composite(&mut store.context(), key_a0_id)
591            .unwrap();
592    }
593
594    #[test]
595    fn test_key_encryption() {
596        let store: KeyStore<TestIds> = KeyStore::default();
597
598        let mut ctx = store.context();
599
600        // Generate and insert a key
601        let key_1_id = TestSymmKey::C(1);
602        let key_1 = SymmetricCryptoKey::make_aes256_cbc_hmac_key();
603
604        ctx.set_symmetric_key(key_1_id, key_1.clone()).unwrap();
605
606        assert!(ctx.has_symmetric_key(key_1_id));
607
608        // Generate and insert a new key
609        let key_2_id = TestSymmKey::C(2);
610        let key_2 = SymmetricCryptoKey::make_aes256_cbc_hmac_key();
611
612        ctx.set_symmetric_key(key_2_id, key_2.clone()).unwrap();
613
614        assert!(ctx.has_symmetric_key(key_2_id));
615
616        // Encrypt the new key with the old key
617        let key_2_enc = ctx.wrap_symmetric_key(key_1_id, key_2_id).unwrap();
618
619        // Decrypt the new key with the old key in a different identifier
620        let new_key_id = TestSymmKey::C(3);
621
622        ctx.unwrap_symmetric_key(key_1_id, new_key_id, &key_2_enc)
623            .unwrap();
624
625        // Now `key_2_id` and `new_key_id` contain the same key, so we should be able to encrypt
626        // with one and decrypt with the other
627
628        let data = DataView("Hello, World!".to_string(), key_2_id);
629        let encrypted = data.encrypt_composite(&mut ctx, key_2_id).unwrap();
630
631        let decrypted1 = encrypted.decrypt(&mut ctx, key_2_id).unwrap();
632        let decrypted2 = encrypted.decrypt(&mut ctx, new_key_id).unwrap();
633
634        // Assert that the decrypted data is the same
635        assert_eq!(decrypted1.0, decrypted2.0);
636    }
637
638    #[test]
639    fn test_wrap_unwrap() {
640        let store: KeyStore<TestIds> = KeyStore::default();
641        let mut ctx = store.context_mut();
642
643        // Aes256 CBC HMAC keys
644        let key_aes_1_id = TestSymmKey::A(1);
645        let key_aes_1 = SymmetricCryptoKey::make_aes256_cbc_hmac_key();
646        ctx.set_symmetric_key(key_aes_1_id, key_aes_1.clone())
647            .unwrap();
648        let key_aes_2_id = TestSymmKey::A(2);
649        let key_aes_2 = SymmetricCryptoKey::make_aes256_cbc_hmac_key();
650        ctx.set_symmetric_key(key_aes_2_id, key_aes_2.clone())
651            .unwrap();
652
653        // XChaCha20 Poly1305 keys
654        let key_xchacha_3_id = TestSymmKey::A(3);
655        let key_xchacha_3 = SymmetricCryptoKey::make_xchacha20_poly1305_key();
656        ctx.set_symmetric_key(key_xchacha_3_id, key_xchacha_3.clone())
657            .unwrap();
658        let key_xchacha_4_id = TestSymmKey::A(4);
659        let key_xchacha_4 = SymmetricCryptoKey::make_xchacha20_poly1305_key();
660        ctx.set_symmetric_key(key_xchacha_4_id, key_xchacha_4.clone())
661            .unwrap();
662
663        // Wrap and unwrap the keys
664        let wrapped_key_1_2 = ctx.wrap_symmetric_key(key_aes_1_id, key_aes_2_id).unwrap();
665        let wrapped_key_1_3 = ctx
666            .wrap_symmetric_key(key_aes_1_id, key_xchacha_3_id)
667            .unwrap();
668        let wrapped_key_3_1 = ctx
669            .wrap_symmetric_key(key_xchacha_3_id, key_aes_1_id)
670            .unwrap();
671        let wrapped_key_3_4 = ctx
672            .wrap_symmetric_key(key_xchacha_3_id, key_xchacha_4_id)
673            .unwrap();
674
675        // Unwrap the keys
676        let unwrapped_key_2 = ctx
677            .unwrap_symmetric_key(key_aes_1_id, key_aes_2_id, &wrapped_key_1_2)
678            .unwrap();
679        let unwrapped_key_3 = ctx
680            .unwrap_symmetric_key(key_aes_1_id, key_xchacha_3_id, &wrapped_key_1_3)
681            .unwrap();
682        let unwrapped_key_1 = ctx
683            .unwrap_symmetric_key(key_xchacha_3_id, key_aes_1_id, &wrapped_key_3_1)
684            .unwrap();
685        let unwrapped_key_4 = ctx
686            .unwrap_symmetric_key(key_xchacha_3_id, key_xchacha_4_id, &wrapped_key_3_4)
687            .unwrap();
688
689        // Assert that the unwrapped keys are the same as the original keys
690        assert_eq!(unwrapped_key_2, key_aes_2_id);
691        assert_eq!(unwrapped_key_3, key_xchacha_3_id);
692        assert_eq!(unwrapped_key_1, key_aes_1_id);
693        assert_eq!(unwrapped_key_4, key_xchacha_4_id);
694    }
695
696    #[test]
697    fn test_signing() {
698        let store: KeyStore<TestIds> = KeyStore::default();
699
700        // Generate and insert a key
701        let key_a0_id = TestSigningKey::A(0);
702        let key_a0 = SigningKey::make(SignatureAlgorithm::Ed25519);
703        let verifying_key = key_a0.to_verifying_key();
704        store
705            .context_mut()
706            .set_signing_key(key_a0_id, key_a0)
707            .unwrap();
708
709        assert!(store.context().has_signing_key(key_a0_id));
710
711        // Sign some data with the key
712        #[derive(Serialize, Deserialize)]
713        struct TestData {
714            data: String,
715        }
716        let signed_object = store
717            .context()
718            .sign(
719                key_a0_id,
720                &TestData {
721                    data: "Hello".to_string(),
722                },
723                &SigningNamespace::ExampleNamespace,
724            )
725            .unwrap();
726        let payload: Result<TestData, CryptoError> =
727            signed_object.verify_and_unwrap(&verifying_key, &SigningNamespace::ExampleNamespace);
728        assert!(payload.is_ok());
729
730        let (signature, serialized_message) = store
731            .context()
732            .sign_detached(
733                key_a0_id,
734                &TestData {
735                    data: "Hello".to_string(),
736                },
737                &SigningNamespace::ExampleNamespace,
738            )
739            .unwrap();
740        assert!(signature.verify(
741            serialized_message.as_bytes(),
742            &verifying_key,
743            &SigningNamespace::ExampleNamespace
744        ))
745    }
746
747    #[test]
748    fn test_account_key_rotation() {
749        let store: KeyStore<TestIds> = KeyStore::default();
750        let mut ctx = store.context_mut();
751
752        // Generate a new user key
753        let current_user_private_key_id = TestAsymmKey::A(0);
754        let current_user_signing_key_id = TestSigningKey::A(0);
755
756        // Make the keys
757        ctx.generate_symmetric_key(TestSymmKey::A(0)).unwrap();
758        ctx.make_signing_key(current_user_signing_key_id).unwrap();
759        ctx.set_asymmetric_key(
760            current_user_private_key_id,
761            AsymmetricCryptoKey::make(PublicKeyEncryptionAlgorithm::RsaOaepSha1),
762        )
763        .unwrap();
764
765        // Get the rotated account keys
766        let rotated_keys = ctx
767            .dangerous_get_v2_rotated_account_keys(
768                current_user_private_key_id,
769                current_user_signing_key_id,
770            )
771            .unwrap();
772
773        // Public/Private key
774        assert_eq!(
775            AsymmetricPublicCryptoKey::from_der(&rotated_keys.public_key)
776                .unwrap()
777                .to_der()
778                .unwrap(),
779            ctx.get_asymmetric_key(current_user_private_key_id)
780                .unwrap()
781                .to_public_key()
782                .to_der()
783                .unwrap()
784        );
785        let decrypted_private_key: Vec<u8> = rotated_keys
786            .private_key
787            .decrypt_with_key(&rotated_keys.user_key)
788            .unwrap();
789        let private_key =
790            AsymmetricCryptoKey::from_der(&Pkcs8PrivateKeyBytes::from(decrypted_private_key))
791                .unwrap();
792        assert_eq!(
793            private_key.to_der().unwrap(),
794            ctx.get_asymmetric_key(current_user_private_key_id)
795                .unwrap()
796                .to_der()
797                .unwrap()
798        );
799
800        // Signing Key
801        let decrypted_signing_key: Vec<u8> = rotated_keys
802            .signing_key
803            .decrypt_with_key(&rotated_keys.user_key)
804            .unwrap();
805        let signing_key =
806            SigningKey::from_cose(&CoseKeyBytes::from(decrypted_signing_key)).unwrap();
807        assert_eq!(
808            signing_key.to_cose(),
809            ctx.get_signing_key(current_user_signing_key_id)
810                .unwrap()
811                .to_cose(),
812        );
813
814        // Signed Public Key
815        let signed_public_key = rotated_keys.signed_public_key;
816        let unwrapped_key = signed_public_key
817            .verify_and_unwrap(
818                &ctx.get_signing_key(current_user_signing_key_id)
819                    .unwrap()
820                    .to_verifying_key(),
821            )
822            .unwrap();
823        assert_eq!(
824            unwrapped_key.to_der().unwrap(),
825            ctx.get_asymmetric_key(current_user_private_key_id)
826                .unwrap()
827                .to_public_key()
828                .to_der()
829                .unwrap()
830        );
831    }
832}