Let state machine handle CheckAuthorizationStatus

components/fxa-client/src/auth.rs

@@ -311,6 +311,14 @@ pub enum FxaEvent {
     ///
     /// This event is valid for the `Authenticating` state.
     CompleteOAuthFlow { code: String, state: String },
+    /// An `fxaccounts:change_password` WebChannel message arrived on the device that just changed
+    /// its password. `json_payload` is the `data` object of that message and contains the new
+    /// session token. The state machine swaps the session token for a new refresh token and
+    /// re-initialises the device record.
+    ///
+    /// This event is valid for the `Connected` and `AuthIssues` states. In `Authenticating` it
+    /// is a no-op so the in-progress OAuth flow is not disrupted.
+    WebChannelPasswordChange { json_payload: String },
     /// Cancel an OAuth flow.
     ///
     /// Use this to cancel an in-progress OAuth, returning to [FxaState::Disconnected] so the

components/fxa-client/src/fxa_client.udl

@@ -992,6 +992,7 @@ interface FxaEvent {
   CompleteOAuthFlow(string code, string state);
   CancelOAuthFlow();
   CheckAuthorizationStatus();
+  WebChannelPasswordChange(string json_payload);
   Disconnect();
   CallGetProfile();
 };

components/fxa-client/src/internal/oauth.rs

@@ -38,15 +38,43 @@ impl FirefoxAccount {
         Ok(())
     }
 
-    /// Extracts the session token from a WebChannel password change JSON payload and exchanges it
-    /// for a new refresh token via a network call.
+    /// Extracts the session token from a WebChannel password change JSON payload, exchanges it
+    /// for a new refresh token.
     pub fn handle_web_channel_password_change(&mut self, json_payload: &str) -> Result<()> {
         let data: serde_json::Value = serde_json::from_str(json_payload)?;
         let token = data
             .get("sessionToken")
             .and_then(|v| v.as_str())
             .ok_or(Error::NoSessionToken)?;
-        self.handle_session_token_change(token)
+        // Grab the current device before the token swap since destroying the old refresh token
+        // tears down its associated device record server-side.
+        let old_device_info = match self.get_current_device() {
+            Ok(maybe_device) => maybe_device,
+            Err(err) => {
+                warn!(
+                    "Error fetching current device before password change: {:?}",
+                    err
+                );
+                None
+            }
+        };
+        self.handle_session_token_change(token)?;
+        if let Some(ref device_info) = old_device_info {
+            if let Err(err) = self.replace_device(
+                &device_info.display_name,
+                &device_info.device_type,
+                &device_info.push_subscription,
+                &device_info.available_commands,
+            ) {
+                warn!(
+                    "Device information restoration failed after password change: {:?}",
+                    err
+                );
+            } else {
+                info!("Restored device information with new refresh token");
+            }
+        }
+        Ok(())
     }
 
     /// Retrieve the current session token from state

components/fxa-client/src/internal/push.rs

@@ -67,15 +67,11 @@ impl FirefoxAccount {
                 })
             }
             PushPayload::PasswordChanged | PushPayload::PasswordReset => {
-                let status = self.check_authorization_status()?;
-                // clear any device or client data due to password change.
                 self.clear_devices_and_attached_clients_cache();
-                Ok(if !status.active {
-                    AccountEvent::AccountAuthStateChanged
-                } else {
-                    info!("Password change event, but no action required");
-                    AccountEvent::Unknown
-                })
+                // Emit the event and let the
+                // CheckAuthorizationStatus arm verify with the
+                // server before committing to AuthIssues.
+                Ok(AccountEvent::AccountAuthStateChanged)
             }
             PushPayload::Unknown => {
                 info!("Unknown Push command.");
@@ -138,14 +134,8 @@ pub struct AccountDestroyedPushPayload {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::internal::http_client::IntrospectResponse;
-    use crate::internal::http_client::MockFxAClient;
-    use crate::internal::oauth::RefreshToken;
     use crate::internal::CachedResponse;
     use crate::internal::Config;
-    use mockall::predicate::always;
-    use mockall::predicate::eq;
-    use std::sync::Arc;
 
     #[test]
     fn test_deserialize_send_tab_command() {
@@ -196,26 +186,12 @@ mod tests {
     fn test_push_password_reset() {
         let mut fxa =
             FirefoxAccount::with_config(Config::stable_dev("12345678", "https://foo.bar"));
-        let mut client = MockFxAClient::new();
-        client
-            .expect_check_refresh_token_status()
-            .with(always(), eq("refresh_token"))
-            .times(1)
-            .returning(|_, _| Ok(IntrospectResponse { active: false }));
-        fxa.set_client(Arc::new(client));
-        let refresh_token_scopes = std::collections::HashSet::new();
-        fxa.state.force_refresh_token(RefreshToken {
-            token: "refresh_token".to_owned(),
-            scopes: refresh_token_scopes,
-        });
-        fxa.state.force_current_device_id("my_id");
         fxa.devices_cache = Some(CachedResponse {
             response: vec![],
             cached_at: 0,
             etag: "".to_string(),
         });
         let json = "{\"version\":1,\"command\":\"fxaccounts:password_reset\"}";
-        assert!(fxa.devices_cache.is_some());
         let event = fxa.handle_push_message(json).unwrap();
         assert!(matches!(event, AccountEvent::AccountAuthStateChanged));
         assert!(fxa.devices_cache.is_none());
@@ -225,28 +201,14 @@ mod tests {
     fn test_push_password_change() {
         let mut fxa =
             FirefoxAccount::with_config(Config::stable_dev("12345678", "https://foo.bar"));
-        let mut client = MockFxAClient::new();
-        client
-            .expect_check_refresh_token_status()
-            .with(always(), eq("refresh_token"))
-            .times(1)
-            .returning(|_, _| Ok(IntrospectResponse { active: true }));
-        fxa.set_client(Arc::new(client));
-        let refresh_token_scopes = std::collections::HashSet::new();
-        fxa.state.force_refresh_token(RefreshToken {
-            token: "refresh_token".to_owned(),
-            scopes: refresh_token_scopes,
-        });
-        fxa.state.force_current_device_id("my_id");
         fxa.devices_cache = Some(CachedResponse {
             response: vec![],
             cached_at: 0,
             etag: "".to_string(),
         });
         let json = "{\"version\":1,\"command\":\"fxaccounts:password_changed\"}";
-        assert!(fxa.devices_cache.is_some());
         let event = fxa.handle_push_message(json).unwrap();
-        assert!(matches!(event, AccountEvent::Unknown));
+        assert!(matches!(event, AccountEvent::AccountAuthStateChanged));
         assert!(fxa.devices_cache.is_none());
     }
     #[test]

components/fxa-client/src/state_machine/display.rs

@@ -35,6 +35,7 @@ impl fmt::Display for FxaEvent {
             Self::CompleteOAuthFlow { .. } => "CompleteOAthFlow",
             Self::CancelOAuthFlow => "CancelOAthFlow",
             Self::CheckAuthorizationStatus => "CheckAuthorizationStatus",
+            Self::WebChannelPasswordChange { .. } => "WebChannelPwdChange",
             Self::Disconnect => "Disconnect",
             Self::CallGetProfile => "CallGetProfile",
         };

components/fxa-client/src/state_machine/helpers.rs

@@ -114,6 +114,10 @@ impl<'a> RetryingAccount<'a> {
         self.with_auth_recovery(|a| a.complete_oauth_flow(code, state))
     }
 
+    pub fn handle_web_channel_password_change(&mut self, json_payload: &str) -> Result<()> {
+        self.with_auth_recovery(|a| a.handle_web_channel_password_change(json_payload))
+    }
+
     /// Cancels any existing OAuth flow before starting a new one.
     pub fn begin_oauth_flow(
         &mut self,

components/fxa-client/src/state_machine/transitions.rs

@@ -142,6 +142,12 @@ pub fn transition(
                 initial_state,
             })
         }
+        // A WebChannel password change while an OAuth flow is in progress
+        // is a no-op; let the flow finish. Should be rare in practice.
+        (s @ S::Authenticating { .. }, FxaEvent::WebChannelPasswordChange { .. }) => {
+            crate::warn!("WebChannel password change received while Authenticating; ignoring");
+            Ok(s)
+        }
 
         // ── From Connected ──────────────────────────────────────────────
         (S::Connected, FxaEvent::Disconnect) => {
@@ -180,6 +186,14 @@ pub fn transition(
                 initial_state: FxaRustAuthState::Connected,
             })
         }
+        (S::Connected, FxaEvent::WebChannelPasswordChange { json_payload }) => {
+            // The inner call swaps the session token for a new refresh token and re-registers
+            // the device record (push subscription, commands, etc) against the new token.
+            account
+                .handle_web_channel_password_change(&json_payload)
+                .to_state_machine_err(|| S::AuthIssues)?;
+            Ok(S::Connected)
+        }
 
         // ── From AuthIssues ─────────────────────────────────────────────
         (
@@ -203,6 +217,14 @@ pub fn transition(
             account.disconnect();
             Ok(S::Disconnected)
         }
+        (S::AuthIssues, FxaEvent::WebChannelPasswordChange { json_payload }) => {
+            // A concurrent sync/401 may have pushed us here before the webchannel ran. The new
+            // session token recovers us; device re-registration will be handled inside the inner call.
+            account
+                .handle_web_channel_password_change(&json_payload)
+                .to_state_machine_err(|| S::AuthIssues)?;
+            Ok(S::Connected)
+        }
 
         // ── Invalid (state, event) pair ─────────────────────────────────
         (state, event) => Err(StateMachineErr::Fatal(Box::new(
@@ -305,4 +327,61 @@ mod tests {
         let result = transition(&mut wrapper, FxaState::Disconnected, FxaEvent::Disconnect);
         assert_fatal_invalid_transition(result);
     }
+
+    fn assert_handled_lands_at(
+        result: std::result::Result<FxaState, StateMachineErr>,
+        expected: FxaState,
+    ) {
+        match result {
+            Err(StateMachineErr::Handled { target, .. }) => assert_eq!(target, expected),
+            Err(StateMachineErr::Fatal(cause)) => panic!("expected Handled, got Fatal({cause:?})"),
+            Ok(s) => panic!("expected Handled, got Ok({s:?})"),
+        }
+    }
+
+    #[test]
+    fn connected_web_channel_password_change_is_valid_transition() {
+        nss_as::ensure_initialized();
+        let mut account = mock_account();
+        let mut wrapper = RetryingAccount::new(&mut account);
+        let result = transition(
+            &mut wrapper,
+            FxaState::Connected,
+            FxaEvent::WebChannelPasswordChange {
+                json_payload: "{}".to_owned(),
+            },
+        );
+        assert_handled_lands_at(result, FxaState::AuthIssues);
+    }
+
+    #[test]
+    fn auth_issues_web_channel_password_change_is_valid_transition() {
+        nss_as::ensure_initialized();
+        let mut account = mock_account();
+        let mut wrapper = RetryingAccount::new(&mut account);
+        let result = transition(
+            &mut wrapper,
+            FxaState::AuthIssues,
+            FxaEvent::WebChannelPasswordChange {
+                json_payload: "{}".to_owned(),
+            },
+        );
+        assert_handled_lands_at(result, FxaState::AuthIssues);
+    }
+
+    #[test]
+    fn authenticating_web_channel_password_change_stays_in_authenticating() {
+        nss_as::ensure_initialized();
+        let mut account = mock_account();
+        let mut wrapper = RetryingAccount::new(&mut account);
+        let from = authenticating_from(FxaRustAuthState::Connected);
+        let result = transition(
+            &mut wrapper,
+            from.clone(),
+            FxaEvent::WebChannelPasswordChange {
+                json_payload: "{}".to_owned(),
+            },
+        );
+        assert_eq!(result.unwrap(), from);
+    }
 }