recovery: Refactor common setup into main()

Move more common setup into the main function.
Main() handles all 1 time setup, such as ui,
logging, and secontext. Recovery_main()
takes in command line arguments, does any
necessary recovery work, and can be called
multiple times from main().

Test: Recovery works
Bug: 78793464
Change-Id: I2d2595fc342b4ddfa80f4e06b30e44263132acd9
Merged-In: I2d2595fc342b4ddfa80f4e06b30e44263132acd9
diff --git a/common.h b/common.h
index 3dc36a9..c24431b 100644
--- a/common.h
+++ b/common.h
@@ -32,6 +32,7 @@
 extern struct selabel_handle* sehandle;
 extern RecoveryUI* ui;
 extern bool modified_flash;
+extern bool has_cache;
 
 // The current stage, e.g. "1/2".
 extern std::string stage;
diff --git a/recovery.cpp b/recovery.cpp
index e427998..b1a2900 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "private/recovery.h"
+#include "recovery.h"
 
 #include <ctype.h>
 #include <dirent.h>
@@ -32,7 +32,6 @@
 #include <sys/stat.h>
 #include <sys/types.h>
 #include <sys/wait.h>
-#include <time.h>
 #include <unistd.h>
 
 #include <algorithm>
@@ -49,12 +48,8 @@
 #include <android-base/strings.h>
 #include <android-base/unique_fd.h>
 #include <bootloader_message/bootloader_message.h>
-#include <cutils/android_reboot.h>
 #include <cutils/properties.h> /* for property_list */
 #include <healthd/BatteryMonitor.h>
-#include <selinux/android.h>
-#include <selinux/label.h>
-#include <selinux/selinux.h>
 #include <ziparchive/zip_archive.h>
 
 #include "adb_install.h"
@@ -70,7 +65,6 @@
 #include "otautil/sysutil.h"
 #include "roots.h"
 #include "screen_ui.h"
-#include "stub_ui.h"
 #include "ui.h"
 
 static constexpr const char* CACHE_LOG_DIR = "/cache/recovery";
@@ -88,13 +82,9 @@
 // into target_files.zip. Assert the version defined in code and in Android.mk are consistent.
 static_assert(kRecoveryApiVersion == RECOVERY_API_VERSION, "Mismatching recovery API versions.");
 
-static bool has_cache = false;
-
-RecoveryUI* ui = nullptr;
 bool modified_flash = false;
 std::string stage;
 const char* reason = nullptr;
