From b0bfd83f33c7f19761d4a81263d3ff200e077374 Mon Sep 17 00:00:00 2001
From: Isaiah Inuwa <iinuwa@bitwarden.com>
Date: Mon, 22 Sep 2025 11:05:15 -0500
Subject: [PATCH] Expose method to derive symmetric key from PRF output

---
 .../src/key_management/crypto.rs              | 15 +++++--
 .../src/key_management/crypto_client.rs       | 13 ++++--
 crates/bitwarden-crypto/src/keys/mod.rs       |  3 ++
 crates/bitwarden-crypto/src/keys/prf.rs       | 44 +++++++++++++++++++
 crates/bitwarden-uniffi/src/crypto.rs         |  6 ++-
 5 files changed, 74 insertions(+), 7 deletions(-)
 create mode 100644 crates/bitwarden-crypto/src/keys/prf.rs

diff --git a/crates/bitwarden-core/src/key_management/crypto.rs b/crates/bitwarden-core/src/key_management/crypto.rs
index f65bd0531..ad41fa255 100644
--- a/crates/bitwarden-core/src/key_management/crypto.rs
+++ b/crates/bitwarden-core/src/key_management/crypto.rs
@@ -10,7 +10,8 @@ use bitwarden_crypto::{
     AsymmetricCryptoKey, CoseSerializable, CryptoError, EncString, Kdf, KeyDecryptable,
     KeyEncryptable, MasterKey, Pkcs8PrivateKeyBytes, PrimitiveEncryptable, SignatureAlgorithm,
     SignedPublicKey, SigningKey, SpkiPublicKeyBytes, SymmetricCryptoKey, UnsignedSharedKey,
-    UserKey, dangerous_get_v2_rotated_account_keys, safe::PasswordProtectedKeyEnvelopeError,
+    UserKey, dangerous_get_v2_rotated_account_keys, derive_symmetric_key_from_prf,
+    safe::PasswordProtectedKeyEnvelopeError,
 };
 use bitwarden_encoding::B64;
 use bitwarden_error::bitwarden_error;
@@ -498,6 +499,14 @@ fn derive_pin_protected_user_key(
     Ok(derived_key.encrypt_user_key(user_key)?)
 }
 
+pub(super) fn derive_prf_key(
+    client: &Client,
+    prf: B64,
+) -> Result<RotateableKeySet, CryptoClientError> {
+    let prf_key = derive_symmetric_key_from_prf(prf.as_bytes())?;
+    create_rotateable_key_set(client, prf_key)
+}
+
 #[allow(missing_docs)]
 #[bitwarden_error(flat)]
 #[derive(Debug, thiserror::Error)]
@@ -594,8 +603,8 @@ pub(super) fn make_key_pair(user_key: B64) -> Result<MakeKeyPairResponse, Crypto
 /// and the `PublicKey` is protected by the `UserKey`. This setup allows:
 ///
 ///   - Access to `UserKey` by knowing the `ExternalKey`
-///   - Rotation to a `NewUserKey` by knowing the current `UserKey`,
-///     without needing access to the `ExternalKey`
+///   - Rotation to a `NewUserKey` by knowing the current `UserKey`, without needing access to the
+///     `ExternalKey`
 
 #[derive(Serialize, Deserialize, Debug)]
 #[serde(rename_all = "camelCase", deny_unknown_fields)]
diff --git a/crates/bitwarden-core/src/key_management/crypto_client.rs b/crates/bitwarden-core/src/key_management/crypto_client.rs
index 08a18168f..d3a11e09d 100644
--- a/crates/bitwarden-core/src/key_management/crypto_client.rs
+++ b/crates/bitwarden-core/src/key_management/crypto_client.rs
@@ -26,9 +26,9 @@ use crate::{
     client::encryption_settings::EncryptionSettingsError,
     error::StatefulCryptoError,
     key_management::crypto::{
-        CryptoClientError, EnrollPinResponse, UpdateKdfResponse, UserCryptoV2KeysResponse,
-        enroll_pin, get_v2_rotated_account_keys, make_update_kdf, make_update_password,
-        make_v2_keys_for_v1_user,
+        CryptoClientError, EnrollPinResponse, RotateableKeySet, UpdateKdfResponse,
+        UserCryptoV2KeysResponse, derive_prf_key, enroll_pin, get_v2_rotated_account_keys,
+        make_update_kdf, make_update_password, make_v2_keys_for_v1_user,
     },
 };
 
@@ -172,6 +172,13 @@ impl CryptoClient {
         derive_pin_user_key(&self.client, encrypted_pin)
     }
 
