[PM-27231] Account registration v2 for key connector

crates/bitwarden-api-api/src/models/set_key_connector_key_request_model.rs

@@ -14,14 +14,22 @@ use crate::models;
 
 #[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)]
 pub struct SetKeyConnectorKeyRequestModel {
-    #[serde(rename = "key", alias = "Key")]
-    pub key: String,
-    #[serde(rename = "keys", alias = "Keys")]
-    pub keys: Box<models::KeysRequestModel>,
-    #[serde(rename = "kdf", alias = "Kdf")]
-    pub kdf: models::KdfType,
-    #[serde(rename = "kdfIterations", alias = "KdfIterations")]
-    pub kdf_iterations: i32,
+    #[serde(rename = "key", alias = "Key", skip_serializing_if = "Option::is_none")]
+    pub key: Option<String>,
+    #[serde(
+        rename = "keys",
+        alias = "Keys",
+        skip_serializing_if = "Option::is_none"
+    )]
+    pub keys: Option<Box<models::KeysRequestModel>>,
+    #[serde(rename = "kdf", alias = "Kdf", skip_serializing_if = "Option::is_none")]
+    pub kdf: Option<models::KdfType>,
+    #[serde(
+        rename = "kdfIterations",
+        alias = "KdfIterations",
+        skip_serializing_if = "Option::is_none"
+    )]
+    pub kdf_iterations: Option<i32>,
     #[serde(
         rename = "kdfMemory",
         alias = "KdfMemory",
@@ -34,25 +42,33 @@ pub struct SetKeyConnectorKeyRequestModel {
         skip_serializing_if = "Option::is_none"
     )]
     pub kdf_parallelism: Option<i32>,
