[PM-27638] Add key id to passwordprotectedkeyenvelope headers

crates/bitwarden-crypto/src/cose.rs

@@ -18,17 +18,25 @@ use crate::{
     xchacha20,
 };
 
+// Custom COSE algorithm values
+// NOTE: Any algorithm value below -65536 is reserved for private use in the IANA allocations and
+// can be used freely.
 /// XChaCha20 <https://datatracker.ietf.org/doc/html/draft-irtf-cfrg-xchacha-03> is used over ChaCha20
 /// to be able to randomly generate nonces, and to not have to worry about key wearout. Since
 /// the draft was never published as an RFC, we use a private-use value for the algorithm.
 pub(crate) const XCHACHA20_POLY1305: i64 = -70000;
-const XCHACHA20_TEXT_PAD_BLOCK_SIZE: usize = 32;
-
 pub(crate) const ALG_ARGON2ID13: i64 = -71000;
+
+// Custom labels for COSE headers
+// NOTE: Any label below -65536 is reserved for private use in the IANA allocations and can be used
+// freely.
 pub(crate) const ARGON2_SALT: i64 = -71001;
 pub(crate) const ARGON2_ITERATIONS: i64 = -71002;
 pub(crate) const ARGON2_MEMORY: i64 = -71003;
 pub(crate) const ARGON2_PARALLELISM: i64 = -71004;
+/// Indicates for any object containing a key (wrapped key, password protected key envelope) which
+/// key ID that contained key has
+pub(crate) const CONTAINED_KEY_ID: i64 = -71005;
 
 // Note: These are in the "unregistered" tree: https://datatracker.ietf.org/doc/html/rfc6838#section-3.4
 // These are only used within Bitwarden, and not meant for exchange with other systems.
@@ -37,13 +45,14 @@ pub(crate) const CONTENT_TYPE_PADDED_CBOR: &str = "application/x.bitwarden.cbor-
 const CONTENT_TYPE_BITWARDEN_LEGACY_KEY: &str = "application/x.bitwarden.legacy-key";
 const CONTENT_TYPE_SPKI_PUBLIC_KEY: &str = "application/x.bitwarden.spki-public-key";
 
-// Labels
-//
+/// Namespaces
 /// The label used for the namespace ensuring strong domain separation when using signatures.
 pub(crate) const SIGNING_NAMESPACE: i64 = -80000;
 /// The label used for the namespace ensuring strong domain separation when using data envelopes.
 pub(crate) const DATA_ENVELOPE_NAMESPACE: i64 = -80001;
 
+const XCHACHA20_TEXT_PAD_BLOCK_SIZE: usize = 32;
+
 /// Encrypts a plaintext message using XChaCha20Poly1305 and returns a COSE Encrypt0 message
 pub(crate) fn encrypt_xchacha20_poly1305(
     plaintext: &[u8],

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

@@ -267,6 +267,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

@@ -30,8 +30,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,
 };
 
@@ -121,7 +122,13 @@ impl PasswordProtectedKeyEnvelope {
                 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());
@@ -221,6 +228,28 @@ impl PasswordProtectedKeyEnvelope {
         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",
+        );
+
+        if let Some(bytes) = key_id_bytes.ok() {
+            let key_id_array: [u8; 16] = bytes.as_slice().try_into().map_err(|_| {
+                PasswordProtectedKeyEnvelopeError::Parsing("Invalid key id".to_string())
+            })?;
+            Ok(Some(KeyId::from(key_id_array)))
+        } else {
+            Ok(None)
+        }
+    }
 }
 
 impl From<&PasswordProtectedKeyEnvelope> for Vec<u8> {
@@ -653,4 +682,39 @@ 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);
+    }
+
+    #[test]
+    fn test_no_key_id() {
+        let key_store = KeyStore::<TestIds>::default();
+        let mut ctx: KeyStoreContext<'_, TestIds> = key_store.context_mut();
+        let test_key = ctx.generate_symmetric_key();
+
+        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!(None, contained_key_id);
+    }
 }