fix(auth): preserve rotated keys during key ID initialization

src/api/core/accounts.rs

@@ -1093,13 +1093,12 @@ struct KeyIdData {
 
 #[post("/accounts/key-management/user-key-id", data = "<data>")]
 async fn post_user_key(data: Json<KeyIdData>, headers: Headers, conn: DbConn) -> EmptyResult {
-    let mut user = headers.user;
+    let user = headers.user;
     if user.key_id.is_some() {
         err_code!("Unexpected data", Status::UnprocessableEntity.code);
     }
 
-    user.key_id = Some(data.into_inner().user_key_id);
-    user.save(&conn).await
+    user.set_key_id_if_current(data.into_inner().user_key_id, &conn).await
 }
 
 #[post("/accounts/security-stamp", data = "<data>")]

src/db/models/user.rs

@@ -362,6 +362,47 @@ impl User {
         }
     }
 
+    /// Initialize the client key ID without overwriting a concurrent key rotation.
+    pub async fn set_key_id_if_current(&self, key_id: KeyId, conn: &DbConn) -> EmptyResult {
+        let expected = (
+            self.uuid.clone(),
+            self.security_stamp.clone(),
+            self.akey.clone(),
+            self.private_key.clone(),
+            self.public_key.clone(),
+        );
+        let updated_at = Utc::now().naive_utc();
+        conn.run(move |conn| {
+            conn.transaction::<_, Error, _>(|conn| {
+                let failed = || Error::new_msg("Unable to update user key ID");
+                // Match the portable account lock used by passkey key rotation.
+                diesel::update(users::table.filter(users::uuid.eq(&expected.0)))
+                    .set(users::uuid.eq(&expected.0))
+                    .execute(conn)
+                    .map_err(|_| failed())?;
+                let current =
+                    users::table.filter(users::uuid.eq(&expected.0)).first::<Self>(conn).map_err(|_| failed())?;
+                if current.key_id.is_some() {
+                    return Err(Error::new_msg("Unexpected data").with_code(422));
+                }
+                if !current.enabled
+                    || current.security_stamp != expected.1
+                    || current.akey != expected.2
+                    || current.private_key != expected.3
+                    || current.public_key != expected.4
+                {
+                    return Err(Error::new_msg("Account keys changed. Sync and try again."));
+                }
+                diesel::update(users::table.filter(users::uuid.eq(&expected.0)))
+                    .set((users::key_id.eq(&key_id), users::updated_at.eq(updated_at)))
+                    .execute(conn)
+                    .map_err(|_| failed())?;
+                Ok(())
+            })
+        })
+        .await
+    }
+
     /// Update only the asymmetric pair, without replaying unrelated fields from
     /// an earlier request snapshot over a completed key rotation.
     pub async fn replace_keypair_if_current(

tests/passkey_unlock.rs

@@ -330,6 +330,22 @@ async fn passkey_enrollment_requires_step_up_and_never_adds_an_sso_bypass() {
     assert_eq!(f.count("web_authn_credentials"), 0);
 }
 
+#[tokio::test]
+async fn user_key_id_initialization_preserves_existing_keys_and_rejects_replacement() {
+    let f = Fixture::with_policy(false).await;
+    let key_id = uuid::Uuid::new_v4().to_string();
+    let data = json!({"userKeyId":key_id});
+    assert_eq!(f.post("/api/accounts/key-management/user-key-id", data.clone()).await.status(), StatusCode::OK);
+    assert_eq!(f.count(&format!("users WHERE key_id='{key_id}' AND akey='synthetic-key' AND private_key='synthetic-private-key' AND security_stamp='synthetic-stamp'")),1);
+    assert_eq!(
+        f.post("/api/accounts/key-management/user-key-id", data).await.status(),
+        StatusCode::UNPROCESSABLE_ENTITY
+    );
+    let sync = f.get("/api/sync").await;
+    assert_eq!(sync["userDecryption"]["userKeyId"], key_id);
+    assert_eq!(sync["userDecryption"]["masterPasswordUnlock"]["masterKeyWrappedUserKey"], "synthetic-key");
+}
+
 #[tokio::test]
 async fn passkeys_are_available_without_opt_in_and_keep_master_password_unlock() {
     let f = Fixture::with_policy(false).await;
@@ -733,7 +749,7 @@ impl Fixture {
 #[tokio::test]
 #[ignore = "requires a disposable PostgreSQL instance; see tests/README.md"]
 async fn passkey_postgres_rotation_serializes_deletion_and_stale_rotation() {
-    for competing_action in ["delete", "rotate"] {
+    for competing_action in ["delete", "rotate", "user-key-id"] {
         let f = Fixture::new().await;
         let (data, id) = f.paused_rotation_data().await;
         let mut pause = PausedRotation::new(&f);
@@ -746,6 +762,13 @@ async fn passkey_postgres_rotation_serializes_deletion_and_stale_rotation() {
                 .bearer_auth(&f.token)
                 .json(&json!({"masterPasswordHash":PASSWORD}));
             tokio::spawn(async move { request.send().await })
+        } else if competing_action == "user-key-id" {
+            let request = f
+                .client
+                .post(format!("{}/api/accounts/key-management/user-key-id", f.url))
+                .bearer_auth(&f.token)
+                .json(&json!({"userKeyId": uuid::Uuid::new_v4().to_string()}));
+            tokio::spawn(async move { request.send().await })
         } else {
             f.start_rotation(&data)
         };
@@ -764,6 +787,9 @@ async fn passkey_postgres_rotation_serializes_deletion_and_stale_rotation() {
         assert_eq!(f.count("users WHERE akey='new-master-wrapper' AND private_key='new-account-private-wrapper'"), 1);
         assert_eq!(f.count("folders WHERE name='rotated-folder-wrapper'"), 1);
         assert_eq!(f.count("web_authn_credentials"), i64::from(competing_action != "delete"));
+        if competing_action == "user-key-id" {
+            assert_eq!(f.count("users WHERE key_id IS NOT NULL"), 0);
+        }
     }
 }