bitwarden_core/key_management/
local_user_data_key_state.rs

1use bitwarden_crypto::EncString;
2use tracing::info;
3
4use crate::{
5    Client, UserId,
6    key_management::{self, SymmetricKeySlotId, local_user_data_key::WrappedLocalUserDataKey},
7};
8
9pub(crate) struct InitLocalUserDataKeyError;
10
11/// Stores [`WrappedLocalUserDataKey`] in state if one does not already exist.
12pub(crate) async fn initialize_local_user_data_key_into_state(
13    client: &Client,
14    user_id: UserId,
15) -> Result<(), InitLocalUserDataKeyError> {
16    let repo = client
17        .platform()
18        .state()
19        .get::<key_management::LocalUserDataKeyState>()
20        .map_err(|_| InitLocalUserDataKeyError)?;
21
22    // Idempotent: only set if no key is present yet.
23    if let Ok(Some(_)) = repo.get(user_id).await {
24        info!("WrappedLocalUserDataKey already exists in state, skipping");
25        return Ok(());
26    }
27
28    info!("Setting WrappedLocalUserDataKey to state from user key");
29    let wrapped_local_user_data_key = {
30        let key_store = client.internal.get_key_store();
31        let mut ctx = key_store.context();
32        WrappedLocalUserDataKey::from_context_user_key(&mut ctx)
33            .map_err(|_| InitLocalUserDataKeyError)?
34    };
35
36    repo.set(user_id, wrapped_local_user_data_key.into())
37        .await
38        .map_err(|_| InitLocalUserDataKeyError)
39}
40
41#[derive(Debug)]
42pub(crate) struct MigrateLocalUserDataKeyForUserKeyUpgradeError;
43
44/// Re-wraps a persisted [`WrappedLocalUserDataKey`] with the current user key after a V1→V2
45/// user-key upgrade, preserving the inner-key plaintext so local data encrypted before the
46/// upgrade remains decryptable. No-ops when migration is unnecessary or impossible.
47pub(crate) async fn migrate_local_user_data_key_for_user_key_upgrade(
48    client: &Client,
49    user_id: UserId,
50) -> Result<(), MigrateLocalUserDataKeyForUserKeyUpgradeError> {
51    // Remove when all host clients implement the state bridge - https://bitwarden.atlassian.net/browse/PM-37189
52    if !client.internal.state_bridge.is_registered() {
53        info!("No state bridge registered, skipping WrappedLocalUserDataKey migration");
54        return Ok(());
55    }
56
57    let Some(token) = client.internal.state_bridge.get_v2_upgrade_token().await else {
58        info!(
59            "No V2 upgrade token available from state bridge, skipping WrappedLocalUserDataKey migration"
60        );
61        return Ok(());
62    };
63
64    let repo = client
65        .platform()
66        .state()
67        .get::<key_management::LocalUserDataKeyState>()
68        .map_err(|_| MigrateLocalUserDataKeyForUserKeyUpgradeError)?;
69    let Some(state) = repo
70        .get(user_id)
71        .await
72        .map_err(|_| MigrateLocalUserDataKeyForUserKeyUpgradeError)?
73    else {
74        return Ok(());
75    };
76    if !matches!(
77        state.wrapped_key,
78        EncString::Aes256Cbc_HmacSha256_B64 { .. }
79    ) {
80        info!("WrappedLocalUserDataKey is not a V1 wrapped key, skipping migration");
81        return Ok(());
82    }
83
84    let rewrapped = {
85        let mut ctx = client.internal.get_key_store().context_mut();
86        let Ok(v1_user_key_id) = token.unwrap_v1(SymmetricKeySlotId::User, &mut ctx) else {
87            info!(
88                "Upgrade token does not apply to current user key, skipping WrappedLocalUserDataKey migration"
89            );
90            return Ok(());
91        };
92
93        let wrapped = WrappedLocalUserDataKey(state.wrapped_key);
94        wrapped
95            .rewrap_with_user_key(v1_user_key_id, &mut ctx)
96            .map_err(|_| MigrateLocalUserDataKeyForUserKeyUpgradeError)?
97    };
98
99    info!("Rewrapping WrappedLocalUserDataKey with current user key");
100    repo.set(user_id, rewrapped.into())
101        .await
102        .map_err(|_| MigrateLocalUserDataKeyForUserKeyUpgradeError)
103}
104
105pub(crate) struct UnableToGetError;
106
107/// Retrieves the [`WrappedLocalUserDataKey`] from state.
108pub(crate) async fn get_local_user_data_key_from_state(
109    client: &Client,
110    user_id: UserId,
111) -> Result<WrappedLocalUserDataKey, UnableToGetError> {
112    info!("Getting the WrappedLocalUserDataKey from state");
113    let user_local_data_key_state = client
114        .platform()
115        .state()
116        .get::<key_management::LocalUserDataKeyState>()
117        .map_err(|_| UnableToGetError)?
118        .get(user_id)
119        .await
120        .map_err(|_| UnableToGetError)?
121        .ok_or(UnableToGetError)?;
122
123    Ok(WrappedLocalUserDataKey(
124        user_local_data_key_state.wrapped_key,
125    ))
126}
127
128#[cfg(test)]
129mod tests {
130    use bitwarden_crypto::{KeyStoreContext, SymmetricCryptoKey};
131    use bitwarden_encoding::B64;
132    use uuid::uuid;
133
134    use super::*;
135    use crate::{
136        Client, UserId,
137        key_management::{
138            KeySlotIds, LocalUserDataKeyState, SymmetricKeySlotId, V2UpgradeToken,
139            local_user_data_key::WrappedLocalUserDataKey,
140            state_bridge::test_support::InMemoryStateBridge,
141        },
142    };
143
144    const V1_USER_KEY: &str =
145        "9j9Ruji/tMHlLZ311I5xJugi4pMLbS7WxApM4yTa4is7c1mEgt4ov8fR6/zA9VvgP+wXfx79HG0C+89xMlqksw==";
146    fn load_v1_user_key(ctx: &mut KeyStoreContext<KeySlotIds>) -> SymmetricKeySlotId {
147        let key = SymmetricCryptoKey::try_from(B64::try_from(V1_USER_KEY).unwrap()).unwrap();
148        ctx.add_local_symmetric_key(key)
149    }
150    const V2_USER_KEY: &str = "pQEEAlCg4GEL17wqaWbSzi7WdH1kAzoAARFvBIQDBAUGIFgg1opRU0oX0Rje8I0ufEOx7Xv6NIoOCSAb1ex312/xDqkB";
151    fn load_v2_user_key(ctx: &mut KeyStoreContext<KeySlotIds>) -> SymmetricKeySlotId {
152        let key = SymmetricCryptoKey::try_from(B64::try_from(V2_USER_KEY).unwrap()).unwrap();
153        ctx.add_local_symmetric_key(key)
154    }
155
156    #[tokio::test]
157    async fn test_migrate_noop_when_state_bridge_not_registered() {
158        let client = test_client(ClientVariants::WithoutStateBridge);
159        initialize_state(
160            &client,
161            UserCryptographyVersion::V1,
162            UpgradeTokenVariant::Present,
163            LocalUserKeyVariant::PreUpgrade,
164        )
165        .await;
166        run_migration_and_assert_noop(&client).await;
167    }
168
169    #[tokio::test]
170    async fn test_migrate_noop_when_no_v2_upgrade_token() {
171        let client = test_client(ClientVariants::WithStateBridge);
172        initialize_state(
173            &client,
174            UserCryptographyVersion::V1,
175            UpgradeTokenVariant::NotPresent,
176            LocalUserKeyVariant::PreUpgrade,
177        )
178        .await;
179        run_migration_and_assert_noop(&client).await;
180    }
181
182    #[tokio::test]
183    async fn test_migrate_noop_when_no_wrapped_key() {
184        let client = test_client(ClientVariants::WithStateBridge);
185        initialize_state(
186            &client,
187            UserCryptographyVersion::V1,
188            UpgradeTokenVariant::Present,
189            LocalUserKeyVariant::NotPresent,
190        )
191        .await;
192        run_migration_and_assert_noop(&client).await;
193    }
194
195    #[tokio::test]
196    async fn test_migrate_noop_when_wrapped_key_is_not_v1() {
197        let client = test_client(ClientVariants::WithStateBridge);
198        initialize_state(
199            &client,
200            UserCryptographyVersion::V2,
201            UpgradeTokenVariant::Present,
202            LocalUserKeyVariant::PostUpgrade,
203        )
204        .await;
205        run_migration_and_assert_noop(&client).await;
206    }
207
208    #[tokio::test]
209    async fn test_migrate_happy_path_rewraps_and_preserves_payload() {
210        let client = test_client(ClientVariants::WithStateBridge);
211        initialize_state(
212            &client,
213            UserCryptographyVersion::V2,
214            UpgradeTokenVariant::Present,
215            LocalUserKeyVariant::PreUpgrade,
216        )
217        .await;
218        let before = read_present_local_user_data_key(&client)
219            .await
220            .expect("LocalUserDataKeyState should be present after initialization");
221        run_migration(&client).await;
222        let after = read_present_local_user_data_key(&client)
223            .await
224            .expect("LocalUserDataKeyState should be present after migration");
225
226        assert_local_user_data_key_is_correct(&client, (&before).into());
227        assert!(matches!(
228            before.wrapped_key,
229            EncString::Aes256Cbc_HmacSha256_B64 { .. }
230        ));
231        assert_local_user_data_key_is_correct(&client, (&after).into());
232        assert!(matches!(
233            after.wrapped_key,
234            EncString::Cose_Encrypt0_B64 { .. }
235        ));
236    }
237
238    // Test helper functions
239
240    fn test_user_id() -> UserId {
241        UserId::new(uuid!("00000000-0000-0000-0000-000000000001"))
242    }
243
244    enum ClientVariants {
245        WithStateBridge,
246        WithoutStateBridge,
247    }
248
249    /// Builds a test client, optionally registering an in-memory state bridge.
250    fn test_client(variant: ClientVariants) -> Client {
251        let client = Client::new_test(None);
252        if let ClientVariants::WithStateBridge = variant {
253            client
254                .km_state_bridge()
255                .register_bridge(Box::new(InMemoryStateBridge::default()));
256        }
257        client
258    }
259
260    enum UserCryptographyVersion {
261        V1,
262        V2,
263    }
264
265    #[derive(PartialEq)]
266    enum UpgradeTokenVariant {
267        Present,
268        NotPresent,
269    }
270
271    #[derive(PartialEq)]
272    enum LocalUserKeyVariant {
273        PreUpgrade,
274        PostUpgrade,
275        NotPresent,
276    }
277
278    /// Persists a freshly-generated user key into `SymmetricKeySlotId::User` and stores a
279    /// `WrappedLocalUserDataKey` wrapped with it.
280    async fn initialize_state(
281        client: &Client,
282        user_cryptography_version: UserCryptographyVersion,
283        upgrade_token_variant: UpgradeTokenVariant,
284        local_user_key_variant: LocalUserKeyVariant,
285    ) {
286        let (upgrade_token, wrapped_key) = {
287            let mut ctx = client.internal.get_key_store().context_mut();
288            let v1_user_key = load_v1_user_key(&mut ctx);
289            ctx.persist_symmetric_key(v1_user_key, SymmetricKeySlotId::User)
290                .expect("persisting V1 user key should succeed");
291
292            // We start out with a v1 state, and the wrapped local user data key is a v1 key
293            let mut wrapped_local_user_data_key =
294                WrappedLocalUserDataKey::from_context_user_key(&mut ctx)
295                    .expect("wrapping should succeed");
296            if let LocalUserKeyVariant::PostUpgrade = local_user_key_variant {
297                let v1_user_key = load_v1_user_key(&mut ctx);
298                let v2_user_key = load_v2_user_key(&mut ctx);
299                ctx.persist_symmetric_key(v2_user_key, SymmetricKeySlotId::User)
300                    .expect("persisting V2 user key should succeed");
301                wrapped_local_user_data_key = wrapped_local_user_data_key
302                    .rewrap_with_user_key(v1_user_key, &mut ctx)
303                    .expect("rewrap with V1 user key should succeed");
304            }
305
306            // Note: Persisting clears the key slots we are persisting from, so we have to reload
307            // the keys
308            let v1_user_key = load_v1_user_key(&mut ctx);
309            let v2_user_key = load_v2_user_key(&mut ctx);
310
311            let upgrade_token = V2UpgradeToken::create(v1_user_key, v2_user_key, &ctx)
312                .expect("upgrade token creation should succeed");
313
314            match user_cryptography_version {
315                UserCryptographyVersion::V1 => {
316                    ctx.persist_symmetric_key(v1_user_key, SymmetricKeySlotId::User)
317                }
318                UserCryptographyVersion::V2 => {
319                    ctx.persist_symmetric_key(v2_user_key, SymmetricKeySlotId::User)
320                }
321            }
322            .expect("persisting user key should succeed");
323
324            (upgrade_token, wrapped_local_user_data_key)
325        };
326
327        if let UpgradeTokenVariant::Present = upgrade_token_variant
328            && client.km_state_bridge().is_bridge_registered()
329        {
330            client
331                .km_state_bridge()
332                .set_v2_upgrade_token(&upgrade_token)
333                .await;
334        }
335
336        if local_user_key_variant == LocalUserKeyVariant::PreUpgrade
337            || local_user_key_variant == LocalUserKeyVariant::PostUpgrade
338        {
339            client
340                .platform()
341                .state()
342                .get::<LocalUserDataKeyState>()
343                .unwrap()
344                .set(test_user_id(), wrapped_key.into())
345                .await
346                .unwrap();
347        }
348    }
349
350    /// Reads the LocalUserDataKeyState for the test user, panicking if absent.
351    async fn read_present_local_user_data_key(client: &Client) -> Option<LocalUserDataKeyState> {
352        client
353            .platform()
354            .state()
355            .get::<LocalUserDataKeyState>()
356            .unwrap()
357            .get(test_user_id())
358            .await
359            .expect("getting LocalUserDataKeyState from state should succeed")
360    }
361
362    async fn run_migration(client: &Client) {
363        migrate_local_user_data_key_for_user_key_upgrade(client, test_user_id())
364            .await
365            .expect("migration should succeed")
366    }
367
368    /// Runs the migration and asserts that the wrapped key in state is unchanged.
369    async fn run_migration_and_assert_noop(client: &Client) {
370        let before = read_present_local_user_data_key(client).await;
371        run_migration(client).await;
372        let after = read_present_local_user_data_key(client).await;
373        assert_eq!(after.map(|k| k.wrapped_key), before.map(|k| k.wrapped_key));
374    }
375
376    /// Asserts that the provided wrapped local user data key can be decrypted, and is the v1 test
377    /// key
378    fn assert_local_user_data_key_is_correct(ctx: &Client, wrapped_key: WrappedLocalUserDataKey) {
379        let mut ctx = ctx.internal.get_key_store().context_mut();
380        let v1_user_key = load_v1_user_key(&mut ctx);
381        let v2_user_key = load_v2_user_key(&mut ctx);
382
383        match wrapped_key.0 {
384            EncString::Aes256Cbc_HmacSha256_B64 { .. } => {
385                let local_user_data_key = ctx
386                    .unwrap_symmetric_key(v1_user_key, &wrapped_key.0)
387                    .expect("unwrapping with V1 user key should succeed");
388                ctx.assert_symmetric_keys_equal(local_user_data_key, v1_user_key)
389            }
390            EncString::Cose_Encrypt0_B64 { .. } => {
391                let local_user_data_key = ctx
392                    .unwrap_symmetric_key(v2_user_key, &wrapped_key.0)
393                    .expect("unwrapping with V2 user key should succeed");
394                ctx.assert_symmetric_keys_equal(local_user_data_key, v1_user_key)
395            }
396            _ => panic!("unexpected encoding variant"),
397        }
398    }
399}
bitwarden_core/key_management/local_user_data_key_state.rs

bitwarden_core/key_management/
local_user_data_key_state.rs
