Skip to main content

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!(
246                    "Unsupported unwrap operation for the given key and data {:?}, {:?}",
247                    wrapping_key,
248                    wrapped_key
249                );
250                return Err(CryptoError::InvalidKey);
251            }
252        };
253
254        let new_key_id = Ids::Symmetric::new_local(LocalId::new());
255
256        #[allow(deprecated)]
257        self.set_symmetric_key(new_key_id, key)?;
258
259        // Returning the new key identifier for convenience
260        Ok(new_key_id)
261    }
262
263    /// Move a symmetric key from a local identifier to a global identifier within the context
264    ///
265    /// The key value is copied to `to` and the original identifier `from` is removed.
266    ///
267    /// # Errors
268    /// Returns an error if the source key does not exist or if setting the destination key
269    /// fails (for example due to read-only global store).
270    pub fn persist_symmetric_key(
271        &mut self,
272        from: Ids::Symmetric,
273        to: Ids::Symmetric,
274    ) -> Result<()> {
275        if !from.is_local() || to.is_local() {
276            return Err(CryptoError::InvalidKeyStoreOperation);
277        }
278        let key = self.get_symmetric_key(from)?.to_owned();
279        self.drop_symmetric_key(from)?;
280        #[allow(deprecated)]
281        self.set_symmetric_key(to, key)?;
282        Ok(())
283    }
284
285    /// Move a private key from a local identifier to a global identifier within the context
286    ///
287    /// The key value is copied to `to` and the original identifier `from` is removed.
288    ///
289    /// # Errors
290    /// Returns an error if the source key does not exist or if setting the destination key
291    /// fails (for example due to read-only global store).
292    pub fn persist_private_key(&mut self, from: Ids::Private, to: Ids::Private) -> Result<()> {
293        if !from.is_local() || to.is_local() {
294            return Err(CryptoError::InvalidKeyStoreOperation);
295        }
296        let key = self.get_private_key(from)?.to_owned();
297        self.drop_private_key(from)?;
298        #[allow(deprecated)]
299        self.set_private_key(to, key)?;
300        Ok(())
301    }
302
303    /// Move a signing key from a local identifier to a global identifier within the context
304    ///
305    /// The key value at `from` will be copied to `to` and the original `from` will be removed.
306    ///
307    /// # Errors
308    /// Returns an error if the source key does not exist or updating the destination fails.
309    pub fn persist_signing_key(&mut self, from: Ids::Signing, to: Ids::Signing) -> Result<()> {
310        if !from.is_local() || to.is_local() {
311            return Err(CryptoError::InvalidKeyStoreOperation);
312        }
313        let key = self.get_signing_key(from)?.to_owned();
314        self.drop_signing_key(from)?;
315        #[allow(deprecated)]
316        self.set_signing_key(to, key)?;
317        Ok(())
318    }
319
320    /// Wrap (encrypt) a signing key with a symmetric key.
321    ///
322    /// The signing key identified by `key_to_wrap` will be serialized to COSE and encrypted
323    /// with the symmetric `wrapping_key`, returning an `EncString` suitable for storage or
324    /// transport.
325    ///
326    /// # Errors
327    /// Returns an error if either key id does not exist or the encryption fails.
328    pub fn wrap_signing_key(
329        &self,
330        wrapping_key: Ids::Symmetric,
331        key_to_wrap: Ids::Signing,
332    ) -> Result<EncString> {
333        let wrapping_key = self.get_symmetric_key(wrapping_key)?;
334        let signing_key = self.get_signing_key(key_to_wrap)?.to_owned();
335        signing_key.to_cose().encrypt_with_key(wrapping_key)
336    }
337
338    /// Wrap (encrypt) a private key with a symmetric key.
339    ///
340    /// The private key identified by `key_to_wrap` will be serialized to DER (PKCS#8) and
341    /// encrypted with `wrapping_key`, returning an `EncString` suitable for storage.
342    ///
343    /// # Errors
344    /// Returns an error if the keys are missing or serialization/encryption fails.
345    pub fn wrap_private_key(
346        &self,
347        wrapping_key: Ids::Symmetric,
348        key_to_wrap: Ids::Private,
349    ) -> Result<EncString> {
350        let wrapping_key = self.get_symmetric_key(wrapping_key)?;
351        let private_key = self.get_private_key(key_to_wrap)?.to_owned();
352        private_key.to_der()?.encrypt_with_key(wrapping_key)
353    }
354
355    /// Decrypt and import a previously wrapped private key into the context.
356    ///
357    /// The `wrapped_key` will be decrypted using `wrapping_key` and parsed as a PKCS#8
358    /// private key; the resulting key will be inserted as a local private key and the
359    /// new local identifier returned.
360    ///
361    /// # Errors
362    /// Returns an error if decryption or parsing fails.
363    #[instrument(skip(self, wrapped_key), err)]
364    pub fn unwrap_private_key(
365        &mut self,
366        wrapping_key: Ids::Symmetric,
367        wrapped_key: &EncString,
368    ) -> Result<Ids::Private> {
369        let wrapping_key = self.get_symmetric_key(wrapping_key)?;
370        let private_key_bytes: Vec<u8> = wrapped_key.decrypt_with_key(wrapping_key)?;
371        let private_key = PrivateKey::from_der(&Pkcs8PrivateKeyBytes::from(private_key_bytes))?;
372        Ok(self.add_local_private_key(private_key))
373    }
374
375    /// Decrypt and import a previously wrapped signing key into the context.
376    ///
377    /// The wrapped COSE key will be decrypted with `wrapping_key` and parsed into a
378    /// `SigningKey` which is inserted as a local signing key. The new local identifier
379    /// is returned.
380    ///
381    /// # Errors
382    /// Returns an error if decryption or parsing fails.
383    pub fn unwrap_signing_key(
384        &mut self,
385        wrapping_key: Ids::Symmetric,
386        wrapped_key: &EncString,
387    ) -> Result<Ids::Signing> {
388        let wrapping_key = self.get_symmetric_key(wrapping_key)?;
389        let signing_key_bytes: Vec<u8> = wrapped_key.decrypt_with_key(wrapping_key)?;
390        let signing_key = SigningKey::from_cose(&CoseKeyBytes::from(signing_key_bytes))?;
391        Ok(self.add_local_signing_key(signing_key))
392    }
393
394    /// Return the verifying (public) key corresponding to a signing key identifier.
395    ///
396    /// This converts the stored `SigningKey` into a `VerifyingKey` suitable for
397    /// signature verification operations.
398    ///
399    /// # Errors
400    /// Returns an error if the signing key id does not exist.
401    pub fn get_verifying_key(&self, signing_key_id: Ids::Signing) -> Result<VerifyingKey> {
402        let signing_key = self.get_signing_key(signing_key_id)?;
403        Ok(signing_key.to_verifying_key())
404    }
405
406    /// Return the public key corresponding to an private key identifier.
407    ///
408    /// This converts the stored private key into its public key representation.
409    ///
410    /// # Errors
411    /// Returns an error if the private key id does not exist.
412    pub fn get_public_key(&self, private_key_id: Ids::Private) -> Result<PublicKey> {
413        let private_key = self.get_private_key(private_key_id)?;
414        Ok(private_key.to_public_key())
415    }
416
417    /// Encrypt and return a symmetric key from the context by using an already existing symmetric
418    /// key
419    ///
420    /// # Arguments
421    ///
422    /// * `wrapping_key` - The key id used to wrap (encrypt) the `key_to_wrap`. It must already
423    ///   exist in the context
424    /// * `key_to_wrap` - The key id to wrap. It must already exist in the context
425    pub fn wrap_symmetric_key(
426        &self,
427        wrapping_key: Ids::Symmetric,
428        key_to_wrap: Ids::Symmetric,
429    ) -> Result<EncString> {
430        use SymmetricCryptoKey::*;
431
432        let wrapping_key_instance = self.get_symmetric_key(wrapping_key)?;
433        let key_to_wrap_instance = self.get_symmetric_key(key_to_wrap)?;
434        // `Aes256CbcHmacKey` can wrap keys by encrypting their byte serialization obtained using
435        // `SymmetricCryptoKey::to_encoded()`. `XChaCha20Poly1305Key` need to specify the
436        // content format to be either octet stream, in case the wrapped key is a Aes256CbcHmacKey
437        // or `Aes256CbcKey`, or by specifying the content format to be CoseKey, in case the
438        // wrapped key is a `XChaCha20Poly1305Key`.
439        match (wrapping_key_instance, key_to_wrap_instance) {
440            (
441                Aes256CbcHmacKey(_),
442                Aes256CbcHmacKey(_) | Aes256CbcKey(_) | XChaCha20Poly1305Key(_),
443            ) => self.encrypt_data_with_symmetric_key(
444                wrapping_key,
445                key_to_wrap_instance
446                    .to_encoded()
447                    .as_ref()
448                    .to_vec()
449                    .as_slice(),
450                ContentFormat::BitwardenLegacyKey,
451            ),
452            (XChaCha20Poly1305Key(_), _) => {
453                let encoded = key_to_wrap_instance.to_encoded_raw();
454                let content_format = encoded.content_format();
455                self.encrypt_data_with_symmetric_key(
456                    wrapping_key,
457                    Into::<Vec<u8>>::into(encoded).as_slice(),
458                    content_format,
459                )
460            }
461            _ => Err(CryptoError::OperationNotSupported(
462                UnsupportedOperationError::EncryptionNotImplementedForKey,
463            )),
464        }
465    }
466
467    /// Returns `true` if the context has a symmetric key with the given identifier
468    pub fn has_symmetric_key(&self, key_id: Ids::Symmetric) -> bool {
469        self.get_symmetric_key(key_id).is_ok()
470    }
471
472    /// Returns `true` if the context has a private key with the given identifier
473    pub fn has_private_key(&self, key_id: Ids::Private) -> bool {
474        self.get_private_key(key_id).is_ok()
475    }
476
477    /// Returns `true` if the context has a signing key with the given identifier
478    pub fn has_signing_key(&self, key_id: Ids::Signing) -> bool {
479        self.get_signing_key(key_id).is_ok()
480    }
481
482    /// Generate a new random symmetric key and store it in the context
483    pub fn generate_symmetric_key(&mut self) -> Ids::Symmetric {
484        self.add_local_symmetric_key(SymmetricCryptoKey::make_aes256_cbc_hmac_key())
485    }
486
487    /// Generate a new symmetric encryption key using the specified algorithm and store it in the
488    /// context as a local key
489    pub fn make_symmetric_key(&mut self, algorithm: SymmetricKeyAlgorithm) -> Ids::Symmetric {
490        self.add_local_symmetric_key(SymmetricCryptoKey::make(algorithm))
491    }
492
493    /// Makes a new private encryption key using the current default algorithm, and stores it in
494    /// the context as a local key
495    pub fn make_private_key(&mut self, algorithm: PublicKeyEncryptionAlgorithm) -> Ids::Private {
496        self.add_local_private_key(PrivateKey::make(algorithm))
497    }
498
499    /// Makes a new signing key using the current default algorithm, and stores it in the context as
500    /// a local key
501    pub fn make_signing_key(&mut self, algorithm: SignatureAlgorithm) -> Ids::Signing {
502        self.add_local_signing_key(SigningKey::make(algorithm))
503    }
504
505    /// Derive a shareable key using hkdf from secret and name and store it in the context.
506    ///
507    /// A specialized variant of this function was called `CryptoService.makeSendKey` in the
508    /// Bitwarden `clients` repository.
509    pub fn derive_shareable_key(
510        &mut self,
511        secret: Zeroizing<[u8; 16]>,
512        name: &str,
513        info: Option<&str>,
514    ) -> Result<Ids::Symmetric> {
515        let key_id = Ids::Symmetric::new_local(LocalId::new());
516        #[allow(deprecated)]
517        self.set_symmetric_key(
518            key_id,
519            SymmetricCryptoKey::Aes256CbcHmacKey(derive_shareable_key(secret, name, info)),
520        )?;
521        Ok(key_id)
522    }
523
524    /// Return a reference to a symmetric key stored in the context.
525    ///
526    /// Deprecated: intended only for internal use and tests. This exposes the underlying
527    /// `SymmetricCryptoKey` reference directly and should not be used by external code. Use
528    /// the higher-level APIs (for example encryption/decryption helpers) or `get_symmetric_key`
529    /// internally when possible.
530    ///
531    /// # Errors
532    /// Returns [`CryptoError::MissingKeyId`] if the key id does not exist in the context.
533    #[deprecated(note = "This function should ideally never be used outside this crate")]
534    pub fn dangerous_get_symmetric_key(
535        &self,
536        key_id: Ids::Symmetric,
537    ) -> Result<&SymmetricCryptoKey> {
538        self.get_symmetric_key(key_id)
539    }
540
541    /// Return a reference to a signing key stored in the context.
542    ///
543    /// Deprecated: intended only for internal use and tests. This exposes the underlying
544    /// `SigningKey` reference directly and should not be used by external code. Use the
545    /// higher-level APIs (for example signing helpers) or `get_signing_key` internally when
546    /// possible
547    ///
548    /// # Errors
549    /// Returns [`CryptoError::MissingKeyId`] if the key id does not exist in
550    /// the context.
551    #[deprecated(note = "This function should ideally never be used outside this crate")]
552    pub fn dangerous_get_signing_key(&self, key_id: Ids::Signing) -> Result<&SigningKey> {
553        self.get_signing_key(key_id)
554    }
555
556    /// Return a reference to an asymmetric (private) key stored in the context.
557    ///
558    /// Deprecated: intended only for internal use and tests. This exposes the underlying
559    /// `PrivateKey` reference directly and should not be used by external code. Prefer
560    /// using the public key via `get_public_key` or other higher-level APIs instead.
561    ///
562    /// # Errors
563    /// Returns [`CryptoError::MissingKeyId`] if the key id does not exist in the context.
564    #[deprecated(note = "This function should ideally never be used outside this crate")]
565    pub fn dangerous_get_private_key(&self, key_id: Ids::Private) -> Result<&PrivateKey> {
566        self.get_private_key(key_id)
567    }
568
569    /// Makes a signed public key from a private key and signing key stored in context.
570    /// Signing a public key asserts ownership, and makes the claim to other users that if they want
571    /// to share with you, they can use this public key.
572    pub fn make_signed_public_key(
573        &self,
574        private_key_id: Ids::Private,
575        signing_key_id: Ids::Signing,
576    ) -> Result<SignedPublicKey> {
577        let public_key = self.get_private_key(private_key_id)?.to_public_key();
578        let signing_key = self.get_signing_key(signing_key_id)?;
579        let signed_public_key =
580            SignedPublicKeyMessage::from_public_key(&public_key)?.sign(signing_key)?;
581        Ok(signed_public_key)
582    }
583
584    pub(crate) fn get_symmetric_key(&self, key_id: Ids::Symmetric) -> Result<&SymmetricCryptoKey> {
585        if key_id.is_local() {
586            self.local_symmetric_keys.get(key_id)
587        } else {
588            self.global_keys.get().symmetric_keys.get(key_id)
589        }
590        .ok_or_else(|| crate::CryptoError::MissingKeyId(format!("{key_id:?}")))
591    }
592
593    pub(super) fn get_private_key(&self, key_id: Ids::Private) -> Result<&PrivateKey> {
594        if key_id.is_local() {
595            self.local_private_keys.get(key_id)
596        } else {
597            self.global_keys.get().private_keys.get(key_id)
598        }
599        .ok_or_else(|| crate::CryptoError::MissingKeyId(format!("{key_id:?}")))
600    }
601
602    pub(super) fn get_signing_key(&self, key_id: Ids::Signing) -> Result<&SigningKey> {
603        if key_id.is_local() {
604            self.local_signing_keys.get(key_id)
605        } else {
606            self.global_keys.get().signing_keys.get(key_id)
607        }
608        .ok_or_else(|| crate::CryptoError::MissingKeyId(format!("{key_id:?}")))
609    }
610
611    /// Set a symmetric key in the context.
612    ///
613    /// # Errors
614    /// Returns [`CryptoError::ReadOnlyKeyStore`] if the context does not have write access when
615    /// attempting to modify the global store.
616    #[deprecated(note = "This function should ideally never be used outside this crate")]
617    pub fn set_symmetric_key(
618        &mut self,
619        key_id: Ids::Symmetric,
620        key: SymmetricCryptoKey,
621    ) -> Result<()> {
622        self.set_symmetric_key_internal(key_id, key)
623    }
624
625    pub(crate) fn set_symmetric_key_internal(
626        &mut self,
627        key_id: Ids::Symmetric,
628        key: SymmetricCryptoKey,
629    ) -> Result<()> {
630        if key_id.is_local() {
631            self.local_symmetric_keys.upsert(key_id, key);
632        } else {
633            self.global_keys
634                .get_mut()?
635                .symmetric_keys
636                .upsert(key_id, key);
637        }
638        Ok(())
639    }
640
641    /// Add a new symmetric key to the local context, returning a new unique identifier for it.
642    pub fn add_local_symmetric_key(&mut self, key: SymmetricCryptoKey) -> Ids::Symmetric {
643        let key_id = Ids::Symmetric::new_local(LocalId::new());
644        self.local_symmetric_keys.upsert(key_id, key);
645        key_id
646    }
647
648    /// Get the type of a symmetric key stored in the context.
649    pub fn get_symmetric_key_algorithm(
650        &self,
651        key_id: Ids::Symmetric,
652    ) -> Result<SymmetricKeyAlgorithm> {
653        let key = self.get_symmetric_key(key_id)?;
654        match key {
655            // Note this is dropped soon
656            SymmetricCryptoKey::Aes256CbcKey(_) => Err(CryptoError::OperationNotSupported(
657                UnsupportedOperationError::EncryptionNotImplementedForKey,
658            )),
659            SymmetricCryptoKey::Aes256CbcHmacKey(_) => Ok(SymmetricKeyAlgorithm::Aes256CbcHmac),
660            SymmetricCryptoKey::XChaCha20Poly1305Key(_) => {
661                Ok(SymmetricKeyAlgorithm::XChaCha20Poly1305)
662            }
663        }
664    }
665
666    /// Set a private key in the context.
667    ///
668    /// # Errors
669    /// Returns [`CryptoError::ReadOnlyKeyStore`] if attempting to write to the global store when
670    /// the context is read-only.
671    #[deprecated(note = "This function should ideally never be used outside this crate")]
672    pub fn set_private_key(&mut self, key_id: Ids::Private, key: PrivateKey) -> Result<()> {
673        if key_id.is_local() {
674            self.local_private_keys.upsert(key_id, key);
675        } else {
676            self.global_keys.get_mut()?.private_keys.upsert(key_id, key);
677        }
678        Ok(())
679    }
680
681    /// Add a new private key to the local context, returning a new unique identifier for it.
682    pub fn add_local_private_key(&mut self, key: PrivateKey) -> Ids::Private {
683        let key_id = Ids::Private::new_local(LocalId::new());
684        self.local_private_keys.upsert(key_id, key);
685        key_id
686    }
687
688    /// Sets a signing key in the context
689    ///
690    /// # Errors
691    /// Returns [`CryptoError::ReadOnlyKeyStore`] if attempting to write to the global store when
692    /// the context is read-only.
693    #[deprecated(note = "This function should ideally never be used outside this crate")]
694    pub fn set_signing_key(&mut self, key_id: Ids::Signing, key: SigningKey) -> Result<()> {
695        if key_id.is_local() {
696            self.local_signing_keys.upsert(key_id, key);
697        } else {
698            self.global_keys.get_mut()?.signing_keys.upsert(key_id, key);
699        }
700        Ok(())
701    }
702
703    /// Add a new signing key to the local context, returning a new unique identifier for it.
704    pub fn add_local_signing_key(&mut self, key: SigningKey) -> Ids::Signing {
705        let key_id = Ids::Signing::new_local(LocalId::new());
706        self.local_signing_keys.upsert(key_id, key);
707        key_id
708    }
709
710    #[instrument(skip(self, data), err)]
711    pub(crate) fn decrypt_data_with_symmetric_key(
712        &self,
713        key: Ids::Symmetric,
714        data: &EncString,
715    ) -> Result<Vec<u8>> {
716        let key = self.get_symmetric_key(key)?;
717
718        match (data, key) {
719            (EncString::Aes256Cbc_B64 { iv, data }, SymmetricCryptoKey::Aes256CbcKey(key)) => {
720                crate::aes::decrypt_aes256(iv, data.clone(), &key.enc_key)
721                    .map_err(|_| CryptoError::Decrypt)
722            }
723            (
724                EncString::Aes256Cbc_HmacSha256_B64 { iv, mac, data },
725                SymmetricCryptoKey::Aes256CbcHmacKey(key),
726            ) => crate::aes::decrypt_aes256_hmac(iv, mac, data.clone(), &key.mac_key, &key.enc_key)
727                .map_err(|_| CryptoError::Decrypt),
728            (
729                EncString::Cose_Encrypt0_B64 { data },
730                SymmetricCryptoKey::XChaCha20Poly1305Key(key),
731            ) => {
732                let (data, _) = crate::cose::decrypt_xchacha20_poly1305(
733                    &CoseEncrypt0Bytes::from(data.clone()),
734                    key,
735                )?;
736                Ok(data)
737            }
738            _ => {
739                tracing::warn!("Unsupported decryption operation for the given key and data");
740                Err(CryptoError::InvalidKey)
741            }
742        }
743    }
744
745    pub(crate) fn encrypt_data_with_symmetric_key(
746        &self,
747        key: Ids::Symmetric,
748        data: &[u8],
749        content_format: ContentFormat,
750    ) -> Result<EncString> {
751        let key = self.get_symmetric_key(key)?;
752        match key {
753            SymmetricCryptoKey::Aes256CbcKey(_) => Err(CryptoError::OperationNotSupported(
754                UnsupportedOperationError::EncryptionNotImplementedForKey,
755            )),
756            SymmetricCryptoKey::Aes256CbcHmacKey(key) => EncString::encrypt_aes256_hmac(data, key),
757            SymmetricCryptoKey::XChaCha20Poly1305Key(key) => {
758                if !key.supported_operations.contains(&KeyOperation::Encrypt) {
759                    return Err(CryptoError::KeyOperationNotSupported(KeyOperation::Encrypt));
760                }
761                EncString::encrypt_xchacha20_poly1305(data, key, content_format)
762            }
763        }
764    }
765
766    /// Signs the given data using the specified signing key, for the given
767    /// [crate::SigningNamespace] and returns the signature and the serialized message. See
768    /// [crate::SigningKey::sign]
769    pub fn sign<Message: Serialize>(
770        &self,
771        key: Ids::Signing,
772        message: &Message,
773        namespace: &crate::SigningNamespace,
774    ) -> Result<SignedObject> {
775        self.get_signing_key(key)?.sign(message, namespace)
776    }
777
778    /// Signs the given data using the specified signing key, for the given
779    /// [crate::SigningNamespace] and returns the signature and the serialized message. See
780    /// [crate::SigningKey::sign_detached]
781    #[allow(unused)]
782    pub(crate) fn sign_detached<Message: Serialize>(
783        &self,
784        key: Ids::Signing,
785        message: &Message,
786        namespace: &crate::SigningNamespace,
787    ) -> Result<(Signature, signing::SerializedMessage)> {
788        self.get_signing_key(key)?.sign_detached(message, namespace)
789    }
790
791    /// Re-encrypts the user's keys with the provided symmetric key for a v2 user.
792    pub fn dangerous_get_v2_rotated_account_keys(
793        &self,
794        current_user_private_key_id: Ids::Private,
795        current_user_signing_key_id: Ids::Signing,
796    ) -> Result<RotatedUserKeys> {
797        #[expect(deprecated)]
798        crate::dangerous_get_v2_rotated_account_keys(
799            current_user_private_key_id,
800            current_user_signing_key_id,
801            self,
802        )
803    }
804}
805
806#[cfg(test)]
807#[allow(deprecated)]
808mod tests {
809    use serde::{Deserialize, Serialize};
810
811    use crate::{
812        CompositeEncryptable, CoseKeyBytes, CoseSerializable, CryptoError, Decryptable,
813        KeyDecryptable, Pkcs8PrivateKeyBytes, PrivateKey, PublicKey, PublicKeyEncryptionAlgorithm,
814        SignatureAlgorithm, SigningKey, SigningNamespace, SymmetricCryptoKey,
815        SymmetricKeyAlgorithm,
816        store::{
817            KeyStore,
818            tests::{Data, DataView},
819        },
820        traits::tests::{TestIds, TestSigningKey, TestSymmKey},
821    };
822
823    #[test]
824    fn test_set_signing_key() {
825        let store: KeyStore<TestIds> = KeyStore::default();
826
827        // Generate and insert a key
828        let key_a0_id = TestSigningKey::A(0);
829        let key_a0 = SigningKey::make(SignatureAlgorithm::Ed25519);
830        store
831            .context_mut()
832            .set_signing_key(key_a0_id, key_a0)
833            .unwrap();
834    }
835
836    #[test]
837    fn test_set_keys_for_encryption() {
838        let store: KeyStore<TestIds> = KeyStore::default();
839
840        // Generate and insert a key
841        let key_a0_id = TestSymmKey::A(0);
842        let mut ctx = store.context_mut();
843        let local_key_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
844        ctx.persist_symmetric_key(local_key_id, TestSymmKey::A(0))
845            .unwrap();
846
847        assert!(ctx.has_symmetric_key(key_a0_id));
848
849        // Encrypt some data with the key
850        let data = DataView("Hello, World!".to_string(), key_a0_id);
851        let _encrypted: Data = data.encrypt_composite(&mut ctx, key_a0_id).unwrap();
852    }
853
854    #[test]
855    fn test_key_encryption() {
856        let store: KeyStore<TestIds> = KeyStore::default();
857
858        let mut ctx = store.context();
859
860        // Generate and insert a key
861        let key_1_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
862
863        assert!(ctx.has_symmetric_key(key_1_id));
864
865        // Generate and insert a new key
866        let key_2_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
867
868        assert!(ctx.has_symmetric_key(key_2_id));
869
870        // Encrypt the new key with the old key
871        let key_2_enc = ctx.wrap_symmetric_key(key_1_id, key_2_id).unwrap();
872
873        // Decrypt the new key with the old key in a different identifier
874        let new_key_id = ctx.unwrap_symmetric_key(key_1_id, &key_2_enc).unwrap();
875
876        // Now `key_2_id` and `new_key_id` contain the same key, so we should be able to encrypt
877        // with one and decrypt with the other
878
879        let data = DataView("Hello, World!".to_string(), key_2_id);
880        let encrypted = data.encrypt_composite(&mut ctx, key_2_id).unwrap();
881
882        let decrypted1 = encrypted.decrypt(&mut ctx, key_2_id).unwrap();
883        let decrypted2 = encrypted.decrypt(&mut ctx, new_key_id).unwrap();
884
885        // Assert that the decrypted data is the same
886        assert_eq!(decrypted1.0, decrypted2.0);
887    }
888
889    #[test]
890    fn test_wrap_unwrap() {
891        let store: KeyStore<TestIds> = KeyStore::default();
892        let mut ctx = store.context_mut();
893
894        // Aes256 CBC HMAC keys
895        let key_aes_1_id = TestSymmKey::A(1);
896        let local_key_1_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
897        ctx.persist_symmetric_key(local_key_1_id, key_aes_1_id)
898            .unwrap();
899        let key_aes_2_id = TestSymmKey::A(2);
900        let local_key_2_id = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
901        ctx.persist_symmetric_key(local_key_2_id, key_aes_2_id)
902            .unwrap();
903
904        // XChaCha20 Poly1305 keys
905        let key_xchacha_3_id = TestSymmKey::A(3);
906        let key_xchacha_3 = SymmetricCryptoKey::make_xchacha20_poly1305_key();
907        ctx.set_symmetric_key(key_xchacha_3_id, key_xchacha_3.clone())
908            .unwrap();
909        let key_xchacha_4_id = TestSymmKey::A(4);
910        let key_xchacha_4 = SymmetricCryptoKey::make_xchacha20_poly1305_key();
911        ctx.set_symmetric_key(key_xchacha_4_id, key_xchacha_4.clone())
912            .unwrap();
913
914        // Wrap and unwrap the keys
915        let wrapped_key_1_2 = ctx.wrap_symmetric_key(key_aes_1_id, key_aes_2_id).unwrap();
916        let wrapped_key_1_3 = ctx
917            .wrap_symmetric_key(key_aes_1_id, key_xchacha_3_id)
918            .unwrap();
919        let wrapped_key_3_1 = ctx
920            .wrap_symmetric_key(key_xchacha_3_id, key_aes_1_id)
921            .unwrap();
922        let wrapped_key_3_4 = ctx
923            .wrap_symmetric_key(key_xchacha_3_id, key_xchacha_4_id)
924            .unwrap();
925
926        // Unwrap the keys
927        let _unwrapped_key_2 = ctx
928            .unwrap_symmetric_key(key_aes_1_id, &wrapped_key_1_2)
929            .unwrap();
930        let _unwrapped_key_3 = ctx
931            .unwrap_symmetric_key(key_aes_1_id, &wrapped_key_1_3)
932            .unwrap();
933        let _unwrapped_key_1 = ctx
934            .unwrap_symmetric_key(key_xchacha_3_id, &wrapped_key_3_1)
935            .unwrap();
936        let _unwrapped_key_4 = ctx
937            .unwrap_symmetric_key(key_xchacha_3_id, &wrapped_key_3_4)
938            .unwrap();
939    }
940
941    #[test]
942    fn test_signing() {
943        let store: KeyStore<TestIds> = KeyStore::default();
944
945        // Generate and insert a key
946        let key_a0_id = TestSigningKey::A(0);
947        let key_a0 = SigningKey::make(SignatureAlgorithm::Ed25519);
948        let verifying_key = key_a0.to_verifying_key();
949        store
950            .context_mut()
951            .set_signing_key(key_a0_id, key_a0)
952            .unwrap();
953
954        assert!(store.context().has_signing_key(key_a0_id));
955
956        // Sign some data with the key
957        #[derive(Serialize, Deserialize)]
958        struct TestData {
959            data: String,
960        }
961        let signed_object = store
962            .context()
963            .sign(
964                key_a0_id,
965                &TestData {
966                    data: "Hello".to_string(),
967                },
968                &SigningNamespace::ExampleNamespace,
969            )
970            .unwrap();
971        let payload: Result<TestData, CryptoError> =
972            signed_object.verify_and_unwrap(&verifying_key, &SigningNamespace::ExampleNamespace);
973        assert!(payload.is_ok());
974
975        let (signature, serialized_message) = store
976            .context()
977            .sign_detached(
978                key_a0_id,
979                &TestData {
980                    data: "Hello".to_string(),
981                },
982                &SigningNamespace::ExampleNamespace,
983            )
984            .unwrap();
985        assert!(signature.verify(
986            serialized_message.as_bytes(),
987            &verifying_key,
988            &SigningNamespace::ExampleNamespace
989        ))
990    }
991
992    #[test]
993    fn test_account_key_rotation() {
994        let store: KeyStore<TestIds> = KeyStore::default();
995        let mut ctx = store.context_mut();
996
997        // Make the keys
998        let current_user_signing_key_id = ctx.make_signing_key(SignatureAlgorithm::Ed25519);
999        let current_user_private_key_id =
1000            ctx.make_private_key(PublicKeyEncryptionAlgorithm::RsaOaepSha1);
1001
1002        // Get the rotated account keys
1003        let rotated_keys = ctx
1004            .dangerous_get_v2_rotated_account_keys(
1005                current_user_private_key_id,
1006                current_user_signing_key_id,
1007            )
1008            .unwrap();
1009
1010        // Public/Private key
1011        assert_eq!(
1012            PublicKey::from_der(&rotated_keys.public_key)
1013                .unwrap()
1014                .to_der()
1015                .unwrap(),
1016            ctx.get_private_key(current_user_private_key_id)
1017                .unwrap()
1018                .to_public_key()
1019                .to_der()
1020                .unwrap()
1021        );
1022        let decrypted_private_key: Vec<u8> = rotated_keys
1023            .private_key
1024            .decrypt_with_key(&rotated_keys.user_key)
1025            .unwrap();
1026        let private_key =
1027            PrivateKey::from_der(&Pkcs8PrivateKeyBytes::from(decrypted_private_key)).unwrap();
1028        assert_eq!(
1029            private_key.to_der().unwrap(),
1030            ctx.get_private_key(current_user_private_key_id)
1031                .unwrap()
1032                .to_der()
1033                .unwrap()
1034        );
1035
1036        // Signing Key
1037        let decrypted_signing_key: Vec<u8> = rotated_keys
1038            .signing_key
1039            .decrypt_with_key(&rotated_keys.user_key)
1040            .unwrap();
1041        let signing_key =
1042            SigningKey::from_cose(&CoseKeyBytes::from(decrypted_signing_key)).unwrap();
1043        assert_eq!(
1044            signing_key.to_cose(),
1045            ctx.get_signing_key(current_user_signing_key_id)
1046                .unwrap()
1047                .to_cose(),
1048        );
1049
1050        // Signed Public Key
1051        let signed_public_key = rotated_keys.signed_public_key;
1052        let unwrapped_key = signed_public_key
1053            .verify_and_unwrap(
1054                &ctx.get_signing_key(current_user_signing_key_id)
1055                    .unwrap()
1056                    .to_verifying_key(),
1057            )
1058            .unwrap();
1059        assert_eq!(
1060            unwrapped_key.to_der().unwrap(),
1061            ctx.get_private_key(current_user_private_key_id)
1062                .unwrap()
1063                .to_public_key()
1064                .to_der()
1065                .unwrap()
1066        );
1067    }
1068
1069    #[test]
1070    fn test_encrypt_fails_when_operation_not_allowed() {
1071        use coset::iana::KeyOperation;
1072        let store = KeyStore::<TestIds>::default();
1073        let mut ctx = store.context_mut();
1074        let key_id = TestSymmKey::A(0);
1075        // Key with only Decrypt allowed
1076        let key = SymmetricCryptoKey::XChaCha20Poly1305Key(crate::XChaCha20Poly1305Key {
1077            key_id: [0u8; 16].into(),
1078            enc_key: Box::pin([0u8; 32].into()),
1079            supported_operations: vec![KeyOperation::Decrypt],
1080        });
1081        ctx.set_symmetric_key(key_id, key).unwrap();
1082        let data = DataView("should fail".to_string(), key_id);
1083        let result = data.encrypt_composite(&mut ctx, key_id);
1084        assert!(
1085            matches!(
1086                result,
1087                Err(CryptoError::KeyOperationNotSupported(KeyOperation::Encrypt))
1088            ),
1089            "Expected encrypt to fail with KeyOperationNotSupported",
1090        );
1091    }
1092
1093    #[test]
1094    fn test_move_key() {
1095        let store: KeyStore<TestIds> = KeyStore::default();
1096        let mut ctx = store.context_mut();
1097
1098        // Generate and insert a key
1099        let key = ctx.make_symmetric_key(SymmetricKeyAlgorithm::Aes256CbcHmac);
1100
1101        assert!(ctx.has_symmetric_key(key));
1102
1103        // Move the key to a new identifier
1104        let new_key_id = TestSymmKey::A(1);
1105        ctx.persist_symmetric_key(key, new_key_id).unwrap();
1106
1107        // Ensure the old key id is gone and the new `one has the key
1108        assert!(!ctx.has_symmetric_key(key));
1109        assert!(ctx.has_symmetric_key(new_key_id));
1110    }
1111}