+    #[serde(
+        rename = "keyConnectorKeyWrappedUserKey",
+        alias = "KeyConnectorKeyWrappedUserKey",
+        skip_serializing_if = "Option::is_none"
+    )]
+    pub key_connector_key_wrapped_user_key: Option<String>,
+    #[serde(
+        rename = "accountKeys",
+        alias = "AccountKeys",
+        skip_serializing_if = "Option::is_none"
+    )]
+    pub account_keys: Option<Box<models::AccountKeysRequestModel>>,
     #[serde(rename = "orgIdentifier", alias = "OrgIdentifier")]
     pub org_identifier: String,
 }
 
 impl SetKeyConnectorKeyRequestModel {
-    pub fn new(
-        key: String,
-        keys: models::KeysRequestModel,
-        kdf: models::KdfType,
-        kdf_iterations: i32,
-        org_identifier: String,
-    ) -> SetKeyConnectorKeyRequestModel {
+    pub fn new(org_identifier: String) -> SetKeyConnectorKeyRequestModel {
         SetKeyConnectorKeyRequestModel {
-            key,
-            keys: Box::new(keys),
-            kdf,
-            kdf_iterations,
+            key: None,
+            keys: None,
+            kdf: None,
+            kdf_iterations: None,
             kdf_memory: None,
             kdf_parallelism: None,
+            key_connector_key_wrapped_user_key: None,
+            account_keys: None,
             org_identifier,
         }
     }

crates/bitwarden-auth/Cargo.toml

@@ -25,12 +25,16 @@ uniffi = ["bitwarden-core/uniffi", "dep:uniffi"] # Uniffi bindings
 
 # Note: dependencies must be alphabetized to pass the cargo sort check in the CI pipeline.
 [dependencies]
+bitwarden-api-api = { workspace = true }
 bitwarden-core = { workspace = true, features = ["internal"] }
+bitwarden-crypto = { workspace = true }
+bitwarden-encoding = { workspace = true }
 bitwarden-error = { workspace = true }
 chrono = { workspace = true }
 reqwest = { workspace = true }
 serde = { workspace = true }
 thiserror = { workspace = true }
+tracing = { workspace = true }
 tsify = { workspace = true, optional = true }
 uniffi = { workspace = true, optional = true }
 wasm-bindgen = { workspace = true, optional = true }

crates/bitwarden-auth/src/registration.rs

@@ -5,7 +5,19 @@
 //! authentication method such as SSO or master password, and a decryption method such as
 //! key-connector, TDE, or master password.
 
-use bitwarden_core::Client;
+use bitwarden_api_api::models::SetKeyConnectorKeyRequestModel;
+use bitwarden_core::{
+    Client, OrganizationId, UserId,
+    key_management::{
+        AccountCryptographyMakeKeysError, KeyConnectorApiClient,
+        account_cryptographic_state::WrappedAccountCryptographicState,
+    },
+};
+use bitwarden_crypto::EncString;
+use bitwarden_encoding::B64;
+use bitwarden_error::bitwarden_error;
+use thiserror::Error;
+use tracing::{error, info};
 #[cfg(feature = "wasm")]
 use wasm_bindgen::prelude::*;
 
@@ -33,4 +45,116 @@ impl RegistrationClient {
         let api_client = &client.get_api_configurations().await.api_client;
         // Do API request here. It will be authenticated using the client's tokens.
     }
+
+    /// Initializes a new cryptographic state for a user and posts it to the server; enrolls the
+    /// user to key connector unlock.
+    pub async fn post_keys_for_key_connector_registration(
+        &self,
+        key_connector_url: String,
+        org_id: OrganizationId,
+        user_id: UserId,
+    ) -> Result<KeyConnectorRegistrationResult, UserRegistrationError> {
+        let client = &self.client.internal;
+        let api_client = &client.get_api_configurations().await.api_client;
+        let key_connector_api_client =
+            KeyConnectorApiClient::new(client, key_connector_url.as_str());
+
+        internal_post_keys_for_key_connector_registration(
+            self,
+            api_client,
+            &key_connector_api_client,
+            org_id,
+            user_id,
+        )
+        .await
+    }
+}
+
+async fn internal_post_keys_for_key_connector_registration(
+    registration_client: &RegistrationClient,
+    api_client: &bitwarden_api_api::apis::ApiClient,
+    key_connector_api_client: &KeyConnectorApiClient,
+    org_id: OrganizationId,
+    user_id: UserId,
+) -> Result<KeyConnectorRegistrationResult, UserRegistrationError> {
+    // First call crypto API to get all keys
+    info!("Initializing account cryptography");
+    let registration_crypto_result = registration_client
+        .client
+        .crypto()
+        .make_user_key_connector_registration(user_id)
+        .map_err(UserRegistrationError::AccountCryptographyMakeKeys)?;
+
+    info!("Posting key connector key to key connector server");
+    key_connector_api_client
+        .post_or_put_key_connector_key(&registration_crypto_result.key_connector_key)
+        .await
+        .map_err(|e| {
+            error!("Failed to post key connector key to key connector server: {e:?}");
+            UserRegistrationError::KeyConnectorApi
+        })?;
+
+    info!("Posting user account cryptographic state to server");
+    let request = SetKeyConnectorKeyRequestModel {
+        key_connector_key_wrapped_user_key: Some(
+            registration_crypto_result
+                .key_connector_key_wrapped_user_key
+                .to_string(),
+        ),
+        account_keys: Some(Box::new(registration_crypto_result.account_keys_request)),
+        ..SetKeyConnectorKeyRequestModel::new(org_id.to_string())
+    };
+    api_client
+        .accounts_key_management_api()
+        .post_set_key_connector_key(Some(request))
+        .await
+        .map_err(|e| {
+            error!("Failed to post account cryptographic state to server: {e:?}");
+            UserRegistrationError::Api
+        })?;
+
+    info!("User initialized!");
+    // Note: This passing out of state and keys is temporary. Once SDK state management is more
+    // mature, the account cryptographic state and keys should be set directly here.
+    Ok(KeyConnectorRegistrationResult {
+        account_cryptographic_state: registration_crypto_result.account_cryptographic_state,
+        key_connector_key: registration_crypto_result.key_connector_key.to_base64(),
+        key_connector_key_wrapped_user_key: registration_crypto_result
+            .key_connector_key_wrapped_user_key,
+        user_key: registration_crypto_result.user_key.to_encoded().into(),
+    })
+}
+
+/// Result of Key Connector registration process.
+#[cfg_attr(
+    feature = "wasm",
+    derive(tsify::Tsify),
+    tsify(into_wasm_abi, from_wasm_abi)
+)]
+#[cfg_attr(feature = "uniffi", derive(uniffi::Record))]
+#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)]
+pub struct KeyConnectorRegistrationResult {
+    /// The account cryptographic state of the user.
+    pub account_cryptographic_state: WrappedAccountCryptographicState,
+    /// The key connector key used for unlocking.
+    pub key_connector_key: B64,
+    /// The encrypted user key, wrapped with the key connector key.
+    pub key_connector_key_wrapped_user_key: EncString,
+    /// The decrypted user key. This can be used to get the consuming client to an unlocked state.
+    pub user_key: B64,
+}
+
+/// Errors that can occur during user registration.
+#[derive(Debug, Error)]
+#[bitwarden_error(flat)]
+pub enum UserRegistrationError {
+    /// Key Connector API call failed.
+    #[error("Key Connector Api call failed")]
+    KeyConnectorApi,
+    /// API call failed.
+    #[error("Api call failed")]
+    Api,
+    /// Account cryptography initialization failed.
+    #[error(transparent)]
+    AccountCryptographyMakeKeys(#[from] AccountCryptographyMakeKeysError),
 }

crates/bitwarden-core/Cargo.toml

@@ -76,8 +76,10 @@ rustls = { version = "0.23.19", default-features = false }
 rustls-platform-verifier = "0.6.0"
 
 [dev-dependencies]
-rand_chacha = "0.3.1"
+mockall = { version = ">=0.13.1, <0.15" }
+rand_chacha = ">=0.3.1, <0.4.0"
 tokio = { workspace = true, features = ["rt"] }
+wiremock = { workspace = true }
 zeroize = { version = ">=1.7.0, <2.0", features = ["derive", "aarch64"] }
 
 [lints]