fxa_client/internal/

state_manager.rs

/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::collections::{HashMap, HashSet};

use crate::{
    internal::{
        oauth::{AccessTokenInfo, RefreshToken},
        profile::Profile,
        state_persistence::state_to_json,
        CachedResponse, Config, OAuthFlow, PersistedState,
    },
    DeviceCapability, FxaRustAuthState, LocalDevice, Result, ScopedKey,
};

/// Stores and manages the current state of the FxA client
///
/// All fields are private, which means that all state mutations must go through this module.  This
/// makes it easier to reason about state changes.
pub struct StateManager {
    /// State that's persisted to disk
    persisted_state: PersistedState,
    /// In-progress OAuth flows
    flow_store: HashMap<String, OAuthFlow>,
}

impl StateManager {
    pub(crate) fn new(persisted_state: PersistedState) -> Self {
        Self {
            persisted_state,
            flow_store: HashMap::new(),
        }
    }

    pub fn serialize_persisted_state(&self) -> Result<String> {
        state_to_json(&self.persisted_state)
    }

    pub fn config(&self) -> &Config {
        &self.persisted_state.config
    }

    pub fn refresh_token(&self) -> Option<&RefreshToken> {
        self.persisted_state.refresh_token.as_ref()
    }

    pub fn session_token(&self) -> Option<&str> {
        self.persisted_state.session_token.as_deref()
    }

    /// Get our device capabilities
    ///
    /// This is the last set of capabilities passed to `initialize_device` or `ensure_capabilities`
    pub fn device_capabilities(&self) -> &HashSet<DeviceCapability> {
        &self.persisted_state.device_capabilities
    }

    /// Set our device capabilities
    pub fn set_device_capabilities(
        &mut self,
        capabilities_set: impl IntoIterator<Item = DeviceCapability>,
    ) {
        self.persisted_state.device_capabilities = HashSet::from_iter(capabilities_set);
    }

    /// Get the last known LocalDevice info sent back from the server
    pub fn server_local_device_info(&self) -> Option<&LocalDevice> {
        self.persisted_state.server_local_device_info.as_ref()
    }

    /// Update the last known LocalDevice info when getting one back from the server
    pub fn update_server_local_device_info(&mut self, local_device: LocalDevice) {
        self.persisted_state.server_local_device_info = Some(local_device)
    }

    /// Clear out the last known LocalDevice info. This means that the next call to
    /// `ensure_capabilities()` will re-send our capabilities to the server
    ///
    /// This is typically called when something may invalidate the server's knowledge of our
    /// local device capabilities, for example replacing our device info.
    pub fn clear_server_local_device_info(&mut self) {
        self.persisted_state.server_local_device_info = None
    }

    pub fn get_commands_data(&self, key: &str) -> Option<&str> {
        self.persisted_state
            .commands_data
            .get(key)
            .map(String::as_str)
    }

    pub fn set_commands_data(&mut self, key: &str, data: String) {
        self.persisted_state
            .commands_data
            .insert(key.to_string(), data);
    }

    pub fn clear_commands_data(&mut self, key: &str) {
        self.persisted_state.commands_data.remove(key);
    }

    pub fn last_handled_command_index(&self) -> Option<u64> {
        self.persisted_state.last_handled_command
    }

    pub fn set_last_handled_command_index(&mut self, idx: u64) {
        self.persisted_state.last_handled_command = Some(idx)
    }

    pub fn current_device_id(&self) -> Option<&str> {
        self.persisted_state.current_device_id.as_deref()
    }

    pub fn set_current_device_id(&mut self, device_id: String) {
        self.persisted_state.current_device_id = Some(device_id);
    }

    pub fn get_scoped_key(&self, scope: &str) -> Option<&ScopedKey> {
        self.persisted_state.scoped_keys.get(scope)
    }

    pub(crate) fn last_seen_profile(&self) -> Option<&CachedResponse<Profile>> {
        self.persisted_state.last_seen_profile.as_ref()
    }

    pub(crate) fn set_last_seen_profile(&mut self, profile: CachedResponse<Profile>) {
        self.persisted_state.last_seen_profile = Some(profile)
    }

    pub fn clear_last_seen_profile(&mut self) {
        self.persisted_state.last_seen_profile = None
    }

    pub fn get_cached_access_token(&mut self, scope: &str) -> Option<&AccessTokenInfo> {
        self.persisted_state.access_token_cache.get(scope)
    }

    pub fn add_cached_access_token(&mut self, scope: impl Into<String>, token: AccessTokenInfo) {
        self.persisted_state
            .access_token_cache
            .insert(scope.into(), token);
    }

    pub fn clear_access_token_cache(&mut self) {
        self.persisted_state.access_token_cache.clear()
    }

    /// Begin an OAuth flow.  This saves the OAuthFlow for later.  `state` must be unique to this
    /// oauth flow process.
    pub fn begin_oauth_flow(&mut self, state: impl Into<String>, flow: OAuthFlow) {
        self.flow_store.insert(state.into(), flow);
    }

    /// Get an OAuthFlow from a previous `begin_oauth_flow()` call
    ///
    /// This operation removes the OAuthFlow from the our internal map.  It can only be called once
    /// per `state` value.
    pub fn pop_oauth_flow(&mut self, state: &str) -> Option<OAuthFlow> {
        self.flow_store.remove(state)
    }

