Skip to main content

recovery_util/ota/
state_machine.rs

1// Copyright 2022 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4use crate::ota::state_machine::DataSharingConsent::Unknown;
5use crate::wlan::NetworkInfo;
6#[cfg(test)]
7use mockall::automock;
8pub use ota_lib::OtaStatus;
9
10// This component maps the current state with a new event to produce a
11// new state. The states, events and state logic have all ben derived
12// from the Recovery OTA UX design.
13// The only state held by the state machine is the current state.
14
15#[derive(Debug, Clone, PartialEq)]
16pub enum Operation {
17    FactoryDataReset,
18    Reinstall,
19}
20
21pub const REBOOT_DELAY_SECONDS: Option<u64> = Some(3);
22
23pub(crate) type Network = String;
24pub(crate) type NetworkInfos = Vec<NetworkInfo>;
25pub(crate) type Password = String;
26pub(crate) type PercentProgress = u8;
27pub(crate) type RebootDelay = Option<u64>;
28
29type Text = String;
30type ErrorMessage = String;
31
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub enum DataSharingConsent {
34    Allow,
35    DontAllow,
36    Unknown,
37}
38
39impl DataSharingConsent {
40    pub fn toggle(&self) -> DataSharingConsent {
41        match self {
42            Self::Allow => Self::DontAllow,
43            Self::DontAllow => Self::Allow,
44            Self::Unknown => Self::Unknown,
45        }
46    }
47}
48
49#[derive(Debug, Clone)]
50pub enum State {
51    Connecting(Network, Password),
52    ConnectionFailed(Network, Password),
53    EnterPassword(Network),
54    EnterWiFi,
55    ExecuteReinstall,
56    FactoryReset,
57    FactoryResetConfirm,
58    FinalizeReinstall(OtaStatus),
59    Failed(Operation, Option<ErrorMessage>),
60    Home,
61    Rebooting(RebootDelay),
62    Reinstall,
63    ReinstallConfirm {
64        desired: DataSharingConsent,
65        reported: DataSharingConsent,
66    },
67    /// Tells the state machine an OTA is in progress
68    /// OtaStatus = None: the reinstall is running and can receive progress updates
69    /// OtaStatus = Some(?): the reinstall has completed with a status
70    ReinstallRunning {
71        status: Option<OtaStatus>,
72        progress: u8,
73    },
74    GetWiFiNetworks,
75    SelectWiFi(NetworkInfos),
76}
77
78impl PartialEq for State {
79    fn eq(&self, other: &Self) -> bool {
80        std::mem::discriminant(self) == std::mem::discriminant(other)
81    }
82}
83
84#[derive(Debug, Clone)]
85pub enum Event {
86    AddNetwork,
87    Cancel,
88    ChooseNetwork,
89    DebugLog(String),
90    Error(ErrorMessage),
91    Networks(NetworkInfos),
92    OtaStatusReceived(OtaStatus),
93    Progress(PercentProgress),
94    //TODO(b/253084947): remove this event when switching to correct update status fidl
95    ProgressComplete(OtaStatus),
96    Reboot,
97    Rebooting,
98    Reinstall,
99    SystemPrivacySetting(DataSharingConsent),
100    SendReports(DataSharingConsent),
101    StartFactoryReset,
102    TryAgain,
103    TryAnotherWay,
104    UserInput(Text),
105    UserInputUnsecuredNetwork(Network),
106    WiFiConnected,
107}
108
109// This tests only for enum entry not the value contained in the enum.
110impl PartialEq for Event {
111    fn eq(&self, other: &Self) -> bool {
112        std::mem::discriminant(self) == std::mem::discriminant(other)
113    }
114}
115
116#[cfg_attr(test, automock)]
117pub trait StateHandler {
118    fn handle_state(&mut self, event: State);
119}
120
121#[cfg_attr(test, automock)]
122pub trait EventProcessor {
123    fn process_event(&mut self, event: Event) -> Option<State>;
124}
125
126pub struct StateMachine {
127    current_state: State,
128}
129
130impl StateMachine {
131    pub fn new(state: State) -> Self {
132        Self { current_state: state }
133    }
134
135    fn event(&mut self, event: Event) -> Option<State> {
136        #[cfg(feature = "debug_logging")]
137        println!("====== SM: state {:?}, event: {:?}", self.current_state, event);
138        let new_state = match (&self.current_state, event) {
139            (_, Event::DebugLog(_)) => None, // Ignore DebugLog Events
140            // Any cancel or error sends us back to the start.
141            (_, Event::Cancel) => Some(State::Home),
142
143            (State::Home, Event::StartFactoryReset) => Some(State::FactoryResetConfirm),
144            (State::Home, Event::TryAnotherWay) => Some(State::Reinstall),
145
146            (State::FactoryResetConfirm, Event::StartFactoryReset) => Some(State::FactoryReset),
147            (State::FactoryReset, Event::Error(_reason)) => {
148                Some(State::Failed(Operation::FactoryDataReset, None))
149            }
150
151            (State::Failed(op, _), Event::TryAgain) => match op {
152                Operation::FactoryDataReset => Some(State::FactoryReset),
153                Operation::Reinstall => Some(State::Reinstall),
154            },
155
156            (State::Reinstall, Event::Reinstall) => Some(State::GetWiFiNetworks),
157            (State::Reinstall, Event::Reboot) => Some(State::Rebooting(REBOOT_DELAY_SECONDS)),
158
159            (State::GetWiFiNetworks, Event::AddNetwork) => Some(State::EnterWiFi),
160            (State::GetWiFiNetworks, Event::Networks(networks)) => {
161                Some(State::SelectWiFi(networks))
162            }
163
164            (State::SelectWiFi(_), Event::UserInput(network)) => {
165                Some(State::EnterPassword(network))
166            }
167            (State::SelectWiFi(_), Event::UserInputUnsecuredNetwork(network)) => {
168                Some(State::Connecting(network, String::new()))
169            }
170            (State::SelectWiFi(_), Event::AddNetwork) => Some(State::EnterWiFi),
171
172            (State::EnterWiFi, Event::UserInput(network)) if network.is_empty() => {
173                Some(State::GetWiFiNetworks)
174            }
175            (State::EnterWiFi, Event::UserInput(network)) => Some(State::EnterPassword(network)),
176
177            (State::EnterPassword(network), Event::UserInput(password)) => {
178                Some(State::Connecting(network.clone(), password))
179            }
180
181            (State::Connecting(_, _), Event::WiFiConnected) => {
182                Some(State::ReinstallConfirm { desired: Unknown, reported: Unknown })
183            }
184            (State::Connecting(network, password), Event::Error(_reason)) => {
185                Some(State::ConnectionFailed(network.clone(), password.clone()))
186            }
187
188            (State::ConnectionFailed(..), Event::ChooseNetwork) => Some(State::GetWiFiNetworks),
189            (State::ConnectionFailed(network, password), Event::TryAgain) => {
190                Some(State::Connecting(network.clone(), password.clone()))
191            }
192
193            (
194                State::ReinstallConfirm { desired: _, reported: _ },
195                Event::SystemPrivacySetting(system_setting),
196            ) => Some(State::ReinstallConfirm {
197                desired: system_setting.clone(),
198                reported: system_setting,
199            }),
200
201            (
202                State::ReinstallConfirm { desired: _, reported },
203                Event::SendReports(desired_setting),
204            ) => Some(State::ReinstallConfirm {
205                desired: desired_setting,
206                reported: reported.clone(),
207            }),
208            (State::ReinstallConfirm { .. }, Event::Reinstall) => Some(State::ExecuteReinstall),
209
210            (State::ExecuteReinstall, Event::Reinstall) => {
211                Some(State::ReinstallRunning { status: None, progress: 0 })
212            }
213            (State::ReinstallRunning { status: None, .. }, Event::Progress(progress)) => {
214                Some(State::ReinstallRunning { status: None, progress: progress })
215            }
216            (State::ReinstallRunning { status: None, .. }, Event::ProgressComplete(status)) => {
217                // TODO(b/253084947): Remove this transition since complete will come from OtaReinstallAction after ota_manager finishes
218                Some(State::ReinstallRunning { status: Some(status), progress: 100 })
219            }
220            (State::ReinstallRunning { .. }, Event::OtaStatusReceived(status)) => {
221                Some(State::FinalizeReinstall(status))
222            }
223            (State::ReinstallRunning { .. }, Event::Error(error))
224            | (State::FinalizeReinstall(_), Event::Error(error)) => {
225                Some(State::Failed(Operation::Reinstall, Some(error)))
226            }
227            // TODO(b/258323217): Add error message to home screen
228            (_, Event::Error(_)) => Some(State::Home),
229            (state, event) => {
230                println!("Error unexpected event {:?} for state {:?}", event, state);
231                None
232            }
233        };
234        if new_state.is_some() {
235            #[cfg(feature = "debug_logging")]
236            println!("====== New state is {:?}", new_state);
237            self.current_state = new_state.as_ref().unwrap().clone();
238        }
239        new_state
240    }
241}
242
243#[cfg_attr(test, automock)]
244impl EventProcessor for StateMachine {
245    fn process_event(&mut self, event: Event) -> Option<State> {
246        self.event(event)
247    }
248}
249
250#[cfg(test)]
251mod test {
252    // TODO(b/258049697): Tests to check the expected flow through more than one state.
253    // c.f. https://cs.opensource.google/fuchsia/fuchsia/+/main:src/recovery/system/src/fdr.rs;l=183.
254
255    use super::OtaStatus;
256    use crate::ota::state_machine::DataSharingConsent::Unknown;
257    use crate::ota::state_machine::{DataSharingConsent, Event, Operation, State, StateMachine};
258    use assert_matches::assert_matches;
259    use std::sync::LazyLock;
260
261    static STATES: LazyLock<Vec<State>> = LazyLock::new(|| {
262        vec![
263            State::Connecting("Network".to_string(), "Password".to_string()),
264            State::ConnectionFailed("Network".to_string(), "Password".to_string()),
265            State::EnterPassword("Network".to_string()),
266            State::EnterWiFi,
267            State::ExecuteReinstall,
268            State::FactoryReset,
269            State::FactoryResetConfirm,
270            State::Failed(Operation::Reinstall, Some("Error message".to_string())),
271            State::FinalizeReinstall(OtaStatus::Succeeded),
272            State::GetWiFiNetworks,
273            State::Home,
274            State::Reinstall,
275            State::ReinstallRunning { status: None, progress: 50 },
276            State::SelectWiFi(vec![]),
277        ]
278    });
279    static EVENTS: LazyLock<Vec<Event>> = LazyLock::new(|| {
280        vec![
281            Event::AddNetwork,
282            Event::Cancel,
283            Event::ChooseNetwork,
284            Event::DebugLog("message".to_string()),
285            Event::Error("Error".to_string()),
286            Event::Networks(Vec::new()),
287            Event::OtaStatusReceived(OtaStatus::Succeeded),
288            Event::Progress(0),
289            Event::ProgressComplete(OtaStatus::Succeeded),
290            Event::Reinstall,
291            Event::SendReports(DataSharingConsent::Allow),
292            Event::StartFactoryReset,
293            Event::TryAgain,
294            Event::TryAnotherWay,
295            Event::UserInput("User Input".to_string()),
296            Event::UserInputUnsecuredNetwork("Network".to_string()),
297            Event::WiFiConnected,
298        ]
299    });
300    // TODO(b/258049617): Enable this when variant_count is in the allowed features list
301    // This will enable a check to make sure all events and states are used
302    // #[test]
303    // fn check_test_validity() {
304    //     assert_eq!(std::mem::variant_count::<State>(), STATES.len());
305    //     assert_eq!(std::mem::variant_count::<Event>(), EVENTS.len());
306    // }
307
308    #[test]
309    fn ensure_all_state_and_event_combos_can_not_crash_state_machine() {
310        let mut sm = StateMachine::new(State::Home);
311        // Generate all possible combinations of States and Events
312        for state in STATES.iter() {
313            for event in EVENTS.iter() {
314                // Set the current state here because sm.event() can change it
315                sm.current_state = state.clone();
316                let _new_state = sm.event(event.clone());
317                if let Some(new_state) = _new_state {
318                    assert_eq!(new_state, sm.current_state);
319                }
320            }
321        }
322    }
323
324    #[test]
325    fn run_through_a_successful_user_flow() {
326        let mut sm = StateMachine::new(State::Home);
327        let mut state = sm.event(Event::TryAnotherWay).unwrap();
328        assert_eq!(state, State::Reinstall);
329        state = sm.event(Event::Reinstall).unwrap();
330        assert_eq!(state, State::GetWiFiNetworks);
331        state = sm.event(Event::AddNetwork).unwrap();
332        assert_eq!(state, State::EnterWiFi);
333        state = sm.event(Event::UserInput("Network".to_string())).unwrap();
334        assert_eq!(state, State::EnterPassword("Network".to_string()));
335        state = sm.event(Event::UserInput("Password".to_string())).unwrap();
336        assert_eq!(state, State::Connecting("Network".to_string(), "Password".to_string()));
337        state = sm.event(Event::WiFiConnected).unwrap();
338        assert_matches!(state, State::ReinstallConfirm { desired: Unknown, reported: Unknown });
339        state = sm.event(Event::Reinstall).unwrap();
340        assert_eq!(state, State::ExecuteReinstall);
341        state = sm.event(Event::Reinstall).unwrap();
342        assert_eq!(state, State::ReinstallRunning { status: None, progress: 0 });
343        state = sm.event(Event::Progress(55)).unwrap();
344        assert_eq!(state, State::ReinstallRunning { status: None, progress: 55 });
345        state = sm.event(Event::ProgressComplete(OtaStatus::Succeeded)).unwrap();
346        assert_eq!(
347            state,
348            State::ReinstallRunning { status: Some(OtaStatus::Succeeded), progress: 100 }
349        );
350        state = sm.event(Event::OtaStatusReceived(OtaStatus::Succeeded)).unwrap();
351        assert_eq!(state, State::FinalizeReinstall(OtaStatus::Succeeded));
352    }
353
354    #[test]
355    fn run_through_factory_reset_with_cancel() {
356        let mut sm = StateMachine::new(State::Home);
357        let mut state = sm.event(Event::StartFactoryReset).unwrap();
358        assert_eq!(state, State::FactoryResetConfirm);
359        state = sm.event(Event::Cancel).unwrap();
360        assert_eq!(state, State::Home);
361        state = sm.event(Event::StartFactoryReset).unwrap();
362        assert_eq!(state, State::FactoryResetConfirm);
363        state = sm.event(Event::StartFactoryReset).unwrap();
364        assert_eq!(state, State::FactoryReset);
365    }
366}