6 files changed, 26 insertions, 29 deletions

diff --git a/device.h b/device.h
index 9510fbedb..bf148f576 100644
--- a/device.h
+++ b/device.h
@@ -19,6 +19,7 @@
 
 #include <stddef.h>
 
+#include <memory>
 #include <string>
 #include <vector>
 
@@ -51,17 +52,15 @@ class Device {
   explicit Device(RecoveryUI* ui) : ui_(ui) {}
   virtual ~Device() {}
 
-  // Called to obtain the UI object that should be used to display the recovery user interface for
-  // this device. You should not have called Init() on the UI object already, the caller will do
-  // that after this method returns.
+  // Returns a raw pointer to the RecoveryUI object.
   virtual RecoveryUI* GetUI() {
-    return ui_;
+    return ui_.get();
   }
 
-  // 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;
+  // Resets the UI object to the given UI. Used to override the default UI in case initialization
+  // failed, or we want a different UI for some reason. The device object will take the ownership.
+  virtual void ResetUI(RecoveryUI* ui) {
+    ui_.reset(ui);
   }
 
   // Called when recovery starts up (after the UI has been obtained and initialized and after the
@@ -70,16 +69,15 @@ class Device {
 
   // Called from the main thread when recovery is at the main menu and waiting for input, and a key
   // is pressed. (Note that "at" the main menu does not necessarily mean the menu is visible;
-  // recovery will be at the main menu with it invisible after an unsuccessful operation [ie OTA
-  // package failure], or if recovery is started with no command.)
+  // recovery will be at the main menu with it invisible after an unsuccessful operation, such as
+  // failed to install an OTA package, or if recovery is started with no command.)
   //
   // 'key' is the code of the key just pressed. (You can call IsKeyPressed() on the RecoveryUI
-  // object you returned from GetUI if you want to find out if other keys are held down.)
+  // object you returned from GetUI() if you want to find out if other keys are held down.)
   //
   // 'visible' is true if the menu is visible.
   //
   // Returns one of the defined constants below in order to:
-  //
   //   - move the menu highlight (kHighlight{Up,Down}: negative value)
   //   - invoke the highlighted item (kInvokeItem: negative value)
   //   - do nothing (kNoAction: negative value)
@@ -87,7 +85,7 @@ class Device {
   virtual int HandleMenuKey(int key, bool visible);
 
   // Returns the list of menu items (a vector of strings). The menu_position passed to
-  // InvokeMenuItem will correspond to the indexes into this array.
+  // InvokeMenuItem() will correspond to the indexes into this array.
   virtual const std::vector<std::string>& GetMenuItems();
 
   // Performs a recovery action selected from the menu. 'menu_position' will be the index of the
@@ -112,7 +110,8 @@ class Device {
   }
 
  private:
-  RecoveryUI* ui_;
+  // The RecoveryUI object that should be used to display the user interface for this device.
+  std::unique_ptr<RecoveryUI> ui_;
 };
 
 // The device-specific library must define this function (or the default one will be used, if there
diff --git a/recovery.cpp b/recovery.cpp
index 5d14e7329..38784b0b3 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -230,8 +230,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(Device* device) {
-  std::string locale = device->GetUI()->GetLocale();
+static void finish_recovery() {
+  std::string locale = ui->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) {
@@ -808,7 +808,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(device);
+    finish_recovery();
     switch (status) {
       case INSTALL_SUCCESS:
       case INSTALL_NONE:
@@ -1193,16 +1193,14 @@ int start_recovery(int argc, char** argv) {
   Device* device = make_device();
   if (android::base::GetBoolProperty("ro.boot.quiescent", false)) {
     printf("Quiescent recovery mode.\n");
-    ui = new StubRecoveryUI();
+    device->ResetUI(new StubRecoveryUI());
   } else {
-    ui = device->GetUI();
-
-    if (!ui->Init(locale)) {
-      printf("Failed to initialize UI, use stub UI instead.\n");
-      ui = new StubRecoveryUI();
+    if (!device->GetUI()->Init(locale)) {
+      printf("Failed to initialize UI; using stub UI instead.\n");
+      device->ResetUI(new StubRecoveryUI());
     }
   }
-  device->SetUI(ui);
+  ui = device->GetUI();
 
   // Set background string to "installing security update" for security update,
   // otherwise set it to "installing system update".
@@ -1369,7 +1367,7 @@ int start_recovery(int argc, char** argv) {
   }
 
   // Save logs and clean up before rebooting or shutting down.
-  finish_recovery(device);
+  finish_recovery();
 
   switch (after) {
     case Device::SHUTDOWN:
diff --git a/screen_ui.cpp b/screen_ui.cpp
index c0fb2cfba..f5dadf7f0 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -751,7 +751,7 @@ bool ScreenRecoveryUI::Init(const std::string& locale) {
   return true;
 }
 
-std::string ScreenRecoveryUI::GetLocale() {
+std::string ScreenRecoveryUI::GetLocale() const {
   return locale_;
 }
 
diff --git a/screen_ui.h b/screen_ui.h
index 293696d22..2d6b621d5 100644
--- a/screen_ui.h
+++ b/screen_ui.h
@@ -114,7 +114,7 @@ class ScreenRecoveryUI : public RecoveryUI {
   explicit ScreenRecoveryUI(bool scrollable_menu);
 
   bool Init(const std::string& locale) override;
-  std::string GetLocale() override;
+  std::string GetLocale() const override;
 
   // overall recovery state ("background image")
   void SetBackground(Icon icon) override;
diff --git a/stub_ui.h b/stub_ui.h
index fddf4e71e..67c338e99 100644
--- a/stub_ui.h
+++ b/stub_ui.h
@@ -28,7 +28,7 @@ class StubRecoveryUI : public RecoveryUI {
  public:
   StubRecoveryUI() = default;
 
-  std::string GetLocale() override {
+  std::string GetLocale() const override {
     return "";
   }
   void SetBackground(Icon /* icon */) override {}
diff --git a/ui.h b/ui.h
index f8677902e..39284268d 100644
--- a/ui.h
+++ b/ui.h
@@ -57,7 +57,7 @@ class RecoveryUI {
   // the given locale. Returns true on success.
   virtual bool Init(const std::string& locale);
 
-  virtual std::string GetLocale() = 0;
+  virtual std::string GetLocale() const = 0;
 
   // Shows a stage indicator. Called immediately after Init().
   virtual void SetStage(int current, int max) = 0;