[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,17 @@ 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 }
+serde_bytes = { 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,23 @@
 //! 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 std::str::FromStr;
+
+use bitwarden_api_api::models::SetKeyConnectorKeyRequestModel;
+use bitwarden_core::{
+    Client, UserId,
+    key_management::{
+        AccountCryptographyMakeKeysError, KeyConnectorApiError,
+        account_cryptographic_state::WrappedAccountCryptographicState,
+        key_connector_api_post_or_put_key_connector_key,
+    },
+};
+use bitwarden_crypto::EncString;
+use bitwarden_encoding::B64;
+use bitwarden_error::bitwarden_error;
+use serde_bytes::ByteBuf;
+use thiserror::Error;
+use tracing::info;
 #[cfg(feature = "wasm")]
 use wasm_bindgen::prelude::*;
 
@@ -33,4 +49,99 @@ 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: String,
+        user_id: String,
+    ) -> Result<KeyConnectorRegistrationResult, UserRegistrationError> {
+        let client = &self.client.internal;
+        let api_client = &client.get_api_configurations().await.api_client;
+        let user_id =
+            UserId::from_str(user_id.as_str()).map_err(|_| UserRegistrationError::Serialization)?;
+
+        // First call crypto API to get all keys
+        info!("Initializing account cryptography");
+        let (
+            cryptography_state,
+            wrapped_user_key,
+            user_key,
+            account_cryptographic_state_request,
+            key_connector_key,
+        ) = self
+            .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_post_or_put_key_connector_key(
+            &self.client,
+            key_connector_url.as_str(),
+            &key_connector_key,
+        )
+        .await
+        .map_err(UserRegistrationError::KeyConnectorApi)?;
+
+        info!("Posting user account cryptographic state to server");
+        let request = SetKeyConnectorKeyRequestModel {
+            key_connector_key_wrapped_user_key: Some(wrapped_user_key.to_string()),
+            account_keys: Some(Box::new(account_cryptographic_state_request)),
+            ..SetKeyConnectorKeyRequestModel::new(org_id)
+        };
+        api_client
+            .accounts_key_management_api()
+            .post_set_key_connector_key(Some(request))
+            .await
+            .map_err(|e| UserRegistrationError::Api(e.into()))?;
+
+        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: cryptography_state,
+            key_connector_key: key_connector_key.to_base64(),
+            key_connector_key_wrapped_user_key: wrapped_user_key,
+            user_key: user_key.to_encoded().to_vec().into(),
+        })
+    }
+}
+
+/// Result of Key Connector registration process.
+#[cfg_attr(
+    feature = "wasm",
+    derive(tsify::Tsify),
+    tsify(into_wasm_abi, from_wasm_abi)
+)]
+#[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: ByteBuf,
+}
+
+/// Errors that can occur during user registration.
+#[derive(Debug, Error)]
+#[bitwarden_error(flat)]
+pub enum UserRegistrationError {
+    /// Key Connector API call failed.
+    #[error(transparent)]
+    KeyConnectorApi(#[from] KeyConnectorApiError),
+    /// API call failed.
+    #[error(transparent)]
+    Api(#[from] bitwarden_core::ApiError),
+    /// Account cryptography initialization failed.
+    #[error(transparent)]
+    AccountCryptographyMakeKeys(#[from] AccountCryptographyMakeKeysError),
+    /// Serialization or deserialization error
+    #[error("Serialization error")]
+    Serialization,
 }

crates/bitwarden-core/Cargo.toml

@@ -51,6 +51,7 @@ chrono = { workspace = true, features = ["std"] }
 # We don't use this directly (it's used by rand), but we need it here to enable WASM support
 getrandom = { version = ">=0.2.9, <0.3", features = ["js"] }
 rand = ">=0.8.5, <0.9"
+rand_chacha = ">=0.3.1, <0.4.0"
 reqwest = { workspace = true }
 schemars = { workspace = true }
 serde = { workspace = true }
@@ -76,7 +77,6 @@ rustls = { version = "0.23.19", default-features = false }
 rustls-platform-verifier = "0.6.0"
 
 [dev-dependencies]
-rand_chacha = "0.3.1"
 tokio = { workspace = true, features = ["rt"] }
 zeroize = { version = ">=1.7.0, <2.0", features = ["derive", "aarch64"] }