Add key id to passwordprotectedkeyenvelope headers

crates/bitwarden-crypto/src/cose.rs

@@ -40,6 +40,7 @@ const CONTENT_TYPE_SPKI_PUBLIC_KEY: &str = "application/x.bitwarden.spki-public-
 //
 /// The label used for the namespace ensuring strong domain separation when using signatures.
 pub(crate) const SIGNING_NAMESPACE: i64 = -80000;
+pub(crate) const CONTAINED_KEY_ID: i64 = -80001;
 
 /// Encrypts a plaintext message using XChaCha20Poly1305 and returns a COSE Encrypt0 message
 pub(crate) fn encrypt_xchacha20_poly1305(

crates/bitwarden-crypto/src/keys/key_id.rs

@@ -9,7 +9,7 @@ pub(crate) const KEY_ID_SIZE: usize = 16;
 /// A key id is a unique identifier for a single key. There is a 1:1 mapping between key ID and key
 /// bytes, so something like a user key rotation is replacing the key with ID A with a new key with
 /// ID B.
-#[derive(Clone)]
+#[derive(Clone, PartialEq)]
 pub(crate) struct KeyId(Uuid);
 
 /// Fixed length identifiers for keys.

crates/bitwarden-crypto/src/keys/symmetric_crypto_key.rs

@@ -230,6 +230,16 @@ impl SymmetricCryptoKey {
     pub fn to_base64(&self) -> B64 {
         B64::from(self.to_encoded().as_ref())
     }
+
+    /// Returns the key ID of the key, if it has one. Only
+    /// [SymmetricCryptoKey::XChaCha20Poly1305Key] has a key ID.
+    pub(crate) fn key_id(&self) -> Option<KeyId> {
+        match self {
+            Self::Aes256CbcKey(_) => None,
+            Self::Aes256CbcHmacKey(_) => None,
+            Self::XChaCha20Poly1305Key(key) => Some(KeyId::from(key.key_id)),
+        }
+    }
 }
 
 impl ConstantTimeEq for SymmetricCryptoKey {

crates/bitwarden-crypto/src/safe/password_protected_key_envelope.rs

@@ -28,8 +28,9 @@ use crate::{
     KeyStoreContext, SymmetricCryptoKey,
     cose::{
         ALG_ARGON2ID13, ARGON2_ITERATIONS, ARGON2_MEMORY, ARGON2_PARALLELISM, ARGON2_SALT,
-        CoseExtractError, extract_bytes, extract_integer,
+        CONTAINED_KEY_ID, CoseExtractError, extract_bytes, extract_integer,
     },
+    keys::KeyId,
     xchacha20,
 };
 
@@ -120,7 +121,13 @@ impl<Ids: KeyIds> PasswordProtectedKeyEnvelope<Ids> {
                 recipient.protected.header.alg = Some(coset::Algorithm::PrivateUse(ALG_ARGON2ID13));
                 recipient
             })
-            .protected(HeaderBuilder::from(content_format).build())
+            .protected({
+                let mut hdr = HeaderBuilder::from(content_format);
+                if let Some(key_id) = key_to_seal.key_id() {
+                    hdr = hdr.value(CONTAINED_KEY_ID, Value::from(Vec::from(&key_id)));
+                }
+                hdr.build()
+            })
             .create_ciphertext(&key_to_seal_bytes, &[], |data, aad| {
                 let ciphertext = xchacha20::encrypt_xchacha20_poly1305(&envelope_key, data, aad);
                 nonce.copy_from_slice(&ciphertext.nonce());
@@ -227,6 +234,23 @@ impl<Ids: KeyIds> PasswordProtectedKeyEnvelope<Ids> {
         let unsealed = self.unseal_ref(password)?;
         Self::seal_ref(&unsealed, new_password)
     }
+
+    /// Get the key ID of the contained key, if the key ID is stored on the envelope headers.
+    /// Only COSE keys have a key ID, legacy keys do not.
+    #[allow(dead_code)]
+    pub(crate) fn contained_key_id(
+        &self,
+    ) -> Result<Option<KeyId>, PasswordProtectedKeyEnvelopeError> {
+        let key_id_bytes = extract_bytes(
+            &self.cose_encrypt.protected.header,
+            CONTAINED_KEY_ID,
+            "key id",
+        )?;
+        let key_id_array: [u8; 16] = key_id_bytes.as_slice().try_into().map_err(|_| {
+            PasswordProtectedKeyEnvelopeError::Parsing("Invalid key id".to_string())
+        })?;
+        Ok(Some(KeyId::from(key_id_array)))
+    }
 }
 
 impl<Ids: KeyIds> From<&PasswordProtectedKeyEnvelope<Ids>> for Vec<u8> {
@@ -642,4 +666,25 @@ mod tests {
             Err(PasswordProtectedKeyEnvelopeError::WrongPassword)
         ));
     }
+
+    #[test]
+    fn test_key_id() {
+        let key_store = KeyStore::<TestIds>::default();
+        let mut ctx: KeyStoreContext<'_, TestIds> = key_store.context_mut();
+        let test_key = ctx.make_cose_symmetric_key(TestSymmKey::A(0)).unwrap();
+        #[allow(deprecated)]
+        let key_id = ctx
+            .dangerous_get_symmetric_key(test_key)
+            .unwrap()
+            .key_id()
+            .unwrap()
+            .clone();
+
+        let password = "test_password";
+
+        // Seal the key with a password
+        let envelope = PasswordProtectedKeyEnvelope::seal(test_key, password, &ctx).unwrap();
+        let contained_key_id = envelope.contained_key_id().unwrap();
+        assert_eq!(Some(key_id), contained_key_id);
+    }
 }