bitwarden_crypto/store/
context.rs

1use std::{
2    cell::Cell,
3    sync::{RwLockReadGuard, RwLockWriteGuard},
4};
5
6use coset::iana::KeyOperation;
7use serde::Serialize;
8use tracing::instrument;
9use zeroize::Zeroizing;
10
11use super::KeyStoreInner;
12use crate::{
13    BitwardenLegacyKeyBytes, ContentFormat, CoseEncrypt0Bytes, CoseKeyBytes, CoseSerializable,
14    CryptoError, EncString, KeyDecryptable, KeyEncryptable, KeyId, KeyIds, LocalId,
15    Pkcs8PrivateKeyBytes, PrivateKey, PublicKey, PublicKeyEncryptionAlgorithm, Result,
16    RotatedUserKeys, Signature, SignatureAlgorithm, SignedObject, SignedPublicKey,
17    SignedPublicKeyMessage, SigningKey, SymmetricCryptoKey, SymmetricKeyAlgorithm, VerifyingKey,
18    derive_shareable_key, error::UnsupportedOperationError, signing, store::backend::StoreBackend,
19};
20
21/// The context of a crypto operation using [super::KeyStore]
22///
23/// This will usually be accessed from an implementation of [crate::Decryptable] or
24/// [crate::CompositeEncryptable], [crate::PrimitiveEncryptable],
25/// but can also be obtained
26/// through [super::KeyStore::context]
27///
28/// This context contains access to the user keys stored in the [super::KeyStore] (sometimes
29/// referred to as `global keys`) and it also contains it's own individual secure backend for key
30/// storage. Keys stored in this individual backend are usually referred to as `local keys`, they
31/// will be cleared when this context goes out of scope and is dropped and they do not affect either
32/// the global [super::KeyStore] or other instances of contexts.
33///
34/// This context-local storage is recommended for ephemeral and temporary keys that are decrypted
35/// during the course of a decrypt/encrypt operation, but won't be used after the operation itself
36/// is complete.
37///
38/// ```rust
39/// # use bitwarden_crypto::*;
40/// # key_ids! {
41/// #     #[symmetric]
42/// #     pub enum SymmKeyId {
43/// #         User,
44/// #         #[local]
45/// #         Local(LocalId),
46/// #     }
47/// #     #[private]
48/// #     pub enum PrivateKeyId {
49/// #         UserPrivate,
50/// #         #[local]
51/// #         Local(LocalId),
52/// #     }
53/// #     #[signing]
54/// #     pub enum SigningKeyId {
55/// #         UserSigning,
56/// #         #[local]
57/// #         Local(LocalId),
58/// #     }
59/// #     pub Ids => SymmKeyId, PrivateKeyId, SigningKeyId;
60/// # }
61/// struct Data {
62///     key: EncString,
63///     name: String,
64/// }
65/// # impl IdentifyKey<SymmKeyId> for Data {
66/// #    fn key_identifier(&self) -> SymmKeyId {
67/// #        SymmKeyId::User
68/// #    }
69/// # }
70///
71///
72/// impl CompositeEncryptable<Ids, SymmKeyId, EncString> for Data {
73///     fn encrypt_composite(&self, ctx: &mut KeyStoreContext<Ids>, key: SymmKeyId) -> Result<EncString, CryptoError> {
74///         let local_key_id = ctx.unwrap_symmetric_key(key, &self.key)?;
75///         self.name.encrypt(ctx, local_key_id)
76///     }
77/// }
78/// ```
79#[must_use]
80pub struct KeyStoreContext<'a, Ids: KeyIds> {
81    pub(super) global_keys: GlobalKeys<'a, Ids>,
82
83    pub(super) local_symmetric_keys: Box<dyn StoreBackend<Ids::Symmetric>>,
84    pub(super) local_private_keys: Box<dyn StoreBackend<Ids::Private>>,
85    pub(super) local_signing_keys: Box<dyn StoreBackend<Ids::Signing>>,
86
87    pub(super) security_state_version: u64,
88
89    // Make sure the context is !Send & !Sync
90    pub(super) _phantom: std::marker::PhantomData<(Cell<()>, RwLockReadGuard<'static, ()>)>,
91}
92
93/// A KeyStoreContext is usually limited to a read only access to the global keys,
94/// which allows us to have multiple read only contexts at the same time and do multitheaded
95/// encryption/decryption. We also have the option to create a read/write context, which allows us
96/// to modify the global keys, but only allows one context at a time. This is controlled by a
97/// [std::sync::RwLock] on the global keys, and this struct stores both types of guards.
98pub(crate) enum GlobalKeys<'a, Ids: KeyIds> {
99    ReadOnly(RwLockReadGuard<'a, KeyStoreInner<Ids>>),
100    ReadWrite(RwLockWriteGuard<'a, KeyStoreInner<Ids>>),
101}
102
103impl<Ids: KeyIds> GlobalKeys<'_, Ids> {
104    /// Get a shared reference to the underlying `KeyStoreInner`.
105    ///
106    /// This returns a shared reference regardless of whether the global keys were locked
107    /// for read-only or read-write access. Callers who need mutable access should use
108    /// `get_mut` which will return an error when the context is read-only.
109    pub fn get(&self) -> &KeyStoreInner<Ids> {
110        match self {
111            GlobalKeys::ReadOnly(keys) => keys,
112            GlobalKeys::ReadWrite(keys) => keys,
113        }
114    }
115
116    /// Get a mutable reference to the underlying `KeyStoreInner`.
117    ///
118    /// This will succeed only when the context was created with write access. If the
119    /// context is read-only an error (`CryptoError::ReadOnlyKeyStore`) is returned.
120    ///
121    /// # Errors
122    /// Returns [`CryptoError::ReadOnlyKeyStore`] when attempting to get mutable access from
123    /// a read-only context.
124    pub fn get_mut(&mut self) -> Result<&mut KeyStoreInner<Ids>> {
125        match self {
126            GlobalKeys::ReadOnly(_) => Err(CryptoError::ReadOnlyKeyStore),
127            GlobalKeys::ReadWrite(keys) => Ok(keys),
128        }
129    }
130}
131
132impl<Ids: KeyIds> KeyStoreContext<'_, Ids> {
133    /// Clears all the local keys stored in this context
134    /// This will not affect the global keys even if this context has write access.
135    /// To clear the global keys, you need to use [super::KeyStore::clear] instead.
136    pub fn clear_local(&mut self) {
137        self.local_symmetric_keys.clear();
138        self.local_private_keys.clear();
139        self.local_signing_keys.clear();
140    }
141
142    /// Returns the version of the security state of the key context. This describes the user's
143    /// encryption version and can be used to disable certain old / dangerous format features
144    /// safely.
145    pub fn get_security_state_version(&self) -> u64 {
146        self.security_state_version
147    }
148
149    /// Remove all symmetric keys from the context for which the predicate returns false
150    /// This will also remove the keys from the global store if this context has write access
151    pub fn retain_symmetric_keys(&mut self, f: fn(Ids::Symmetric) -> bool) {
152        if let Ok(keys) = self.global_keys.get_mut() {
153            keys.symmetric_keys.retain(f);
154        }
155        self.local_symmetric_keys.retain(f);
156    }
157
158    /// Remove all private keys from the context for which the predicate returns false
159    /// This will also remove the keys from the global store if this context has write access
160    pub fn retain_private_keys(&mut self, f: fn(Ids::Private) -> bool) {
161        if let Ok(keys) = self.global_keys.get_mut() {
162            keys.private_keys.retain(f);
163        }
164        self.local_private_keys.retain(f);
165    }
166
167    fn drop_symmetric_key(&mut self, key_id: Ids::Symmetric) -> Result<()> {
168        if key_id.is_local() {
169            self.local_symmetric_keys.remove(key_id);
170        } else {
171            self.global_keys.get_mut()?.symmetric_keys.remove(key_id);
172        }
173        Ok(())
174    }
175
176    fn drop_private_key(&mut self, key_id: Ids::Private) -> Result<()> {
177        if key_id.is_local() {
178            self.local_private_keys.remove(key_id);
179        } else {
180            self.global_keys.get_mut()?.private_keys.remove(key_id);
181        }
182        Ok(())
183    }
184
185    fn drop_signing_key(&mut self, key_id: Ids::Signing) -> Result<()> {
186        if key_id.is_local() {
187            self.local_signing_keys.remove(key_id);
188        } else {
189            self.global_keys.get_mut()?.signing_keys.remove(key_id);
190        }
191        Ok(())
192    }
193
194    // TODO: All these encrypt x key with x key look like they need to be made generic,
195    // but I haven't found the best way to do that yet.
196
197    /// Decrypt a symmetric key into the context by using an already existing symmetric key
198    ///
199    /// # Arguments
200    ///
201    /// * `wrapping_key` - The key id used to decrypt the `wrapped_key`. It must already exist in
202    ///   the context
203    /// * `new_key_id` - The key id where the decrypted key will be stored. If it already exists, it
204    ///   will be overwritten
205    /// * `wrapped_key` - The key to decrypt
206    #[instrument(skip(self, wrapped_key), err)]
207    pub fn unwrap_symmetric_key(
208        &mut self,
209        wrapping_key: Ids::Symmetric,
210        wrapped_key: &EncString,
211    ) -> Result<Ids::Symmetric> {
212        let wrapping_key = self.get_symmetric_key(wrapping_key)?;
213
214        let key = match (wrapped_key, wrapping_key) {
215            (EncString::Aes256Cbc_B64 { iv, data }, SymmetricCryptoKey::Aes256CbcKey(key)) => {
216                SymmetricCryptoKey::try_from(&BitwardenLegacyKeyBytes::from(
217                    crate::aes::decrypt_aes256(iv, data.clone(), &key.enc_key)
218                        .map_err(|_| CryptoError::Decrypt)?,
219                ))?
220            }
221            (
222                EncString::Aes256Cbc_HmacSha256_B64 { iv, mac, data },
223                SymmetricCryptoKey::Aes256CbcHmacKey(key),
224            ) => SymmetricCryptoKey::try_from(&BitwardenLegacyKeyBytes::from(
225                crate::aes::decrypt_aes256_hmac(iv, mac, data.clone(), &key.mac_key, &key.enc_key)
226                    .map_err(|_| CryptoError::Decrypt)?,
227            ))?,
228            (
229                EncString::Cose_Encrypt0_B64 { data },
230                SymmetricCryptoKey::XChaCha20Poly1305Key(key),
231            ) => {
232                let (content_bytes, content_format) = crate::cose::decrypt_xchacha20_poly1305(
233                    &CoseEncrypt0Bytes::from(data.clone()),
234                    key,
235                )?;
236                match content_format {
237                    ContentFormat::BitwardenLegacyKey => {
238                        SymmetricCryptoKey::try_from(&BitwardenLegacyKeyBytes::from(content_bytes))?
239                    }
240                    ContentFormat::CoseKey => SymmetricCryptoKey::try_from_cose(&content_bytes)?,
241                    _ => return Err(CryptoError::InvalidKey),
242                }
243            }
244            _ => {
245                tracing::warn!("Unsupported unwrap operation for the given key and data");
246                return Err(CryptoError::InvalidKey);
247            }
248        };
249
250        let new_key_id = Ids::Symmetric::new_local(LocalId::new());
251
252        #[allow(deprecated)]
253        self.set_symmetric_key(new_key_id, key)?;
254
255        // Returning the new key identifier for convenience
256        Ok(new_key_id)
257    }
258
259    /// Move a symmetric key from a local identifier to a global identifier within the context
260    ///
261    /// The key value is copied to `to` and the original identifier `from` is removed.
262    ///
263    /// # Errors
264    /// Returns an error if the source key does not exist or if setting the destination key
265    /// fails (for example due to read-only global store).
266    pub fn persist_symmetric_key(
267        &mut self,
268        from: Ids::Symmetric,
269        to: Ids::Symmetric,
270    ) -> Result<()> {
271        if !from.is_local() || to.is_local() {
272            return Err(CryptoError::InvalidKeyStoreOperation);
273        }
274        let key = self.get_symmetric_key(from)?.to_owned();
275        self.drop_symmetric_key(from)?;
276        #[allow(deprecated)]
277        self.set_symmetric_key(to, key)?;
278        Ok(())
279    }
280
281    /// Move a private key from a local identifier to a global identifier within the context
282    ///
283    /// The key value is copied to `to` and the original identifier `from` is removed.
284    ///
285    /// # Errors
286    /// Returns an error if the source key does not exist or if setting the destination key
287    /// fails (for example due to read-only global store).
288    pub fn persist_private_key(&mut self, from: Ids::Private, to: Ids::Private) -> Result<()> {
289        if !from.is_local() || to.is_local() {
290            return Err(CryptoError::InvalidKeyStoreOperation);
291        }
292        let key = self.get_private_key(from)?.to_owned();
293        self.drop_private_key(from)?;
294        #[allow(deprecated)]
295        self.set_private_key(to, key)?;
296        Ok(())
297    }
298
299    /// Move a signing key from a local identifier to a global identifier within the context
300    ///
301    /// The key value at `from` will be copied to `to` and the original `from` will be removed.
302    ///
303    /// # Errors
304    /// Returns an error if the source key does not exist or updating the destination fails.
305    pub fn persist_signing_key(&mut self, from: Ids::Signing, to: Ids::Signing) -> Result<()> {
306        if !from.is_local() || to.is_local() {
307            return Err(CryptoError::InvalidKeyStoreOperation);
308        }
309        let key = self.get_signing_key(from)?.to_owned();
310        self.drop_signing_key(from)?;
311        #[allow(deprecated)]
312        self.set_signing_key(to, key)?;
313        Ok(())
314    }
315
316    /// Wrap (encrypt) a signing key with a symmetric key.
317    ///
318    /// The signing key identified by `key_to_wrap` will be serialized to COSE and encrypted
319    /// with the symmetric `wrapping_key`, returning an `EncString` suitable for storage or
320    /// transport.
321    ///
322    /// # Errors
323    /// Returns an error if either key id does not exist or the encryption fails.
324    pub fn wrap_signing_key(
325        &self,
326        wrapping_key: Ids::Symmetric,
327        key_to_wrap: Ids::Signing,
328    ) -> Result<EncString> {
329        let wrapping_key = self.get_symmetric_key(wrapping_key)?;
330        let signing_key = self.get_signing_key(key_to_wrap)?.to_owned();
331        signing_key.to_cose().encrypt_with_key(wrapping_key)
332    }
333
334    /// Wrap (encrypt) a private key with a symmetric key.
335    ///
336    /// The private key identified by `key_to_wrap` will be serialized to DER (PKCS#8) and
337    /// encrypted with `wrapping_key`, returning an `EncString` suitable for storage.
338    ///
339    /// # Errors
340    /// Returns an error if the keys are missing or serialization/encryption fails.
341    pub fn wrap_private_key(
342        &self,
343        wrapping_key: Ids::Symmetric,
344        key_to_wrap: Ids::Private,
345    ) -> Result<EncString> {
346        let wrapping_key = self.get_symmetric_key(wrapping_key)?;
347        let private_key = self.get_private_key(key_to_wrap)?.to_owned();
348        private_key.to_der()?.encrypt_with_key(wrapping_key)
349    }
350
351    /// Decrypt and import a previously wrapped private key into the context.
352    ///
353    /// The `wrapped_key` will be decrypted using `wrapping_key` and parsed as a PKCS#8
354    /// private key; the resulting key will be inserted as a local private key and the
355    /// new local identifier returned.
356    ///
357    /// # Errors
358    /// Returns an error if decryption or parsing fails.
359    #[instrument(skip(self, wrapped_key), err)]
360    pub fn unwrap_private_key(
361        &mut self,
362        wrapping_key: Ids::Symmetric,
363        wrapped_key: &EncString,
364    ) -> Result<Ids::Private> {
365        let wrapping_key = self.get_symmetric_key(wrapping_key)?;
366        let private_key_bytes: Vec<u8> = wrapped_key.decrypt_with_key(wrapping_key)?;
367        let private_key = PrivateKey::from_der(&Pkcs8PrivateKeyBytes::from(private_key_bytes))?;
368        Ok(self.add_local_private_key(private_key))
369    }
370
371    /// Decrypt and import a previously wrapped signing key into the context.
372    ///
373    /// The wrapped COSE key will be decrypted with `wrapping_key` and parsed into a
374    /// `SigningKey` which is inserted as a local signing key. The new local identifier
375    /// is returned.
376    ///
377    /// # Errors
378    /// Returns an error if decryption or parsing fails.
379    pub fn unwrap_signing_key(
380        &mut self,
381        wrapping_key: Ids::Symmetric,
382        wrapped_key: &EncString,
383    ) -> Result<Ids::Signing> {
384        let wrapping_key = self.get_symmetric_key(wrapping_key)?;
385        let signing_key_bytes: Vec<u8> = wrapped_key.decrypt_with_key(wrapping_key)?;
386        let signing_key = SigningKey::from_cose(&CoseKeyBytes::from(signing_key_bytes))?;
387        Ok(self.add_local_signing_key(signing_key))
388    }
389
390    /// Return the verifying (public) key corresponding to a signing key identifier.
391    ///
392    /// This converts the stored `SigningKey` into a `VerifyingKey` suitable for
393    /// signature verification operations.
394    ///
395    /// # Errors
396    /// Returns an error if the signing key id does not exist.
397    pub fn get_verifying_key(&self, signing_key_id: Ids::Signing) -> Result<VerifyingKey> {
398        let signing_key = self.get_signing_key(signing_key_id)?;
399        Ok(signing_key.to_verifying_key())
400    }
401
402    /// Return the public key corresponding to an private key identifier.
403    ///
404    /// This converts the stored private key into its public key representation.
405    ///
406    /// # Errors
407    /// Returns an error if the private key id does not exist.
408    pub fn get_public_key(&self, private_key_id: Ids::Private) -> Result<PublicKey> {
409        let private_key = self.get_private_key(private_key_id)?;
410        Ok(private_key.to_public_key())
411    }
412
413    /// Encrypt and return a symmetric key from the context by using an already existing symmetric
414    /// key
415    ///
416    /// # Arguments
417    ///
418    /// * `wrapping_key` - The key id used to wrap (encrypt) the `key_to_wrap`. It must already
419    ///   exist in the context
420    /// * `key_to_wrap` - The key id to wrap. It must already exist in the context
421    pub fn wrap_symmetric_key(
422        &self,
423        wrapping_key: Ids::Symmetric,
424        key_to_wrap: Ids::Symmetric,
425    ) -> Result<EncString> {
426        use SymmetricCryptoKey::*;
427
428        let wrapping_key_instance = self.get_symmetric_key(wrapping_key)?;
429        let key_to_wrap_instance = self.get_symmetric_key(key_to_wrap)?;
430        // `Aes256CbcHmacKey` can wrap keys by encrypting their byte serialization obtained using
431        // `SymmetricCryptoKey::to_encoded()`. `XChaCha20Poly1305Key` need to specify the
432        // content format to be either octet stream, in case the wrapped key is a Aes256CbcHmacKey
433        // or `Aes256CbcKey`, or by specifying the content format to be CoseKey, in case the
434        // wrapped key is a `XChaCha20Poly1305Key`.
435        match (wrapping_key_instance, key_to_wrap_instance) {
436            (
437                Aes256CbcHmacKey(_),
438                Aes256CbcHmacKey(_) | Aes256CbcKey(_) | XChaCha20Poly1305Key(_),
439            ) => self.encrypt_data_with_symmetric_key(
440                wrapping_key,
441                key_to_wrap_instance
442                    .to_encoded()
443                    .as_ref()
444                    .to_vec()
445                    .as_slice(),
446                ContentFormat::BitwardenLegacyKey,
447            ),
448            (XChaCha20Poly1305Key(_), _) => {
449                let encoded = key_to_wrap_instance.to_encoded_raw();
450                let content_format = encoded.content_format();
451                self.encrypt_data_with_symmetric_key(
452                    wrapping_key,
453                    Into::<Vec<u8>>::into(encoded).as_slice(),
454                    content_format,
455                )
456            }
457            _ => Err(CryptoError::OperationNotSupported(
458                UnsupportedOperationError::EncryptionNotImplementedForKey,
459            )),
460        }
461    }
462
463    /// Returns `true` if the context has a symmetric key with the given identifier
464    pub fn has_symmetric_key(&self, key_id: Ids::Symmetric) -> bool {
465        self.get_symmetric_key(key_id).is_ok()
466    }
467
468    /// Returns `true` if the context has a private key with the given identifier
469    pub fn has_private_key(&self, key_id: Ids::Private) -> bool {
470        self.get_private_key(key_id).is_ok()
471    }
472
473    /// Returns `true` if the context has a signing key with the given identifier
474    pub fn has_signing_key(&self, key_id: Ids::Signing) -> bool {
475        self.get_signing_key(key_id).is_ok()
476    }
477
478    /// Generate a new random symmetric key and store it in the context
479    pub fn generate_symmetric_key(&mut self) -> Ids::Symmetric {
480        self.add_local_symmetric_key(SymmetricCryptoKey::make_aes256_cbc_hmac_key())
481    }
482
483    /// Generate a new symmetric encryption key using the specified algorithm and store it in the
484    /// context as a local key
485    pub fn make_symmetric_key(&mut self, algorithm: SymmetricKeyAlgorithm) -> Ids::Symmetric {
486        self.add_local_symmetric_key(SymmetricCryptoKey::make(algorithm))
487    }
488
489    /// Makes a new private encryption key using the current default algorithm, and stores it in
490    /// the context as a local key
491    pub fn make_private_key(&mut self, algorithm: PublicKeyEncryptionAlgorithm) -> Ids::Private {
492        self.add_local_private_key(PrivateKey::make(algorithm))
493    }
494
495    /// Makes a new signing key using the current default algorithm, and stores it in the context as
496    /// a local key
497    pub fn make_signing_key(&mut self, algorithm: SignatureAlgorithm) -> Ids::Signing {
498        self.add_local_signing_key(SigningKey::make(algorithm))
499    }
500
501    /// Derive a shareable key using hkdf from secret and name and store it in the context.
502    ///
503    /// A specialized variant of this function was called `CryptoService.makeSendKey` in the
504    /// Bitwarden `clients` repository.
505    pub fn derive_shareable_key(
506        &mut self,
507        secret: Zeroizing<[u8; 16]>,
508        name: &str,
509        info: Option<&str>,
510    ) -> Result<Ids::Symmetric> {
511        let key_id = Ids::Symmetric::new_local(LocalId::new());
512        #[allow(deprecated)]
513        self.set_symmetric_key(
514            key_id,
515            SymmetricCryptoKey::Aes256CbcHmacKey(derive_shareable_key(secret, name, info)),
516        )?;
517        Ok(key_id)
518    }
519
520    /// Return a reference to a symmetric key stored in the context.
521    ///
522    /// Deprecated: intended only for internal use and tests. This exposes the underlying
523    /// `SymmetricCryptoKey` reference directly and should not be used by external code. Use
524    /// the higher-level APIs (for example encryption/decryption helpers) or `get_symmetric_key`
525    /// internally when possible.
526    ///
527    /// # Errors
528    /// Returns [`CryptoError::MissingKeyId`] if the key id does not exist in the context.
529    #[deprecated(note = "This function should ideally never be used outside this crate")]
530    pub fn dangerous_get_symmetric_key(
531        &self,
532        key_id: Ids::Symmetric,
533    ) -> Result<&SymmetricCryptoKey> {
534        self.get_symmetric_key(key_id)
535    }
536
537    /// Return a reference to an asymmetric (private) key stored in the context.
538    ///
539    /// Deprecated: intended only for internal use and tests. This exposes the underlying
540    /// `PrivateKey` reference directly and should not be used by external code. Prefer
541    /// using the public key via `get_public_key` or other higher-level APIs instead.
542    ///
543    /// # Errors
544    /// Returns [`CryptoError::MissingKeyId`] if the key id does not exist in the context.
545    #[deprecated(note = "This function should ideally never be used outside this crate")]
546    pub fn dangerous_get_private_key(&self, key_id: Ids::Private) -> Result<&PrivateKey> {
547        self.get_private_key(key_id)
548    }
549
550    /// Makes a signed public key from a private key and signing key stored in context.
551    /// Signing a public key asserts ownership, and makes the claim to other users that if they want
552    /// to share with you, they can use this public key.
553    pub fn make_signed_public_key(
554        &self,
555        private_key_id: Ids::Private,
556        signing_key_id: Ids::Signing,
557    ) -> Result<SignedPublicKey> {
558        let public_key = self.get_private_key(private_key_id)?.to_public_key();
559        let signing_key = self.get_signing_key(signing_key_id)?;
560        let signed_public_key =
561            SignedPublicKeyMessage::from_public_key(&public_key)?.sign(signing_key)?;
562        Ok(signed_public_key)
563    }
564
565    pub(crate) fn get_symmetric_key(&self, key_id: Ids::Symmetric) -> Result<&SymmetricCryptoKey> {
566        if key_id.is_local() {
567            self.local_symmetric_keys.get(key_id)
568        } else {
569            self.global_keys.get().symmetric_keys.get(key_id)
570        }
571        .ok_or_else(|| crate::CryptoError::MissingKeyId(format!("{key_id:?}")))
572    }
573
574    pub(super) fn get_private_key(&self, key_id: Ids::Private) -> Result<&PrivateKey> {
575        if key_id.is_local() {
576            self.local_private_keys.get(key_id)
577        } else {
578            self.global_keys.get().private_keys.get(key_id)
579        }
580        .ok_or_else(|| crate::CryptoError::MissingKeyId(format!("{key_id:?}")))
581    }
582
583    pub(super) fn get_signing_key(&self, key_id: Ids::Signing) -> Result<&SigningKey> {
584        if key_id.is_local() {
585            self.local_signing_keys.get(key_id)
586        } else {
587            self.global_keys.get().signing_keys.get(key_id)
588        }
589        .ok_or_else(|| crate::CryptoError::MissingKeyId(format!("{key_id:?}")))
590    }
591
592    /// Set a symmetric key in the context.
593    ///
594    /// # Errors
595    /// Returns [`CryptoError::ReadOnlyKeyStore`] if the context does not have write access when
596    /// attempting to modify the global store.
597    #[deprecated(note = "This function should ideally never be used outside this crate")]
598    pub fn set_symmetric_key(
599        &mut self,
600        key_id: Ids::Symmetric,
601        key: SymmetricCryptoKey,
602    ) -> Result<()> {
603        self.set_symmetric_key_internal(key_id, key)
604    }
605
606    pub(crate) fn set_symmetric_key_internal(
607        &mut self,
608        key_id: Ids::Symmetric,
609        key: SymmetricCryptoKey,
610    ) -> Result<()> {
611        if key_id.is_local() {
612            self.local_symmetric_keys.upsert(key_id, key);
613        } else {
614            self.global_keys
615                .get_mut()?
616                .symmetric_keys
617                .upsert(key_id, key);
618        }
619        Ok(())
620    }
621
622    /// Add a new symmetric key to the local context, returning a new unique identifier for it.
623    pub fn add_local_symmetric_key(&mut self, key: SymmetricCryptoKey) -> Ids::Symmetric {
624        let key_id = Ids::Symmetric::new_local(LocalId::new());
625        self.local_symmetric_keys.upsert(key_id, key);
626        key_id
627    }
628
629    /// Get the type of a symmetric key stored in the context.
630    pub fn get_symmetric_key_algorithm(
631        &self,
632        key_id: Ids::Symmetric,
633    ) -> Result<SymmetricKeyAlgorithm> {
634        let key = self.get_symmetric_key(key_id)?;
635        match key {
636            // Note this is dropped soon
637            SymmetricCryptoKey::Aes256CbcKey(_) => Err(CryptoError::OperationNotSupported(
638                UnsupportedOperationError::EncryptionNotImplementedForKey,
639            )),
640            SymmetricCryptoKey::Aes256CbcHmacKey(_) => Ok(SymmetricKeyAlgorithm::Aes256CbcHmac),
641            SymmetricCryptoKey::XChaCha20Poly1305Key(_) => {
642                Ok(SymmetricKeyAlgorithm::XChaCha20Poly1305)
643            }
644        }
645    }
646
647    /// Set a private key in the context.
648    ///
649    /// # Errors
650    /// Returns [`CryptoError::ReadOnlyKeyStore`] if attempting to write to the global store when
651    /// the context is read-only.
652    #[deprecated(note = "This function should ideally never be used outside this crate")]
653    pub fn set_private_key(&mut self, key_id: Ids::Private, key: PrivateKey) -> Result<()> {
654        if key_id.is_local() {
655            self.local_private_keys.upsert(key_id, key);
656        } else {
657            self.global_keys.get_mut()?.private_keys.upsert(key_id, key);
658        }
659        Ok(())
660    }
661
662    /// Add a new private key to the local context, returning a new unique identifier for it.
663    pub fn add_local_private_key(&mut self, key: PrivateKey) -> Ids::Private {
664        let key_id = Ids::Private::new_local(LocalId::new());
665        self.local_private_keys.upsert(key_id, key);
666        key_id
667    }
668
669    /// Sets a signing key in the context
670    ///
671    /// # Errors
672    /// Returns [`CryptoError::ReadOnlyKeyStore`] if attempting to write to the global store when
673    /// the context is read-only.
674    #[deprecated(note = "This function should ideally never be used outside this crate")]
675    pub fn set_signing_key(&mut self, key_id: Ids::Signing, key: SigningKey) -> Result<()> {
676        if key_id.is_local() {
677            self.local_signing_keys.upsert(key_id, key);
678        } else {
679            self.global_keys.get_mut()?.signing_keys.upsert(key_id, key);
680        }
681        Ok(())
682    }
683
684    /// Add a new signing key to the local context, returning a new unique identifier for it.
685    pub fn add_local_signing_key(&mut self, key: SigningKey) -> Ids::Signing {
686        let key_id = Ids::Signing::new_local(LocalId::new());
687        self.local_signing_keys.upsert(key_id, key);
688        key_id
689    }
690
691    #[instrument(skip(self, data), err)]
692    pub(crate) fn decrypt_data_with_symmetric_key(
693        &self,
694        key: Ids::Symmetric,
695        data: &EncString,
696    ) -> Result<Vec<u8>> {
697        let key = self.get_symmetric_key(key)?;
698
699        match (data, key) {
700            (EncString::Aes256Cbc_B64 { iv, data }, SymmetricCryptoKey::Aes256CbcKey(key)) => {
701                crate::aes::decrypt_aes256(iv, data.clone(), &key.enc_key)
702                    .map_err(|_| CryptoError::Decrypt)
703            }
704            (
705                EncString::Aes256Cbc_HmacSha256_B64 { iv, mac, data },
706                SymmetricCryptoKey::Aes256CbcHmacKey(key),
707            ) => crate::aes::decrypt_aes256_hmac(iv, mac, data.clone(), &key.mac_key, &key.enc_key)
708                .map_err(|_| CryptoError::Decrypt),
709            (
710                EncString::Cose_Encrypt0_B64 { data },
711                SymmetricCryptoKey::XChaCha20Poly1305Key(key),
712            ) => {
713                let (data, _) = crate::cose::decrypt_xchacha20_poly1305(
714                    &CoseEncrypt0Bytes::from(data.clone()),
715                    key,
716                )?;
717                Ok(data)
718            }
719            _ => {
720                tracing::warn!("Unsupported decryption operation for the given key and data");
721                Err(CryptoError::InvalidKey)
722            }
723        }
724    }
725
726    pub(crate) fn encrypt_data_with_symmetric_key(
727        &self,
728        key: Ids::Symmetric,
729        data: &[u8],
730        content_format: ContentFormat,
731    ) -> Result<EncString> {
732        let key = self.get_symmetric_key(key)?;
733        match key {
734            SymmetricCryptoKey::Aes256CbcKey(_) => Err(CryptoError::OperationNotSupported(
735                UnsupportedOperationError::EncryptionNotImplementedForKey,
736            )),
737            SymmetricCryptoKey::Aes256CbcHmacKey(key) => EncString::encrypt_aes256_hmac(data, key),
738            SymmetricCryptoKey::XChaCha20Poly1305Key(key) => {
739                if !key.supported_operations.contains(&KeyOperation::Encrypt) {
740                    return Err(CryptoError::KeyOperationNotSupported(KeyOperation::Encrypt));
741                }
742                EncString::encrypt_xchacha20_poly1305(data, key, content_format)
743            }
744        }
745    }
746
747    /// Signs the given data using the specified signing key, for the given
748    /// [crate::SigningNamespace] and returns the signature and the serialized message. See
749    /// [crate::SigningKey::sign]
750    pub fn sign<Message: Serialize>(
751        &self,
752        key: Ids::Signing,
753        message: &Message,
754        namespace: &crate::SigningNamespace,
755    ) -> Result<SignedObject> {
756        self.get_signing_key(key)?.sign(message, namespace)
757    }
758
759    /// Signs the given data using the specified signing key, for the given
760    /// [crate::SigningNamespace] and returns the signature and the serialized message. See
761    /// [crate::SigningKey::sign_detached]
762    #[allow(unused)]
763    pub(crate) fn sign_detached<Message: Serialize>(
764        &self,
765        key: Ids::Signing,
766        message: &Message,
767        namespace: &crate::SigningNamespace,
768    ) -> Result<(Signature, signing::SerializedMessage)> {
769        self.get_signing_key(key)?.sign_detached(message, namespace)
770    }
771
772    /// Re-encrypts the user's keys with the provided symmetric key for a v2 user.
773    pub fn dangerous_get_v2_rotated_account_keys(
774        &self,
775        current_user_private_key_id: Ids::Private,
776        current_user_signing_key_id: Ids::Signing,
777    ) -> Result<RotatedUserKeys> {
778        #[expect(deprecated)]
779        crate::dangerous_get_v2_rotated_account_keys(
780            current_user_private_key_id,
781            current_user_signing_key_id,
782            self,
783        )
784    }
785}
786
787#[cfg(test)]
788#[allow(deprecated)]
789mod tests {
790    use serde::{Deserialize, Serialize};
791
792    use crate::{
793        CompositeEncryptable, CoseKeyBytes, CoseSerializable, CryptoError, Decryptable,
794        KeyDecryptable, Pkcs8PrivateKeyBytes, PrivateKey, PublicKey, PublicKeyEncryptionAlgorithm,
795        SignatureAlgorithm, SigningKey, SigningNamespace, SymmetricCryptoKey,
796        SymmetricKeyAlgorithm,
797        store::{
798            KeyStore,
799            tests::{Data, DataView},
800        },
801        traits::tests::{TestIds, TestSigningKey, TestSymmKey},
802    };
803
804    #[test]
805    fn test_set_signing_key() {
806        let store: KeyStore<TestIds> = KeyStore::default();
807
808        // Generate and insert a key
809        let key_a0_id = TestSigningKey::A(0);
810        let key_a0 = SigningKey::make(SignatureAlgorithm::Ed25519);
811        store
812            .context_mut()
813            .set_signing_key(key_a0_id, key_a0)
814            .unwrap();
815    }
816
817    #[test]
818    fn test_set_keys_for_encryption() {
819        let store: KeyStore<TestIds> = KeyStore::default();
820
821        // Generate and insert a key
822        let key_a0_id = TestSymmKey::A(0);
823        let mut ctx = store.context_mut();
824        let local_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
825        ctx.persist_symmetric_key(local_key_id, TestSymmKey::A(0))
826            .unwrap();
827
828        assert!(ctx.has_symmetric_key(key_a0_id));
829
830        // Encrypt some data with the key
831        let data = DataView("Hello, World!".to_string(), key_a0_id);
832        let _encrypted: Data = data.encrypt_composite(&mut ctx, key_a0_id).unwrap();
833    }
834
835    #[test]
836    fn test_key_encryption() {
837        let store: KeyStore<TestIds> = KeyStore::default();
838
839        let mut ctx = store.context();
840
841        // Generate and insert a key
842        let key_1_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
843
844        assert!(ctx.has_symmetric_key(key_1_id));
845
846        // Generate and insert a new key
847        let key_2_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
848
849        assert!(ctx.has_symmetric_key(key_2_id));
850
851        // Encrypt the new key with the old key
852        let key_2_enc = ctx.wrap_symmetric_key(key_1_id, key_2_id).unwrap();
853
854        // Decrypt the new key with the old key in a different identifier
855        let new_key_id = ctx.unwrap_symmetric_key(key_1_id, &key_2_enc).unwrap();
856
857        // Now `key_2_id` and `new_key_id` contain the same key, so we should be able to encrypt
858        // with one and decrypt with the other
859
860        let data = DataView("Hello, World!".to_string(), key_2_id);
861        let encrypted = data.encrypt_composite(&mut ctx, key_2_id).unwrap();
862
863        let decrypted1 = encrypted.decrypt(&mut ctx, key_2_id).unwrap();
864        let decrypted2 = encrypted.decrypt(&mut ctx, new_key_id).unwrap();
865
866        // Assert that the decrypted data is the same
867        assert_eq!(decrypted1.0, decrypted2.0);
868    }
869
870    #[test]
871    fn test_wrap_unwrap() {
872        let store: KeyStore<TestIds> = KeyStore::default();
873        let mut ctx = store.context_mut();
874
875        // Aes256 CBC HMAC keys
876        let key_aes_1_id = TestSymmKey::A(1);
877        let local_key_1_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
878        ctx.persist_symmetric_key(local_key_1_id, key_aes_1_id)
879            .unwrap();
880        let key_aes_2_id = TestSymmKey::A(2);
881        let local_key_2_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
882        ctx.persist_symmetric_key(local_key_2_id, key_aes_2_id)
883            .unwrap();
884
885        // XChaCha20 Poly1305 keys
886        let key_xchacha_3_id = TestSymmKey::A(3);
887        let key_xchacha_3 = SymmetricCryptoKey::make_xchacha20_poly1305_key();
888        ctx.set_symmetric_key(key_xchacha_3_id, key_xchacha_3.clone())
889            .unwrap();
890        let key_xchacha_4_id = TestSymmKey::A(4);
891        let key_xchacha_4 = SymmetricCryptoKey::make_xchacha20_poly1305_key();
892        ctx.set_symmetric_key(key_xchacha_4_id, key_xchacha_4.clone())
893            .unwrap();
894
895        // Wrap and unwrap the keys
896        let wrapped_key_1_2 = ctx.wrap_symmetric_key(key_aes_1_id, key_aes_2_id).unwrap();
897        let wrapped_key_1_3 = ctx
898            .wrap_symmetric_key(key_aes_1_id, key_xchacha_3_id)
899            .unwrap();
900        let wrapped_key_3_1 = ctx
901            .wrap_symmetric_key(key_xchacha_3_id, key_aes_1_id)
902            .unwrap();
903        let wrapped_key_3_4 = ctx
904            .wrap_symmetric_key(key_xchacha_3_id, key_xchacha_4_id)
905            .unwrap();
906
907        // Unwrap the keys
908        let _unwrapped_key_2 = ctx
909            .unwrap_symmetric_key(key_aes_1_id, &wrapped_key_1_2)
910            .unwrap();
911        let _unwrapped_key_3 = ctx
912            .unwrap_symmetric_key(key_aes_1_id, &wrapped_key_1_3)
913            .unwrap();
914        let _unwrapped_key_1 = ctx
915            .unwrap_symmetric_key(key_xchacha_3_id, &wrapped_key_3_1)
916            .unwrap();
917        let _unwrapped_key_4 = ctx
918            .unwrap_symmetric_key(key_xchacha_3_id, &wrapped_key_3_4)
919            .unwrap();
920    }
921
922    #[test]
923    fn test_signing() {
924        let store: KeyStore<TestIds> = KeyStore::default();
925
926        // Generate and insert a key
927        let key_a0_id = TestSigningKey::A(0);
928        let key_a0 = SigningKey::make(SignatureAlgorithm::Ed25519);
929        let verifying_key = key_a0.to_verifying_key();
930        store
931            .context_mut()
932            .set_signing_key(key_a0_id, key_a0)
933            .unwrap();
934
935        assert!(store.context().has_signing_key(key_a0_id));
936
937        // Sign some data with the key
938        #[derive(Serialize, Deserialize)]
939        struct TestData {
940            data: String,
941        }
942        let signed_object = store
943            .context()
944            .sign(
945                key_a0_id,
946                &TestData {
947                    data: "Hello".to_string(),
948                },
949                &SigningNamespace::ExampleNamespace,
950            )
951            .unwrap();
952        let payload: Result<TestData, CryptoError> =
953            signed_object.verify_and_unwrap(&verifying_key, &SigningNamespace::ExampleNamespace);
954        assert!(payload.is_ok());
955
956        let (signature, serialized_message) = store
957            .context()
958            .sign_detached(
959                key_a0_id,
960                &TestData {
961                    data: "Hello".to_string(),
962                },
963                &SigningNamespace::ExampleNamespace,
964            )
965            .unwrap();
966        assert!(signature.verify(
967            serialized_message.as_bytes(),
968            &verifying_key,
969            &SigningNamespace::ExampleNamespace
970        ))
971    }
972
973    #[test]
974    fn test_account_key_rotation() {
975        let store: KeyStore<TestIds> = KeyStore::default();
976        let mut ctx = store.context_mut();
977
978        // Make the keys
979        let current_user_signing_key_id = ctx.make_signing_key(SignatureAlgorithm::Ed25519);
980        let current_user_private_key_id =
981            ctx.make_private_key(PublicKeyEncryptionAlgorithm::RsaOaepSha1);
982
983        // Get the rotated account keys
984        let rotated_keys = ctx
985            .dangerous_get_v2_rotated_account_keys(
986                current_user_private_key_id,
987                current_user_signing_key_id,
988            )
989            .unwrap();
990
991        // Public/Private key
992        assert_eq!(
993            PublicKey::from_der(&rotated_keys.public_key)
994                .unwrap()
995                .to_der()
996                .unwrap(),
997            ctx.get_private_key(current_user_private_key_id)
998                .unwrap()
999                .to_public_key()
1000                .to_der()
1001                .unwrap()
1002        );
1003        let decrypted_private_key: Vec<u8> = rotated_keys
1004            .private_key
1005            .decrypt_with_key(&rotated_keys.user_key)
1006            .unwrap();
1007        let private_key =
1008            PrivateKey::from_der(&Pkcs8PrivateKeyBytes::from(decrypted_private_key)).unwrap();
1009        assert_eq!(
1010            private_key.to_der().unwrap(),
1011            ctx.get_private_key(current_user_private_key_id)
1012                .unwrap()
1013                .to_der()
1014                .unwrap()
1015        );
1016
1017        // Signing Key
1018        let decrypted_signing_key: Vec<u8> = rotated_keys
1019            .signing_key
1020            .decrypt_with_key(&rotated_keys.user_key)
1021            .unwrap();
1022        let signing_key =
1023            SigningKey::from_cose(&CoseKeyBytes::from(decrypted_signing_key)).unwrap();
1024        assert_eq!(
1025            signing_key.to_cose(),
1026            ctx.get_signing_key(current_user_signing_key_id)
1027                .unwrap()
1028                .to_cose(),
1029        );
1030
1031        // Signed Public Key
1032        let signed_public_key = rotated_keys.signed_public_key;
1033        let unwrapped_key = signed_public_key
1034            .verify_and_unwrap(
1035                &ctx.get_signing_key(current_user_signing_key_id)
1036                    .unwrap()
1037                    .to_verifying_key(),
1038            )
1039            .unwrap();
1040        assert_eq!(
1041            unwrapped_key.to_der().unwrap(),
1042            ctx.get_private_key(current_user_private_key_id)
1043                .unwrap()
1044                .to_public_key()
1045                .to_der()
1046                .unwrap()
1047        );
1048    }
1049
1050    #[test]
1051    fn test_encrypt_fails_when_operation_not_allowed() {
1052        use coset::iana::KeyOperation;
1053        let store = KeyStore::<TestIds>::default();
1054        let mut ctx = store.context_mut();
1055        let key_id = TestSymmKey::A(0);
1056        // Key with only Decrypt allowed
1057        let key = SymmetricCryptoKey::XChaCha20Poly1305Key(crate::XChaCha20Poly1305Key {
1058            key_id: [0u8; 16],
1059            enc_key: Box::pin([0u8; 32].into()),
1060            supported_operations: vec![KeyOperation::Decrypt],
1061        });
1062        ctx.set_symmetric_key(key_id, key).unwrap();
1063        let data = DataView("should fail".to_string(), key_id);
1064        let result = data.encrypt_composite(&mut ctx, key_id);
1065        assert!(
1066            matches!(
1067                result,
1068                Err(CryptoError::KeyOperationNotSupported(KeyOperation::Encrypt))
1069            ),
1070            "Expected encrypt to fail with KeyOperationNotSupported",
1071        );
1072    }
1073
1074    #[test]
1075    fn test_move_key() {
1076        let store: KeyStore<TestIds> = KeyStore::default();
1077        let mut ctx = store.context_mut();
1078
1079        // Generate and insert a key
1080        let key = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
1081
1082        assert!(ctx.has_symmetric_key(key));
1083
1084        // Move the key to a new identifier
1085        let new_key_id = TestSymmKey::A(1);
1086        ctx.persist_symmetric_key(key, new_key_id).unwrap();
1087
1088        // Ensure the old key id is gone and the new `one has the key
1089        assert!(!ctx.has_symmetric_key(key));
1090        assert!(ctx.has_symmetric_key(new_key_id));
1091    }
1092}