7 files changed, 29 insertions, 10 deletions

diff --git a/fastboot/fastboot.cpp b/fastboot/fastboot.cpp
index 14f5e4bdc..202334997 100644
--- a/fastboot/fastboot.cpp
+++ b/fastboot/fastboot.cpp
@@ -30,10 +30,10 @@
 #include "recovery_ui/ui.h"
 
 static const std::vector<std::pair<std::string, Device::BuiltinAction>> kFastbootMenuActions{
-  { "Reboot system now", Device::REBOOT },
+  { "Reboot system now", Device::REBOOT_FROM_FASTBOOT },
   { "Enter recovery", Device::ENTER_RECOVERY },
   { "Reboot to bootloader", Device::REBOOT_BOOTLOADER },
-  { "Power off", Device::SHUTDOWN },
+  { "Power off", Device::SHUTDOWN_FROM_FASTBOOT },
 };
 
 Device::BuiltinAction StartFastboot(Device* device, const std::vector<std::string>& /* args */) {
diff --git a/otautil/include/otautil/sysutil.h b/otautil/include/otautil/sysutil.h
index 48e9011e5..326db8644 100644
--- a/otautil/include/otautil/sysutil.h
+++ b/otautil/include/otautil/sysutil.h
@@ -102,11 +102,11 @@ class MemMapping {
 };
 
 // Reboots the device into the specified target, by additionally handling quiescent reboot mode.
-// 'target' can be an empty string, which indicates booting into Android.
+// All unknown targets reboot into Android.
 bool Reboot(std::string_view target);
 
 // Triggers a shutdown.
-bool Shutdown();
+bool Shutdown(std::string_view target);
 
 // Returns a null-terminated char* array, where the elements point to the C-strings in the given
 // vector, plus an additional nullptr at the end. This is a helper function that facilitates
diff --git a/otautil/sysutil.cpp b/otautil/sysutil.cpp
index 420db4cac..a8829858d 100644
--- a/otautil/sysutil.cpp
+++ b/otautil/sysutil.cpp
@@ -229,9 +229,9 @@ bool Reboot(std::string_view target) {
   return android::base::SetProperty(ANDROID_RB_PROPERTY, cmd);
 }
 
-bool Shutdown() {
-  // "shutdown" doesn't need a "reason" arg nor a comma.
-  return android::base::SetProperty(ANDROID_RB_PROPERTY, "shutdown");
+bool Shutdown(std::string_view target) {
+  std::string cmd = "shutdown," + std::string(target);
+  return android::base::SetProperty(ANDROID_RB_PROPERTY, cmd);
 }
 
 std::vector<char*> StringVectorToNullTerminatedArray(const std::vector<std::string>& args) {
diff --git a/recovery.cpp b/recovery.cpp
index 20e5a1be3..eb0c2b216 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -356,6 +356,8 @@ static Device::BuiltinAction PromptAndWait(Device* device, InstallResult status)
             : device->InvokeMenuItem(chosen_item);
 
     switch (chosen_action) {
+      case Device::REBOOT_FROM_FASTBOOT:    // Can not happen
+      case Device::SHUTDOWN_FROM_FASTBOOT:  // Can not happen
       case Device::NO_ACTION:
         break;
 
diff --git a/recovery_main.cpp b/recovery_main.cpp
index aba9c5d75..f78c1c788 100644
--- a/recovery_main.cpp
+++ b/recovery_main.cpp
@@ -471,7 +471,12 @@ int main(int argc, char** argv) {
     switch (ret) {
       case Device::SHUTDOWN:
         ui->Print("Shutting down...\n");
-        Shutdown();
+        Shutdown("recovery");
+        break;
+
+      case Device::SHUTDOWN_FROM_FASTBOOT:
+        ui->Print("Shutting down...\n");
+        Shutdown("fastboot");
         break;
 
       case Device::REBOOT_BOOTLOADER:
@@ -520,9 +525,19 @@ int main(int argc, char** argv) {
         fastboot = false;
         break;
 
+      case Device::REBOOT:
+        ui->Print("Rebooting...\n");
+        Reboot("recovery_menu");
+        break;
+
+      case Device::REBOOT_FROM_FASTBOOT:
+        ui->Print("Rebooting...\n");
+        Reboot("fastboot_menu");
+        break;
+
       default:
         ui->Print("Rebooting...\n");
-        Reboot("");
+        Reboot("unknown" + std::to_string(ret));
         break;
     }
   }
diff --git a/recovery_ui/include/recovery_ui/device.h b/recovery_ui/include/recovery_ui/device.h
index 7c76cdb0a..9a4edf261 100644
--- a/recovery_ui/include/recovery_ui/device.h
+++ b/recovery_ui/include/recovery_ui/device.h
@@ -58,6 +58,8 @@ class Device {
     REBOOT_FASTBOOT = 17,
     REBOOT_RECOVERY = 18,
     REBOOT_RESCUE = 19,
+    REBOOT_FROM_FASTBOOT = 20,
+    SHUTDOWN_FROM_FASTBOOT = 21,
   };
 
   explicit Device(RecoveryUI* ui);
diff --git a/recovery_ui/ui.cpp b/recovery_ui/ui.cpp
index 7ea9307c9..fc5177576 100644
--- a/recovery_ui/ui.cpp
+++ b/recovery_ui/ui.cpp
@@ -375,7 +375,7 @@ void RecoveryUI::ProcessKey(int key_code, int updown) {
 
       case RecoveryUI::REBOOT:
         if (reboot_enabled) {
-          Reboot("");
+          Reboot("recovery_ui");
           while (true) {
             pause();
           }