    /// Complete an OAuth flow.
    pub fn complete_oauth_flow(
        &mut self,
        scoped_keys: Vec<(String, ScopedKey)>,
        refresh_token: RefreshToken,
        new_session_token: Option<String>,
    ) {
        // When our keys change, we might need to re-register device capabilities with the server.
        // Ensure that this happens on the next call to ensure_capabilities.
        self.clear_server_local_device_info();

        for (scope, key) in scoped_keys {
            self.persisted_state.scoped_keys.insert(scope, key);
        }
        self.persisted_state.refresh_token = Some(refresh_token);
        // We prioritize the existing session token if we already have one, because we might have
        // acquired a session token before the oauth flow
        if let (None, Some(new_session_token)) = (self.session_token(), new_session_token) {
            self.set_session_token(new_session_token)
        }
        self.persisted_state.logged_out_from_auth_issues = false;
        self.flow_store.clear();
    }

    /// Called when the account is disconnected.  This clears most of the auth state, but keeps
    /// some information in order to eventually reconnect to the same user account later.
    pub fn disconnect(&mut self) {
        self.persisted_state.current_device_id = None;
        self.persisted_state.refresh_token = None;
        self.persisted_state.scoped_keys = HashMap::new();
        self.persisted_state.last_handled_command = None;
        self.persisted_state.commands_data = HashMap::new();
        self.persisted_state.access_token_cache = HashMap::new();
        self.persisted_state.device_capabilities = HashSet::new();
        self.persisted_state.server_local_device_info = None;
        self.persisted_state.session_token = None;
        self.persisted_state.logged_out_from_auth_issues = false;
        self.flow_store.clear();
    }

    /// Called when we notice authentication issues with the account state.
    ///
    /// This clears the auth state, but leaves some fields untouched. That way, if the user
    /// re-authenticates they can continue using the account without unexpected behavior.  The
    /// fields that don't change compared to `disconnect()` are:
    ///
    ///   * `current_device_id`
    ///   * `device_capabilities`
    ///   * `last_handled_command`
    pub fn on_auth_issues(&mut self) {
        self.persisted_state.refresh_token = None;
        self.persisted_state.scoped_keys = HashMap::new();
        self.persisted_state.commands_data = HashMap::new();
        self.persisted_state.access_token_cache = HashMap::new();
        self.persisted_state.server_local_device_info = None;
        self.persisted_state.session_token = None;
        self.persisted_state.logged_out_from_auth_issues = true;
        self.flow_store.clear();
    }

    /// Called when we begin an OAuth flow.
    ///
    /// This clears out tokens/keys set from the previous time we completed an oauth flow.  In
    /// particular, it clears the session token to avoid
    /// https://bugzilla.mozilla.org/show_bug.cgi?id=1887071.
    pub fn on_begin_oauth(&mut self) {
        self.persisted_state.refresh_token = None;
        self.persisted_state.scoped_keys = HashMap::new();
        self.persisted_state.commands_data = HashMap::new();
        self.persisted_state.access_token_cache = HashMap::new();
        self.persisted_state.session_token = None;
    }

    pub fn get_auth_state(&self) -> FxaRustAuthState {
        if self.persisted_state.refresh_token.is_some() {
            FxaRustAuthState::Connected
        } else if self.persisted_state.logged_out_from_auth_issues {
            FxaRustAuthState::AuthIssues
        } else {
            FxaRustAuthState::Disconnected
        }
    }

    /// Handle the auth tokens changing
    ///
    /// This method updates the token data and clears out data that may be invalidated with the
    /// token changes.
    pub fn update_tokens(&mut self, session_token: String, refresh_token: RefreshToken) {
        self.persisted_state.session_token = Some(session_token);
        self.persisted_state.refresh_token = Some(refresh_token);
        self.persisted_state.access_token_cache.clear();
        self.persisted_state.server_local_device_info = None;
    }

    /// Used by the application to test auth token issues
    pub fn simulate_temporary_auth_token_issue(&mut self) {
        for (_, access_token) in self.persisted_state.access_token_cache.iter_mut() {
            "invalid-data".clone_into(&mut access_token.token)
        }
    }

    /// Used by the application to test auth token issues
    pub fn simulate_permanent_auth_token_issue(&mut self) {
        self.persisted_state.session_token = None;
        self.persisted_state.refresh_token = None;
        self.persisted_state.access_token_cache.clear();
    }
    pub fn set_session_token(&mut self, token: String) {
        self.persisted_state.session_token = Some(token)
    }
}

#[cfg(test)]
impl StateManager {
    pub fn is_access_token_cache_empty(&self) -> bool {
        self.persisted_state.access_token_cache.is_empty()
    }

    pub fn force_refresh_token(&mut self, token: RefreshToken) {
        self.persisted_state.refresh_token = Some(token)
    }

    pub fn force_current_device_id(&mut self, device_id: impl Into<String>) {
        self.persisted_state.current_device_id = Some(device_id.into())
    }

    pub fn insert_scoped_key(&mut self, scope: impl Into<String>, key: ScopedKey) {
        self.persisted_state.scoped_keys.insert(scope.into(), key);
    }
}