1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
45use std::collections::{HashMap, HashSet};
67use crate::{
8 internal::{
9 oauth::{AccessTokenInfo, RefreshToken},
10 profile::Profile,
11 state_persistence::state_to_json,
12 CachedResponse, Config, OAuthFlow, PersistedState,
13 },
14 DeviceCapability, FxaRustAuthState, LocalDevice, Result, ScopedKey,
15};
1617/// Stores and manages the current state of the FxA client
18///
19/// All fields are private, which means that all state mutations must go through this module. This
20/// makes it easier to reason about state changes.
21pub struct StateManager {
22/// State that's persisted to disk
23persisted_state: PersistedState,
24/// In-progress OAuth flows
25flow_store: HashMap<String, OAuthFlow>,
26}
2728impl StateManager {
29pub(crate) fn new(persisted_state: PersistedState) -> Self {
30Self {
31 persisted_state,
32 flow_store: HashMap::new(),
33 }
34 }
3536pub fn serialize_persisted_state(&self) -> Result<String> {
37 state_to_json(&self.persisted_state)
38 }
3940pub fn config(&self) -> &Config {
41&self.persisted_state.config
42 }
4344pub fn refresh_token(&self) -> Option<&RefreshToken> {
45self.persisted_state.refresh_token.as_ref()
46 }
4748pub fn session_token(&self) -> Option<&str> {
49self.persisted_state.session_token.as_deref()
50 }
5152/// Get our device capabilities
53 ///
54 /// This is the last set of capabilities passed to `initialize_device` or `ensure_capabilities`
55pub fn device_capabilities(&self) -> &HashSet<DeviceCapability> {
56&self.persisted_state.device_capabilities
57 }
5859/// Set our device capabilities
60pub fn set_device_capabilities(
61&mut self,
62 capabilities_set: impl IntoIterator<Item = DeviceCapability>,
63 ) {
64self.persisted_state.device_capabilities = HashSet::from_iter(capabilities_set);
65 }
6667/// Get the last known LocalDevice info sent back from the server
68pub fn server_local_device_info(&self) -> Option<&LocalDevice> {
69self.persisted_state.server_local_device_info.as_ref()
70 }
7172/// Update the last known LocalDevice info when getting one back from the server
73pub fn update_server_local_device_info(&mut self, local_device: LocalDevice) {
74self.persisted_state.server_local_device_info = Some(local_device)
75 }
7677/// Clear out the last known LocalDevice info. This means that the next call to
78 /// `ensure_capabilities()` will re-send our capabilities to the server
79 ///
80 /// This is typically called when something may invalidate the server's knowledge of our
81 /// local device capabilities, for example replacing our device info.
82pub fn clear_server_local_device_info(&mut self) {
83self.persisted_state.server_local_device_info = None
84}
8586pub fn get_commands_data(&self, key: &str) -> Option<&str> {
87self.persisted_state
88 .commands_data
89 .get(key)
90 .map(String::as_str)
91 }
9293pub fn set_commands_data(&mut self, key: &str, data: String) {
94self.persisted_state
95 .commands_data
96 .insert(key.to_string(), data);
97 }
9899pub fn clear_commands_data(&mut self, key: &str) {
100self.persisted_state.commands_data.remove(key);
101 }
102103pub fn last_handled_command_index(&self) -> Option<u64> {
104self.persisted_state.last_handled_command
105 }
106107pub fn set_last_handled_command_index(&mut self, idx: u64) {
108self.persisted_state.last_handled_command = Some(idx)
109 }
110111pub fn current_device_id(&self) -> Option<&str> {
112self.persisted_state.current_device_id.as_deref()
113 }
114115pub fn set_current_device_id(&mut self, device_id: String) {
116self.persisted_state.current_device_id = Some(device_id);
117 }
118119pub fn get_scoped_key(&self, scope: &str) -> Option<&ScopedKey> {
120self.persisted_state.scoped_keys.get(scope)
121 }
122123pub(crate) fn last_seen_profile(&self) -> Option<&CachedResponse<Profile>> {
124self.persisted_state.last_seen_profile.as_ref()
125 }
126127pub(crate) fn set_last_seen_profile(&mut self, profile: CachedResponse<Profile>) {
128self.persisted_state.last_seen_profile = Some(profile)
129 }
130131pub fn clear_last_seen_profile(&mut self) {
132self.persisted_state.last_seen_profile = None
133}
134135pub fn get_cached_access_token(&mut self, scope: &str) -> Option<&AccessTokenInfo> {
136self.persisted_state.access_token_cache.get(scope)
137 }
138139pub fn add_cached_access_token(&mut self, scope: impl Into<String>, token: AccessTokenInfo) {
140self.persisted_state
141 .access_token_cache
142 .insert(scope.into(), token);
143 }
144145pub fn clear_access_token_cache(&mut self) {
146self.persisted_state.access_token_cache.clear()
147 }
148149/// Begin an OAuth flow. This saves the OAuthFlow for later. `state` must be unique to this
150 /// oauth flow process.
151pub fn begin_oauth_flow(&mut self, state: impl Into<String>, flow: OAuthFlow) {
152self.flow_store.insert(state.into(), flow);
153 }
154155/// Get an OAuthFlow from a previous `begin_oauth_flow()` call
156 ///
157 /// This operation removes the OAuthFlow from the our internal map. It can only be called once
158 /// per `state` value.
159pub fn pop_oauth_flow(&mut self, state: &str) -> Option<OAuthFlow> {
160self.flow_store.remove(state)
161 }
162163/// Complete an OAuth flow.
164pub fn complete_oauth_flow(
165&mut self,
166 scoped_keys: Vec<(String, ScopedKey)>,
167 refresh_token: RefreshToken,
168 new_session_token: Option<String>,
169 ) {
170// When our keys change, we might need to re-register device capabilities with the server.
171 // Ensure that this happens on the next call to ensure_capabilities.
172self.clear_server_local_device_info();
173174for (scope, key) in scoped_keys {
175self.persisted_state.scoped_keys.insert(scope, key);
176 }
177self.persisted_state.refresh_token = Some(refresh_token);
178// We prioritize the existing session token if we already have one, because we might have
179 // acquired a session token before the oauth flow
180if let (None, Some(new_session_token)) = (self.session_token(), new_session_token) {
181self.set_session_token(new_session_token)
182 }
183self.persisted_state.logged_out_from_auth_issues = false;
184self.flow_store.clear();
185 }
186187/// Called when the account is disconnected. This clears most of the auth state, but keeps
188 /// some information in order to eventually reconnect to the same user account later.
189pub fn disconnect(&mut self) {
190self.persisted_state.current_device_id = None;
191self.persisted_state.refresh_token = None;
192self.persisted_state.scoped_keys = HashMap::new();
193self.persisted_state.last_handled_command = None;
194self.persisted_state.commands_data = HashMap::new();
195self.persisted_state.access_token_cache = HashMap::new();
196self.persisted_state.device_capabilities = HashSet::new();
197self.persisted_state.server_local_device_info = None;
198self.persisted_state.session_token = None;
199self.persisted_state.logged_out_from_auth_issues = false;
200self.flow_store.clear();
201 }
202203/// Called when we notice authentication issues with the account state.
204 ///
205 /// This clears the auth state, but leaves some fields untouched. That way, if the user
206 /// re-authenticates they can continue using the account without unexpected behavior. The
207 /// fields that don't change compared to `disconnect()` are:
208 ///
209 /// * `current_device_id`
210 /// * `device_capabilities`
211 /// * `last_handled_command`
212pub fn on_auth_issues(&mut self) {
213self.persisted_state.refresh_token = None;
214self.persisted_state.scoped_keys = HashMap::new();
215self.persisted_state.commands_data = HashMap::new();
216self.persisted_state.access_token_cache = HashMap::new();
217self.persisted_state.server_local_device_info = None;
218self.persisted_state.session_token = None;
219self.persisted_state.logged_out_from_auth_issues = true;
220self.flow_store.clear();
221 }
222223/// Called when we begin an OAuth flow.
224 ///
225 /// This clears out tokens/keys set from the previous time we completed an oauth flow. In
226 /// particular, it clears the session token to avoid
227 /// https://bugzilla.mozilla.org/show_bug.cgi?id=1887071.
228pub fn on_begin_oauth(&mut self) {
229self.persisted_state.refresh_token = None;
230self.persisted_state.scoped_keys = HashMap::new();
231self.persisted_state.commands_data = HashMap::new();
232self.persisted_state.access_token_cache = HashMap::new();
233self.persisted_state.session_token = None;
234 }
235236pub fn get_auth_state(&self) -> FxaRustAuthState {
237if self.persisted_state.refresh_token.is_some() {
238 FxaRustAuthState::Connected
239 } else if self.persisted_state.logged_out_from_auth_issues {
240 FxaRustAuthState::AuthIssues
241 } else {
242 FxaRustAuthState::Disconnected
243 }
244 }
245246/// Handle the auth tokens changing
247 ///
248 /// This method updates the token data and clears out data that may be invalidated with the
249 /// token changes.
250pub fn update_tokens(&mut self, session_token: String, refresh_token: RefreshToken) {
251self.persisted_state.session_token = Some(session_token);
252self.persisted_state.refresh_token = Some(refresh_token);
253self.persisted_state.access_token_cache.clear();
254self.persisted_state.server_local_device_info = None;
255 }
256257/// Used by the application to test auth token issues
258pub fn simulate_temporary_auth_token_issue(&mut self) {
259for (_, access_token) in self.persisted_state.access_token_cache.iter_mut() {
260"invalid-data".clone_into(&mut access_token.token)
261 }
262 }
263264/// Used by the application to test auth token issues
265pub fn simulate_permanent_auth_token_issue(&mut self) {
266self.persisted_state.session_token = None;
267self.persisted_state.refresh_token = None;
268self.persisted_state.access_token_cache.clear();
269 }
270pub fn set_session_token(&mut self, token: String) {
271self.persisted_state.session_token = Some(token)
272 }
273}
274275#[cfg(test)]
276impl StateManager {
277pub fn is_access_token_cache_empty(&self) -> bool {
278self.persisted_state.access_token_cache.is_empty()
279 }
280281pub fn force_refresh_token(&mut self, token: RefreshToken) {
282self.persisted_state.refresh_token = Some(token)
283 }
284285pub fn force_current_device_id(&mut self, device_id: impl Into<String>) {
286self.persisted_state.current_device_id = Some(device_id.into())
287 }
288289pub fn insert_scoped_key(&mut self, scope: impl Into<String>, key: ScopedKey) {
290self.persisted_state.scoped_keys.insert(scope.into(), key);
291 }
292}