From 4bc838301673a780a626f33e933f17c93e5f27c4 Mon Sep 17 00:00:00 2001
From: Blizzard Finnegan <blizzardfinnegan@gmail.com>
Date: Wed, 14 Jun 2023 13:31:39 -0400
Subject: [PATCH] Continue tweaking reboot algorithm

---
 src/device.rs | 24 ++++++++++++++++++------
 src/tty.rs    |  4 +++-
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/src/device.rs b/src/device.rs
index 10e43a2..5d9ae0f 100755
--- a/src/device.rs
+++ b/src/device.rs
@@ -107,8 +107,8 @@ impl Device{
                         _ = usb_port.read_from_device(None);
                         initial_state = State::LoginPrompt;
                     },
-                    Response::Other | Response::Empty | Response::ShellPrompt 
-                        | Response::LoginPrompt | Response::Rebooting => 
+                    Response::Other | Response::Empty | Response::ShellPrompt |
+                    Response::LoginPrompt | Response::ShuttingDown | Response::Rebooting => 
                             initial_state = State::LoginPrompt,
                     Response::BPOn | Response::BPOff | Response::TempCount(_) =>
                             initial_state = State::LifecycleMenu,
@@ -417,9 +417,22 @@ impl Device{
     }
     pub fn reboot(&mut self) -> () {
         self.usb_tty.write_to_device(Command::Quit);
-        while self.usb_tty.read_from_device(None) != Response::Rebooting {}
-        self.current_state = State::Shutdown;
-        while self.usb_tty.read_from_device(None) != Response::LoginPrompt {}
+        let mut successful_reboot:bool = false;
+        loop{
+            match self.usb_tty.read_from_device(None){
+                Response::LoginPrompt => break,
+                Response::Rebooting => {
+                    log::trace!("Successful reboot detected for device {}.",self.serial);
+                    successful_reboot = true;
+                },
+                Response::ShuttingDown => {
+                    log::warn!("Failed reboot on device {}!",self.serial);
+                    successful_reboot = false;
+                },
+                _ => {}
+            }
+        };
+        if successful_reboot { self.reboots += 1; }
         self.current_state = State::LoginPrompt;
     }
     pub fn is_rebooted(&mut self) -> bool {
@@ -463,7 +476,6 @@ impl Device{
         }
         log::info!("Rebooting {} for the {}th time",self.serial, self.reboots);
         self.reboot();
-        self.reboots += 1;
         self.save_values();
     }
 }
diff --git a/src/tty.rs b/src/tty.rs
index 87dea27..7925221 100755
--- a/src/tty.rs
+++ b/src/tty.rs
@@ -39,6 +39,7 @@ pub enum Response{
     Rebooting,
     Other,
     Empty,
+    ShuttingDown
 }
 
 
@@ -58,10 +59,11 @@ const COMMAND_MAP:Lazy<HashMap<Command,&str>> = Lazy::new(||HashMap::from([
     (Command::Newline,"\n"),
 ]));
 
-const RESPONSES:[(&str,Response);9] = [
+const RESPONSES:[(&str,Response);10] = [
     ("reboot: Restarting",Response::Rebooting),
     ("login:",Response::LoginPrompt),
     ("Password:",Response::PasswordPrompt),
+    ("EXIT Debug menu",Response::ShuttingDown),
     ("root@",Response::ShellPrompt),
     ("Check NIBP In Progress: True",Response::BPOn),
     ("Check NIBP In Progress: False",Response::BPOff),