Skip to main content

bitwarden_auth/registration/
post_keys_for_key_connector_registration.rs

1//! Initializes a new cryptographic state for a user and posts it to the server; enrolls the
2//! user to key connector unlock.
3use bitwarden_api_api::models::SetKeyConnectorKeyRequestModel;
4use bitwarden_core::key_management::account_cryptographic_state::WrappedAccountCryptographicState;
5use bitwarden_crypto::EncString;
6use bitwarden_encoding::B64;
7use tracing::{error, info};
8
9use crate::registration::{RegistrationClient, RegistrationError};
10
11/// Result of Key Connector registration process.
12#[bitwarden_ffi::wasm_record]
13#[cfg_attr(feature = "uniffi", derive(uniffi::Record))]
14#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)]
15pub struct KeyConnectorRegistrationResult {
16    /// The account cryptographic state of the user.
17    pub account_cryptographic_state: WrappedAccountCryptographicState,
18    /// The key connector key used for unlocking.
19    pub key_connector_key: B64,
20    /// The encrypted user key, wrapped with the key connector key.
21    pub key_connector_key_wrapped_user_key: EncString,
22    /// The decrypted user key. This can be used to get the consuming client to an unlocked state.
23    pub user_key: B64,
24}
25
26#[bitwarden_ffi::wasm_export]
27impl RegistrationClient {
28    /// Initializes a new cryptographic state for a user and posts it to the server; enrolls the
29    /// user to key connector unlock.
30    pub async fn post_keys_for_key_connector_registration(
31        &self,
32        key_connector_url: String,
33        sso_org_identifier: String,
34    ) -> Result<KeyConnectorRegistrationResult, RegistrationError> {
35        let client = &self.client.internal;
36        let configuration = &client.get_api_configurations();
37        let key_connector_client = client.get_key_connector_client(key_connector_url);
38
39        internal_post_keys_for_key_connector_registration(
40            self,
41            &configuration.api_client,
42            &key_connector_client,
43            sso_org_identifier,
44        )
45        .await
46    }
47}
48
49async fn internal_post_keys_for_key_connector_registration(
50    registration_client: &RegistrationClient,
51    api_client: &bitwarden_api_api::apis::ApiClient,
52    key_connector_api_client: &bitwarden_api_key_connector::apis::ApiClient,
53    sso_org_identifier: String,
54) -> Result<KeyConnectorRegistrationResult, RegistrationError> {
55    // First call crypto API to get all keys
56    info!("Initializing account cryptography");
57    let registration_crypto_result = registration_client
58        .client
59        .crypto()
60        .make_user_key_connector_registration()
61        .map_err(|_| RegistrationError::Crypto)?;
62
63    info!("Posting key connector key to key connector server");
64    let key_connector_key: B64 = registration_crypto_result.key_connector_key.into();
65    post_key_to_key_connector(key_connector_api_client, &key_connector_key).await?;
66
67    info!("Posting user account cryptographic state to server");
68    let request = SetKeyConnectorKeyRequestModel {
69        key_connector_key_wrapped_user_key: Some(
70            registration_crypto_result
71                .key_connector_key_wrapped_user_key
72                .to_string(),
73        ),
74        account_keys: Some(Box::new(registration_crypto_result.account_keys_request)),
75        ..SetKeyConnectorKeyRequestModel::new(sso_org_identifier.to_string())
76    };
77    api_client
78        .accounts_key_management_api()
79        .post_set_key_connector_key(Some(request))
80        .await
81        .map_err(|e| {
82            error!("Failed to post account cryptographic state to server: {e:?}");
83            RegistrationError::Api
84        })?;
85
86    info!("User initialized!");
87    // Note: This passing out of state and keys is temporary. Once SDK state management is more
88    // mature, the account cryptographic state and keys should be set directly here.
89    Ok(KeyConnectorRegistrationResult {
90        account_cryptographic_state: registration_crypto_result.account_cryptographic_state,
91        key_connector_key,
92        key_connector_key_wrapped_user_key: registration_crypto_result
93            .key_connector_key_wrapped_user_key,
94        user_key: registration_crypto_result.user_key.to_encoded().into(),
95    })
96}
97
98async fn post_key_to_key_connector(
99    key_connector_api_client: &bitwarden_api_key_connector::apis::ApiClient,
100    key_connector_key: &B64,
101) -> Result<(), RegistrationError> {
102    let request =
103        bitwarden_api_key_connector::models::user_key_request_model::UserKeyKeyRequestModel {
104            key: key_connector_key.to_string(),
105        };
106
107    let result = if key_connector_api_client
108        .user_keys_api()
109        .get_user_key()
110        .await
111        .is_ok()
112    {
113        info!("User's key connector key exists, updating");
114        key_connector_api_client
115            .user_keys_api()
116            .put_user_key(request)
117            .await
118    } else {
119        info!("User's key connector key does not exist, creating");
120        key_connector_api_client
121            .user_keys_api()
122            .post_user_key(request)
123            .await
124    };
125
126    result.map_err(|e| {
127        error!("Failed to post key connector key to key connector server: {e:?}");
128        RegistrationError::KeyConnectorApi
129    })
130}
131
132#[cfg(test)]
133mod tests {
134    use bitwarden_api_api::apis::ApiClient;
135    use bitwarden_core::Client;
136
137    use super::*;
138
139    const TEST_SSO_ORG_IDENTIFIER: &str = "test-org";
140
141    #[tokio::test]
142    async fn test_post_keys_for_key_connector_registration_success() {
143        let client = Client::new(None);
144        let registration_client = RegistrationClient::new(client);
145
146        let api_client = ApiClient::new_mocked(|mock| {
147            mock.accounts_key_management_api
148                .expect_post_set_key_connector_key()
149                .once()
150                .returning(move |_body| Ok(()));
151        });
152
153        let key_connector_api_client =
154            bitwarden_api_key_connector::apis::ApiClient::new_mocked(|mock| {
155                mock.user_keys_api
156                    .expect_get_user_key()
157                    .once()
158                    .returning(move || {
159                        Err(bitwarden_api_key_connector::apis::Error::ResponseError(
160                            bitwarden_api_key_connector::apis::ResponseContent {
161                                status: reqwest::StatusCode::NOT_FOUND,
162                                content: "Not Found".to_string(),
163                            },
164                        ))
165                    });
166                mock.user_keys_api
167                    .expect_post_user_key()
168                    .once()
169                    .returning(move |_body| Ok(()));
170            });
171
172        let result = internal_post_keys_for_key_connector_registration(
173            &registration_client,
174            &api_client,
175            &key_connector_api_client,
176            TEST_SSO_ORG_IDENTIFIER.to_string(),
177        )
178        .await;
179        assert!(result.is_ok());
180
181        // Assert that the mock expectations were met
182        if let ApiClient::Mock(mut mock) = api_client {
183            mock.accounts_key_management_api.checkpoint();
184        }
185        if let bitwarden_api_key_connector::apis::ApiClient::Mock(mut mock) =
186            key_connector_api_client
187        {
188            mock.user_keys_api.checkpoint();
189        }
190    }
191
192    #[tokio::test]
193    async fn test_post_keys_for_key_connector_registration_key_connector_api_failure() {
194        let client = Client::new(None);
195        let registration_client = RegistrationClient::new(client);
196
197        let api_client = ApiClient::new_mocked(|mock| {
198            // Should not be called if Key Connector API fails
199            mock.accounts_key_management_api
200                .expect_post_set_key_connector_key()
201                .never();
202        });
203
204        let key_connector_api_client =
205            bitwarden_api_key_connector::apis::ApiClient::new_mocked(|mock| {
206                mock.user_keys_api
207                    .expect_get_user_key()
208                    .once()
209                    .returning(move || {
210                        Err(bitwarden_api_key_connector::apis::Error::ResponseError(
211                            bitwarden_api_key_connector::apis::ResponseContent {
212                                status: reqwest::StatusCode::NOT_FOUND,
213                                content: "Not Found".to_string(),
214                            },
215                        ))
216                    });
217                mock.user_keys_api
218                    .expect_post_user_key()
219                    .once()
220                    .returning(move |_body| {
221                        Err(bitwarden_api_key_connector::apis::Error::Serde(
222                            serde_json::Error::io(std::io::Error::other("API error")),
223                        ))
224                    });
225            });
226
227        let result = internal_post_keys_for_key_connector_registration(
228            &registration_client,
229            &api_client,
230            &key_connector_api_client,
231            TEST_SSO_ORG_IDENTIFIER.to_string(),
232        )
233        .await;
234
235        assert!(result.is_err());
236        assert!(matches!(
237            result.unwrap_err(),
238            RegistrationError::KeyConnectorApi
239        ));
240
241        // Assert that the mock expectations were met
242        if let ApiClient::Mock(mut mock) = api_client {
243            mock.accounts_key_management_api.checkpoint();
244        }
245        if let bitwarden_api_key_connector::apis::ApiClient::Mock(mut mock) =
246            key_connector_api_client
247        {
248            mock.user_keys_api.checkpoint();
249        }
250    }
251
252    #[tokio::test]
253    async fn test_post_keys_for_key_connector_registration_api_failure() {
254        let client = Client::new(None);
255        let registration_client = RegistrationClient::new(client);
256
257        let api_client = ApiClient::new_mocked(|mock| {
258            mock.accounts_key_management_api
259                .expect_post_set_key_connector_key()
260                .once()
261                .returning(move |_body| {
262                    Err(serde_json::Error::io(std::io::Error::other("API error")).into())
263                });
264        });
265
266        let key_connector_api_client =
267            bitwarden_api_key_connector::apis::ApiClient::new_mocked(|mock| {
268                mock.user_keys_api
269                    .expect_get_user_key()
270                    .once()
271                    .returning(move || {
272                        Err(bitwarden_api_key_connector::apis::Error::ResponseError(
273                            bitwarden_api_key_connector::apis::ResponseContent {
274                                status: reqwest::StatusCode::NOT_FOUND,
275                                content: "Not Found".to_string(),
276                            },
277                        ))
278                    });
279                mock.user_keys_api
280                    .expect_post_user_key()
281                    .once()
282                    .returning(move |_body| Ok(()));
283            });
284
285        let result = internal_post_keys_for_key_connector_registration(
286            &registration_client,
287            &api_client,
288            &key_connector_api_client,
289            TEST_SSO_ORG_IDENTIFIER.to_string(),
290        )
291        .await;
292
293        assert!(result.is_err());
294        assert!(matches!(result.unwrap_err(), RegistrationError::Api));
295
296        // Assert that the mock expectations were met
297        if let ApiClient::Mock(mut mock) = api_client {
298            mock.accounts_key_management_api.checkpoint();
299        }
300        if let bitwarden_api_key_connector::apis::ApiClient::Mock(mut mock) =
301            key_connector_api_client
302        {
303            mock.user_keys_api.checkpoint();
304        }
305    }
306}