-struct selabel_handle* sehandle;
 
 /*
  * The recovery tool communicates with the main system through /cache files.
@@ -146,77 +136,6 @@
     return android::base::GetBoolProperty("ro.debuggable", false);
 }
 
-// command line args come from, in decreasing precedence:
-//   - the actual command line
-//   - the bootloader control block (one per line, after "recovery")
-//   - the contents of COMMAND_FILE (one per line)
-static std::vector<std::string> get_args(const int argc, char** const argv) {
-  CHECK_GT(argc, 0);
-
-  bootloader_message boot = {};
-  std::string err;
-  if (!read_bootloader_message(&boot, &err)) {
-    LOG(ERROR) << err;
-    // If fails, leave a zeroed bootloader_message.
-    boot = {};
-  }
-  stage = std::string(boot.stage);
-
-  if (boot.command[0] != 0) {
-    std::string boot_command = std::string(boot.command, sizeof(boot.command));
-    LOG(INFO) << "Boot command: " << boot_command;
-  }
-
-  if (boot.status[0] != 0) {
-    std::string boot_status = std::string(boot.status, sizeof(boot.status));
-    LOG(INFO) << "Boot status: " << boot_status;
-  }
-
-  std::vector<std::string> args(argv, argv + argc);
-
-  // --- if arguments weren't supplied, look in the bootloader control block
-  if (args.size() == 1) {
-    boot.recovery[sizeof(boot.recovery) - 1] = '\0';  // Ensure termination
-    std::string boot_recovery(boot.recovery);
-    std::vector<std::string> tokens = android::base::Split(boot_recovery, "\n");
-    if (!tokens.empty() && tokens[0] == "recovery") {
-      for (auto it = tokens.begin() + 1; it != tokens.end(); it++) {
-        // Skip empty and '\0'-filled tokens.
-        if (!it->empty() && (*it)[0] != '\0') args.push_back(std::move(*it));
-      }
-      LOG(INFO) << "Got " << args.size() << " arguments from boot message";
-    } else if (boot.recovery[0] != 0) {
-      LOG(ERROR) << "Bad boot message: \"" << boot_recovery << "\"";
-    }
-  }
-
-  // --- if that doesn't work, try the command file (if we have /cache).
-  if (args.size() == 1 && has_cache) {
-    std::string content;
-    if (ensure_path_mounted(COMMAND_FILE) == 0 &&
-        android::base::ReadFileToString(COMMAND_FILE, &content)) {
-      std::vector<std::string> tokens = android::base::Split(content, "\n");
-      // All the arguments in COMMAND_FILE are needed (unlike the BCB message,
-      // COMMAND_FILE doesn't use filename as the first argument).
-      for (auto it = tokens.begin(); it != tokens.end(); it++) {
-        // Skip empty and '\0'-filled tokens.
-        if (!it->empty() && (*it)[0] != '\0') args.push_back(std::move(*it));
-      }
-      LOG(INFO) << "Got " << args.size() << " arguments from " << COMMAND_FILE;
-    }
-  }
-
-  // Write the arguments (excluding the filename in args[0]) back into the
-  // bootloader control block. So the device will always boot into recovery to
-  // finish the pending work, until finish_recovery() is called.
-  std::vector<std::string> options(args.cbegin() + 1, args.cend());
-  if (!update_bootloader_message(options, &err)) {
-    LOG(ERROR) << "Failed to set BCB message: " << err;
-  }
-
-  return args;
-}
-
 // Set the BCB to reboot back into recovery (it won't resume the install from
 // sdcard though).
 static void set_sdcard_update_bootloader_message() {
@@ -921,21 +840,6 @@
   printf("%s=%s\n", key, name);
 }
 
-static std::string load_locale_from_cache() {
-    if (ensure_path_mounted(LOCALE_FILE) != 0) {
-        LOG(ERROR) << "Can't mount " << LOCALE_FILE;
-        return "";
-    }
-
-    std::string content;
-    if (!android::base::ReadFileToString(LOCALE_FILE, &content)) {
-        PLOG(ERROR) << "Can't read " << LOCALE_FILE;
-        return "";
-    }
-
-    return android::base::Trim(content);
-}
-
 void ui_print(const char* format, ...) {
     std::string buffer;
     va_list ap;
@@ -1059,15 +963,7 @@
   LOG(INFO) << log_content;
 }
 
-int start_recovery(int argc, char** argv) {
-  time_t start = time(nullptr);
-
-  printf("Starting recovery (pid %d) on %s", getpid(), ctime(&start));
-
-  load_volume_table();
-  has_cache = volume_for_mount_point(CACHE_ROOT) != nullptr;
-
-  std::vector<std::string> args = get_args(argc, argv);
+Device::BuiltinAction start_recovery(Device* device, const std::vector<std::string>& args) {
   std::vector<char*> args_to_parse(args.size());
   std::transform(args.cbegin(), args.cend(), args_to_parse.begin(),
                  [](const std::string& arg) { return const_cast<char*>(arg.c_str()); });
@@ -1097,7 +993,6 @@
   bool should_wipe_cache = false;
   bool should_wipe_ab = false;
   size_t wipe_package_size = 0;
-  bool show_text = false;
   bool sideload = false;
   bool sideload_auto_reboot = false;
   bool just_exit = false;
@@ -1112,7 +1007,7 @@
                             &option_index)) != -1) {
     switch (arg) {
       case 't':
-        show_text = true;
+        // Handled in recovery_main.cpp
         break;
       case 'x':
         just_exit = true;
@@ -1120,7 +1015,7 @@
       case 0: {
         std::string option = OPTIONS[option_index].name;
         if (option == "locale") {
-          locale = optarg;
+          // Handled in recovery_main.cpp
         } else if (option == "prompt_and_wipe_data") {
           should_prompt_and_wipe_data = true;
         } else if (option == "reason") {
@@ -1155,37 +1050,9 @@
     }
   }
 
-  if (locale.empty()) {
-    if (has_cache) {
-      locale = load_locale_from_cache();
-    }
-
-    if (locale.empty()) {
-      static constexpr const char* DEFAULT_LOCALE = "en-US";
-      locale = DEFAULT_LOCALE;
-    }
-  }
-
-  printf("locale is [%s]\n", locale.c_str());
   printf("stage is [%s]\n", stage.c_str());
   printf("reason is [%s]\n", reason);
 
-  Device* device = make_device();
-  if (android::base::GetBoolProperty("ro.boot.quiescent", false)) {
-    printf("Quiescent recovery mode.\n");
-    device->ResetUI(new StubRecoveryUI());
-  } else {
-    if (!device->GetUI()->Init(locale)) {
-      printf("Failed to initialize UI; using stub UI instead.\n");
-      device->ResetUI(new StubRecoveryUI());
-    }
-  }
-  ui = device->GetUI();
-
-  if (!has_cache) {
-    device->RemoveMenuItemForAction(Device::WIPE_CACHE);
-  }
-
   // Set background string to "installing security update" for security update,
   // otherwise set it to "installing system update".
   ui->SetSystemUpdateText(security_update);
@@ -1195,15 +1062,6 @@
     ui->SetStage(st_cur, st_max);
   }
 
-  ui->SetBackground(RecoveryUI::NONE);
-  if (show_text) ui->ShowText(true);
-
-  sehandle = selinux_android_file_context_handle();
-  selinux_android_set_sehandle(sehandle);
-  if (!sehandle) {
-    ui->Print("Warning: No file_contexts\n");
-  }
-
   device->StartRecovery();
 
   printf("Command:");
@@ -1353,25 +1211,5 @@
   // Save logs and clean up before rebooting or shutting down.
   finish_recovery();
 
-  switch (after) {
-    case Device::SHUTDOWN:
-      ui->Print("Shutting down...\n");
-      android::base::SetProperty(ANDROID_RB_PROPERTY, "shutdown,");
-      break;
-
-    case Device::REBOOT_BOOTLOADER:
-      ui->Print("Rebooting to bootloader...\n");
-      android::base::SetProperty(ANDROID_RB_PROPERTY, "reboot,bootloader");
-      break;
-
-    default:
-      ui->Print("Rebooting...\n");
-      reboot("reboot,");
-      break;
-  }
-  while (true) {
-    pause();
-  }
-  // Should be unreachable.
-  return EXIT_SUCCESS;
+  return after;
 }
diff --git a/private/recovery.h b/recovery.h
similarity index 80%
rename from private/recovery.h
rename to recovery.h
index 5b2ca4b..00e22da 100644
--- a/private/recovery.h
+++ b/recovery.h
@@ -16,4 +16,9 @@
 
 #pragma once
 
-int start_recovery(int argc, char** argv);
+#include <string>
+#include <vector>
+
+#include "device.h"
+
+Device::BuiltinAction start_recovery(Device* device, const std::vector<std::string>& args);
diff --git a/recovery_main.cpp b/recovery_main.cpp
index 3147511..e21c782 100644
--- a/recovery_main.cpp
+++ b/recovery_main.cpp
@@ -14,22 +14,57 @@
  * limitations under the License.
  */
 
