6 files changed, 23 insertions, 5 deletions

diff --git a/device.h b/device.h
index 8788b2d14..9510fbedb 100644
--- a/device.h
+++ b/device.h
@@ -58,6 +58,12 @@ class Device {
     return ui_;
   }
 
+  // Sets the UI object to the given UI. Used to override the default UI in case initialization
+  // failed, or we want a stub for some reason.
+  virtual void SetUI(RecoveryUI* ui) {
+    ui_ = ui;
+  }
+
   // Called when recovery starts up (after the UI has been obtained and initialized and after the
   // arguments have been parsed, but before anything else).
   virtual void StartRecovery() {};
diff --git a/recovery.cpp b/recovery.cpp
index 0ab34197f..37298d792 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -89,7 +89,6 @@ static constexpr const char* SDCARD_ROOT = "/sdcard";
 // 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 std::string locale;
 static bool has_cache = false;
 
 RecoveryUI* ui = nullptr;
@@ -232,7 +231,8 @@ static void set_sdcard_update_bootloader_message() {
 // Clear the recovery command and prepare to boot a (hopefully working) system,
 // copy our log file to cache as well (for the system to read). This function is
 // idempotent: call it as many times as you like.
-static void finish_recovery() {
+static void finish_recovery(Device* device) {
+  std::string locale = device->GetUI()->GetLocale();
   // Save the locale to cache, so if recovery is next started up without a '--locale' argument
   // (e.g., directly from the bootloader) it will use the last-known locale.
   if (!locale.empty() && has_cache) {
@@ -809,7 +809,7 @@ static int apply_from_sdcard(Device* device, bool* wipe_cache) {
 // which is to reboot or shutdown depending on if the --shutdown_after flag was passed to recovery.
 static Device::BuiltinAction prompt_and_wait(Device* device, int status) {
   for (;;) {
-    finish_recovery();
+    finish_recovery(device);
     switch (status) {
       case INSTALL_SUCCESS:
       case INSTALL_NONE:
@@ -897,7 +897,7 @@ static Device::BuiltinAction prompt_and_wait(Device* device, int status) {
 
       case Device::RUN_LOCALE_TEST: {
         ScreenRecoveryUI* screen_ui = static_cast<ScreenRecoveryUI*>(ui);
-        screen_ui->CheckBackgroundTextImages(locale);
+        screen_ui->CheckBackgroundTextImages(screen_ui->GetLocale());
         break;
       }
       case Device::MOUNT_SYSTEM:
@@ -1125,6 +1125,7 @@ int start_recovery(int argc, char** argv) {
   bool shutdown_after = false;
   int retry_count = 0;
   bool security_update = false;
+  std::string locale;
 
   int arg;
   int option_index;
@@ -1202,6 +1203,7 @@ int start_recovery(int argc, char** argv) {
       ui = new StubRecoveryUI();
     }
   }
+  device->SetUI(ui);
 
   // Set background string to "installing security update" for security update,
   // otherwise set it to "installing system update".
@@ -1368,7 +1370,7 @@ int start_recovery(int argc, char** argv) {
   }
 
   // Save logs and clean up before rebooting or shutting down.
-  finish_recovery();
+  finish_recovery(device);
 
   switch (after) {
     case Device::SHUTDOWN:
diff --git a/screen_ui.cpp b/screen_ui.cpp
index 7ae81e55f..4a1a5b99c 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -748,6 +748,10 @@ bool ScreenRecoveryUI::Init(const std::string& locale) {
   return true;
 }
 
+std::string ScreenRecoveryUI::GetLocale() {
+  return locale_;
+}
+
 void ScreenRecoveryUI::LoadAnimation() {
   std::unique_ptr<DIR, decltype(&closedir)> dir(opendir("/res/images"), closedir);
   dirent* de;
diff --git a/screen_ui.h b/screen_ui.h
index fb811ce70..3e391beba 100644
--- a/screen_ui.h
+++ b/screen_ui.h
@@ -114,6 +114,7 @@ class ScreenRecoveryUI : public RecoveryUI {
   explicit ScreenRecoveryUI(bool scrollable_menu);
 
   bool Init(const std::string& locale) override;
+  std::string GetLocale() override;
 
   // overall recovery state ("background image")
   void SetBackground(Icon icon) override;
diff --git a/stub_ui.h b/stub_ui.h
index 2ccd49115..fddf4e71e 100644
--- a/stub_ui.h
+++ b/stub_ui.h
@@ -28,6 +28,9 @@ class StubRecoveryUI : public RecoveryUI {
  public:
   StubRecoveryUI() = default;
 
+  std::string GetLocale() override {
+    return "";
+  }
   void SetBackground(Icon /* icon */) override {}
   void SetSystemUpdateText(bool /* security_update */) override {}
 
diff --git a/ui.h b/ui.h
index 35cc36e70..f8677902e 100644
--- a/ui.h
+++ b/ui.h
@@ -57,6 +57,8 @@ class RecoveryUI {
   // the given locale. Returns true on success.
   virtual bool Init(const std::string& locale);
 
+  virtual std::string GetLocale() = 0;
+
   // Shows a stage indicator. Called immediately after Init().
   virtual void SetStage(int current, int max) = 0;