+    /// Generates a PRF-protected user key from the provided PRF secret. The result can be stored
+    /// and later used to initialize another client instance by using the PRF and the PRF key
+    /// with `initialize_user_crypto`.
+    pub fn derive_prf_key(&self, prf: B64) -> Result<RotateableKeySet, CryptoClientError> {
+        derive_prf_key(&self.client, prf)
+    }
+
     /// Prepares the account for being enrolled in the admin password reset feature. This encrypts
     /// the users [UserKey][bitwarden_crypto::UserKey] with the organization's public key.
     pub fn enroll_admin_password_reset(
diff --git a/crates/bitwarden-crypto/src/keys/mod.rs b/crates/bitwarden-crypto/src/keys/mod.rs
index 1e6cda4db..6a802e853 100644
--- a/crates/bitwarden-crypto/src/keys/mod.rs
+++ b/crates/bitwarden-crypto/src/keys/mod.rs
@@ -34,3 +34,6 @@ pub use kdf::{
 };
 pub(crate) use key_id::{KEY_ID_SIZE, KeyId};
 pub(crate) mod utils;
+
+mod prf;
+pub use prf::derive_symmetric_key_from_prf;
diff --git a/crates/bitwarden-crypto/src/keys/prf.rs b/crates/bitwarden-crypto/src/keys/prf.rs
new file mode 100644
index 000000000..cc83af01c
--- /dev/null
+++ b/crates/bitwarden-crypto/src/keys/prf.rs
@@ -0,0 +1,44 @@
+use crate::{utils::stretch_key, CryptoError, SymmetricCryptoKey};
+
+/// Takes the output of a PRF and derives a symmetric key
+pub fn derive_symmetric_key_from_prf(prf: &[u8]) -> Result<SymmetricCryptoKey, CryptoError> {
+    let (secret, _) = prf
+        .split_at_checked(32)
+        .ok_or_else(|| CryptoError::InvalidKeyLen)?;
+    let secret: [u8; 32] = secret.try_into().unwrap();
+    if secret.iter().all(|b| *b == b'\0') {
+        return Err(CryptoError::ZeroNumber);
+    }
+    Ok(SymmetricCryptoKey::Aes256CbcHmacKey(stretch_key(
+        &Box::pin(secret.into()),
+    )?))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn test_prf_succeeds() {
+        let prf = [
+            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
+            24, 25, 26, 27, 28, 29, 30, 31,
+        ];
+        derive_symmetric_key_from_prf(&prf).unwrap();
+    }
+
+    #[test]
+    fn test_zero_key_fails() {
+        let prf = [
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0,
+        ];
+        let err = derive_symmetric_key_from_prf(&prf).unwrap_err();
+        assert!(matches!(err, CryptoError::ZeroNumber));
+    }
+    #[test]
+    fn test_short_prf_fails() {
+        let prf = [0, 1, 2, 3, 4, 5, 6, 7, 8];
+        let err = derive_symmetric_key_from_prf(&prf).unwrap_err();
+        assert!(matches!(err, CryptoError::InvalidKeyLen));
+    }
+}
diff --git a/crates/bitwarden-uniffi/src/crypto.rs b/crates/bitwarden-uniffi/src/crypto.rs
index d43881ed0..29b660ef4 100644
--- a/crates/bitwarden-uniffi/src/crypto.rs
+++ b/crates/bitwarden-uniffi/src/crypto.rs
@@ -1,6 +1,6 @@
 use bitwarden_core::key_management::crypto::{
     DeriveKeyConnectorRequest, DerivePinKeyResponse, EnrollPinResponse, InitOrgCryptoRequest,
-    InitUserCryptoRequest, UpdateKdfResponse, UpdatePasswordResponse,
+    InitUserCryptoRequest, RotateableKeySet, UpdateKdfResponse, UpdatePasswordResponse,
 };
 use bitwarden_crypto::{EncString, Kdf, UnsignedSharedKey};
 use bitwarden_encoding::B64;
@@ -67,6 +67,10 @@ impl CryptoClient {
         Ok(self.0.enroll_pin(pin)?)
     }
 
+    pub fn derive_prf_key(&self, prf: B64) -> Result<RotateableKeySet> {
+        Ok(self.0.derive_prf_key(prf)?)
+    }
+
     /// Protects the current user key with the provided PIN. The result can be stored and later
     /// used to initialize another client instance by using the PIN and the PIN key with
     /// `initialize_user_crypto`. The provided pin is encrypted with the user key.