+#include <errno.h>
+#include <fcntl.h>
+#include <getopt.h>
+#include <inttypes.h>
+#include <limits.h>
+#include <linux/fs.h>
+#include <stdarg.h>
 #include <stdio.h>
 #include <stdlib.h>
+#include <string.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <time.h>
 #include <unistd.h>
 
-#include <chrono>
+#include <algorithm>
+#include <string>
+#include <vector>
 
+#include <android-base/file.h>
 #include <android-base/logging.h>
+#include <android-base/properties.h>
+#include <android-base/strings.h>
+#include <bootloader_message/bootloader_message.h>
+#include <cutils/android_reboot.h>
 #include <private/android_logger.h> /* private pmsg functions */
+#include <selinux/android.h>
+#include <selinux/label.h>
+#include <selinux/selinux.h>
 
 #include "common.h"
+#include "device.h"
 #include "logging.h"
 #include "minadbd/minadbd.h"
 #include "otautil/paths.h"
-#include "private/recovery.h"
+#include "otautil/sysutil.h"
+#include "recovery.h"
+#include "roots.h"
+#include "stub_ui.h"
 #include "ui.h"
 
+static constexpr const char* COMMAND_FILE = "/cache/recovery/command";
+static constexpr const char* LOCALE_FILE = "/cache/recovery/last_locale";
+
+static constexpr const char* CACHE_ROOT = "/cache";
+
+bool has_cache = false;
+
+RecoveryUI* ui = nullptr;
+struct selabel_handle* sehandle;
+
 static void UiLogger(android::base::LogId /* id */, android::base::LogSeverity severity,
                      const char* /* tag */, const char* /* file */, unsigned int /* line */,
                      const char* message) {
@@ -41,6 +76,92 @@
   }
 }
 
+// command line args come from, in decreasing precedence:
+//   - the actual command line
+//   - the bootloader control block (one per line, after "recovery")
+//   - the contents of COMMAND_FILE (one per line)
+static std::vector<std::string> get_args(const int argc, char** const argv) {
+  CHECK_GT(argc, 0);
+
+  bootloader_message boot = {};
+  std::string err;
+  if (!read_bootloader_message(&boot, &err)) {
+    LOG(ERROR) << err;
+    // If fails, leave a zeroed bootloader_message.
+    boot = {};
+  }
+  stage = std::string(boot.stage);
+
+  if (boot.command[0] != 0) {
+    std::string boot_command = std::string(boot.command, sizeof(boot.command));
+    LOG(INFO) << "Boot command: " << boot_command;
+  }
+
+  if (boot.status[0] != 0) {
+    std::string boot_status = std::string(boot.status, sizeof(boot.status));
+    LOG(INFO) << "Boot status: " << boot_status;
+  }
+
+  std::vector<std::string> args(argv, argv + argc);
+
+  // --- if arguments weren't supplied, look in the bootloader control block
+  if (args.size() == 1) {
+    boot.recovery[sizeof(boot.recovery) - 1] = '\0';  // Ensure termination
+    std::string boot_recovery(boot.recovery);
+    std::vector<std::string> tokens = android::base::Split(boot_recovery, "\n");
+    if (!tokens.empty() && tokens[0] == "recovery") {
+      for (auto it = tokens.begin() + 1; it != tokens.end(); it++) {
+        // Skip empty and '\0'-filled tokens.
+        if (!it->empty() && (*it)[0] != '\0') args.push_back(std::move(*it));
+      }
+      LOG(INFO) << "Got " << args.size() << " arguments from boot message";
+    } else if (boot.recovery[0] != 0) {
+      LOG(ERROR) << "Bad boot message: \"" << boot_recovery << "\"";
+    }
+  }
+
+  // --- if that doesn't work, try the command file (if we have /cache).
+  if (args.size() == 1 && has_cache) {
+    std::string content;
+    if (ensure_path_mounted(COMMAND_FILE) == 0 &&
+        android::base::ReadFileToString(COMMAND_FILE, &content)) {
+      std::vector<std::string> tokens = android::base::Split(content, "\n");
+      // All the arguments in COMMAND_FILE are needed (unlike the BCB message,
+      // COMMAND_FILE doesn't use filename as the first argument).
+      for (auto it = tokens.begin(); it != tokens.end(); it++) {
+        // Skip empty and '\0'-filled tokens.
+        if (!it->empty() && (*it)[0] != '\0') args.push_back(std::move(*it));
+      }
+      LOG(INFO) << "Got " << args.size() << " arguments from " << COMMAND_FILE;
+    }
+  }
+
+  // Write the arguments (excluding the filename in args[0]) back into the
+  // bootloader control block. So the device will always boot into recovery to
+  // finish the pending work, until finish_recovery() is called.
+  std::vector<std::string> options(args.cbegin() + 1, args.cend());
+  if (!update_bootloader_message(options, &err)) {
+    LOG(ERROR) << "Failed to set BCB message: " << err;
+  }
+
+  return args;
+}
+
+static std::string load_locale_from_cache() {
+  if (ensure_path_mounted(LOCALE_FILE) != 0) {
+    LOG(ERROR) << "Can't mount " << LOCALE_FILE;
+    return "";
+  }
+
+  std::string content;
+  if (!android::base::ReadFileToString(LOCALE_FILE, &content)) {
+    PLOG(ERROR) << "Can't read " << LOCALE_FILE;
+    return "";
+  }
+
+  return android::base::Trim(content);
+}
+
 static void redirect_stdio(const char* filename) {
   int pipefd[2];
   if (pipe(pipefd) == -1) {
@@ -154,9 +275,108 @@
     return 0;
   }
 
+  time_t start = time(nullptr);
+
   // redirect_stdio should be called only in non-sideload mode. Otherwise we may have two logger
   // instances with different timestamps.
   redirect_stdio(Paths::Get().temporary_log_file().c_str());
 
-  return start_recovery(argc, argv);
+  printf("Starting recovery (pid %d) on %s", getpid(), ctime(&start));
+
+  load_volume_table();
+  has_cache = volume_for_mount_point(CACHE_ROOT) != nullptr;
+
+  std::vector<std::string> args = get_args(argc, argv);
+  std::vector<char*> args_to_parse(args.size());
+  std::transform(args.cbegin(), args.cend(), args_to_parse.begin(),
+                 [](const std::string& arg) { return const_cast<char*>(arg.c_str()); });
+
+  static constexpr struct option OPTIONS[] = {
+    { "locale", required_argument, nullptr, 0 },
+    { "show_text", no_argument, nullptr, 't' },
+    { nullptr, 0, nullptr, 0 },
+  };
+
+  bool show_text = false;
+  std::string locale;
+
+  int arg;
+  int option_index;
+  while ((arg = getopt_long(args_to_parse.size(), args_to_parse.data(), "", OPTIONS,
+                            &option_index)) != -1) {
+    switch (arg) {
+      case 't':
+        show_text = true;
+        break;
+      case 0: {
+        std::string option = OPTIONS[option_index].name;
+        if (option == "locale") {
+          locale = optarg;
+        }
+        break;
+      }
+    }
+  }
+
+  if (locale.empty()) {
+    if (has_cache) {
+      locale = load_locale_from_cache();
+    }
+
+    if (locale.empty()) {
+      static constexpr const char* DEFAULT_LOCALE = "en-US";
+      locale = DEFAULT_LOCALE;
+    }
+  }
+
+  printf("locale is [%s]\n", locale.c_str());
+
+  Device* device = make_device();
+  if (android::base::GetBoolProperty("ro.boot.quiescent", false)) {
+    printf("Quiescent recovery mode.\n");
+    device->ResetUI(new StubRecoveryUI());
+  } else {
+    if (!device->GetUI()->Init(locale)) {
+      printf("Failed to initialize UI; using stub UI instead.\n");
+      device->ResetUI(new StubRecoveryUI());
+    }
+  }
+  ui = device->GetUI();
+
+  if (!has_cache) {
+    device->RemoveMenuItemForAction(Device::WIPE_CACHE);
+  }
+
+  ui->SetBackground(RecoveryUI::NONE);
+  if (show_text) ui->ShowText(true);
+
+  sehandle = selinux_android_file_context_handle();
+  selinux_android_set_sehandle(sehandle);
+  if (!sehandle) {
+    ui->Print("Warning: No file_contexts\n");
+  }
+
+  Device::BuiltinAction after = start_recovery(device, args);
+
+  switch (after) {
+    case Device::SHUTDOWN:
+      ui->Print("Shutting down...\n");
+      android::base::SetProperty(ANDROID_RB_PROPERTY, "shutdown,");
+      break;
+
+    case Device::REBOOT_BOOTLOADER:
+      ui->Print("Rebooting to bootloader...\n");
+      android::base::SetProperty(ANDROID_RB_PROPERTY, "reboot,bootloader");
+      break;
+
+    default:
+      ui->Print("Rebooting...\n");
+      reboot("reboot,");
+      break;
+  }
+  while (true) {
+    pause();
+  }
+  // Should be unreachable.
+  return EXIT_SUCCESS;
 }