diff --git a/common.h b/common.h
index 319af3d..bda7935 100644
--- a/common.h
+++ b/common.h
@@ -20,6 +20,8 @@
 #include <stdio.h>
 #include <stdarg.h>
 
+#include <string>
+
 #define STRINGIFY(x) #x
 #define EXPAND(x) STRINGIFY(x)
 
@@ -29,7 +31,7 @@
 extern bool modified_flash;
 
 // The current stage, e.g. "1/2".
-extern const char* stage;
+extern std::string stage;
 
 // The reason argument provided in "--reason=".
 extern const char* reason;
diff --git a/device.cpp b/device.cpp
index e717ddd..6150186 100644
--- a/device.cpp
+++ b/device.cpp
@@ -60,7 +60,7 @@
   return menu_position < 0 ? NO_ACTION : MENU_ACTIONS[menu_position];
 }
 
-int Device::HandleMenuKey(int key, int visible) {
+int Device::HandleMenuKey(int key, bool visible) {
   if (!visible) {
     return kNoAction;
   }
diff --git a/device.h b/device.h
index 08dfcdf..639e2bf 100644
--- a/device.h
+++ b/device.h
@@ -20,96 +20,90 @@
 #include "ui.h"
 
 class Device {
-  public:
-    explicit Device(RecoveryUI* ui) : ui_(ui) { }
-    virtual ~Device() { }
+ public:
+  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.
-    virtual RecoveryUI* GetUI() { return ui_; }
+  // 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.
+  virtual RecoveryUI* GetUI() {
+    return 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() { };
+  // 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() {};
 
-    // 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.)
-    //
-    // 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.)
-    //
-    // visible is true if the menu is visible.
-    //
-    // Return one of the defined constants below in order to:
-    //
-    //   - move the menu highlight (kHighlight{Up,Down})
-    //   - invoke the highlighted item (kInvokeItem)
-    //   - do nothing (kNoAction)
-    //   - invoke a specific action (a menu position: any non-negative number)
-    virtual int HandleMenuKey(int key, int visible);
+  // 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.)
+  //
+  // '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.)
+  //
+  // '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})
+  //   - invoke the highlighted item (kInvokeItem)
+  //   - do nothing (kNoAction)
+  //   - invoke a specific action (a menu position: any non-negative number)
+  virtual int HandleMenuKey(int key, bool visible);
 
-    enum BuiltinAction {
-        NO_ACTION = 0,
-        REBOOT = 1,
-        APPLY_SDCARD = 2,
-        // APPLY_CACHE was 3.
-        APPLY_ADB_SIDELOAD = 4,
-        WIPE_DATA = 5,
-        WIPE_CACHE = 6,
-        REBOOT_BOOTLOADER = 7,
-        SHUTDOWN = 8,
-        VIEW_RECOVERY_LOGS = 9,
-        MOUNT_SYSTEM = 10,
-        RUN_GRAPHICS_TEST = 11,
-    };
+  enum BuiltinAction {
+    NO_ACTION = 0,
+    REBOOT = 1,
+    APPLY_SDCARD = 2,
+    // APPLY_CACHE was 3.
+    APPLY_ADB_SIDELOAD = 4,
+    WIPE_DATA = 5,
+    WIPE_CACHE = 6,
+    REBOOT_BOOTLOADER = 7,
+    SHUTDOWN = 8,
+    VIEW_RECOVERY_LOGS = 9,
+    MOUNT_SYSTEM = 10,
+    RUN_GRAPHICS_TEST = 11,
+  };
 
-    // Return the list of menu items (an array of strings,
-    // NULL-terminated).  The menu_position passed to InvokeMenuItem
-    // will correspond to the indexes into this array.
-    virtual const char* const* GetMenuItems();
+  // Return the list of menu items (an array of strings, NULL-terminated). The menu_position passed
+  // to InvokeMenuItem will correspond to the indexes into this array.
+  virtual const char* const* GetMenuItems();
 
-    // Perform a recovery action selected from the menu.
-    // 'menu_position' will be the item number of the selected menu
-    // item, or a non-negative number returned from
-    // device_handle_key().  The menu will be hidden when this is
-    // called; implementations can call ui_print() to print
-    // information to the screen.  If the menu position is one of the
-    // builtin actions, you can just return the corresponding enum
-    // value.  If it is an action specific to your device, you
-    // actually perform it here and return NO_ACTION.
-    virtual BuiltinAction InvokeMenuItem(int menu_position);
+  // Perform a recovery action selected from the menu. 'menu_position' will be the item number of
+  // the selected menu item, or a non-negative number returned from HandleMenuKey(). The menu will
+  // be hidden when this is called; implementations can call ui_print() to print information to the
+  // screen. If the menu position is one of the builtin actions, you can just return the
+  // corresponding enum value. If it is an action specific to your device, you actually perform it
+  // here and return NO_ACTION.
+  virtual BuiltinAction InvokeMenuItem(int menu_position);
 
-    static const int kNoAction = -1;
-    static const int kHighlightUp = -2;
-    static const int kHighlightDown = -3;
-    static const int kInvokeItem = -4;
+  static const int kNoAction = -1;
+  static const int kHighlightUp = -2;
+  static const int kHighlightDown = -3;
+  static const int kInvokeItem = -4;
 
-    // Called before and after we do a wipe data/factory reset operation,
-    // either via a reboot from the main system with the --wipe_data flag,
-    // or when the user boots into recovery image manually and selects the
-    // option from the menu, to perform whatever device-specific wiping
-    // actions are needed.
-    // Return true on success; returning false from PreWipeData will prevent
-    // the regular wipe, and returning false from PostWipeData will cause
-    // the wipe to be considered a failure.
-    virtual bool PreWipeData() { return true; }
-    virtual bool PostWipeData() { return true; }
+  // Called before and after we do a wipe data/factory reset operation, either via a reboot from the
+  // main system with the --wipe_data flag, or when the user boots into recovery image manually and
+  // selects the option from the menu, to perform whatever device-specific wiping actions as needed.
+  // Returns true on success; returning false from PreWipeData will prevent the regular wipe, and
+  // returning false from PostWipeData will cause the wipe to be considered a failure.
+  virtual bool PreWipeData() {
+    return true;
+  }
 
-  private:
-    RecoveryUI* ui_;
+  virtual bool PostWipeData() {
+    return true;
+  }
+
+ private:
+  RecoveryUI* ui_;
 };
 
-// The device-specific library must define this function (or the
-// default one will be used, if there is no device-specific library).
-// It returns the Device object that recovery should use.
+// The device-specific library must define this function (or the default one will be used, if there
+// is no device-specific library). It returns the Device object that recovery should use.
 Device* make_device();
 
 #endif  // _DEVICE_H
diff --git a/recovery.cpp b/recovery.cpp
index 5c60ce6..b5cd348 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -83,7 +83,6 @@
   { "sideload_auto_reboot", no_argument, NULL, 'a' },
   { "just_exit", no_argument, NULL, 'x' },
   { "locale", required_argument, NULL, 'l' },
-  { "stages", required_argument, NULL, 'g' },
   { "shutdown_after", no_argument, NULL, 'p' },
   { "reason", required_argument, NULL, 'r' },
   { "security", no_argument, NULL, 'e'},
@@ -129,7 +128,7 @@
 
 RecoveryUI* ui = nullptr;
 bool modified_flash = false;
-const char* stage = nullptr;
+std::string stage;
 const char* reason = nullptr;
 struct selabel_handle* sehandle;
 
@@ -309,7 +308,7 @@
     // If fails, leave a zeroed bootloader_message.
     boot = {};
   }
-  stage = strndup(boot.stage, sizeof(boot.stage));
+  stage = std::string(boot.stage);
 
   if (boot.command[0] != 0) {
     std::string boot_command = std::string(boot.command, sizeof(boot.command));
@@ -608,54 +607,55 @@
   return (result == 0);
 }
 
-static int
-get_menu_selection(const char* const * headers, const char* const * items,
-                   int menu_only, int initial_selection, Device* device) {
-    // throw away keys pressed previously, so user doesn't
-    // accidentally trigger menu items.
-    ui->FlushKeys();
+// Display a menu with the specified 'headers' and 'items'. Device specific HandleMenuKey() may
+// return a positive number beyond the given range. Caller sets 'menu_only' to true to ensure only
+// a menu item gets selected. 'initial_selection' controls the initial cursor location.
+static int get_menu_selection(const char* const* headers, const char* const* items, bool menu_only,
+                              int initial_selection, Device* device) {
+  // Throw away keys pressed previously, so user doesn't accidentally trigger menu items.
+  ui->FlushKeys();
 
-    ui->StartMenu(headers, items, initial_selection);
-    int selected = initial_selection;
-    int chosen_item = -1;
+  ui->StartMenu(headers, items, initial_selection);
+  int selected = initial_selection;
+  int chosen_item = -1;
 
-    while (chosen_item < 0) {
-        int key = ui->WaitKey();
-        int visible = ui->IsTextVisible();
+  while (chosen_item < 0) {
+    int key = ui->WaitKey();
 
-        if (key == -1) {   // ui_wait_key() timed out
-            if (ui->WasTextEverVisible()) {
-                continue;
-            } else {
-                LOG(INFO) << "timed out waiting for key input; rebooting.";
-                ui->EndMenu();
-                return 0; // XXX fixme
-            }
-        }
-
-        int action = device->HandleMenuKey(key, visible);
-
-        if (action < 0) {
-            switch (action) {
-                case Device::kHighlightUp:
-                    selected = ui->SelectMenu(--selected);
-                    break;
-                case Device::kHighlightDown:
-                    selected = ui->SelectMenu(++selected);
-                    break;
-                case Device::kInvokeItem:
-                    chosen_item = selected;
-                    break;
-                case Device::kNoAction:
-                    break;
-            }
-        } else if (!menu_only) {
-            chosen_item = action;
-        }
+    if (key == -1) {  // ui_wait_key() timed out
+      if (ui->WasTextEverVisible()) {
+        continue;
+      } else {
+        LOG(INFO) << "timed out waiting for key input; rebooting.";
+        ui->EndMenu();
+        return 0;  // XXX fixme
+      }
     }
 
-    ui->EndMenu();
-    return chosen_item;
+    bool visible = ui->IsTextVisible();
+    int action = device->HandleMenuKey(key, visible);
+
+    if (action < 0) {
+      switch (action) {
+        case Device::kHighlightUp:
+          selected = ui->SelectMenu(--selected);
+          break;
+        case Device::kHighlightDown:
+          selected = ui->SelectMenu(++selected);
+          break;
+        case Device::kInvokeItem:
+          chosen_item = selected;
+          break;
+        case Device::kNoAction:
+          break;
+      }
+    } else if (!menu_only) {
+      chosen_item = action;
+    }
+  }
+
+  ui->EndMenu();
+  return chosen_item;
 }
 
 // Returns the selected filename, or an empty string.
@@ -700,7 +700,7 @@
 
   int chosen_item = 0;
   while (true) {
-    chosen_item = get_menu_selection(headers, entries, 1, chosen_item, device);
+    chosen_item = get_menu_selection(headers, entries, true, chosen_item, device);
 
     const std::string& item = zips[chosen_item];
     if (chosen_item == 0) {
@@ -727,7 +727,7 @@
     const char* headers[] = { question1, question2, NULL };
     const char* items[] = { " No", " Yes", NULL };
 
-    int chosen_item = get_menu_selection(headers, items, 1, 0, device);
+    int chosen_item = get_menu_selection(headers, items, true, 0, device);
     return (chosen_item == 1);
 }
 
@@ -750,25 +750,25 @@
 }
 
 static bool prompt_and_wipe_data(Device* device) {
-    const char* const headers[] = {
-        "Boot halted, user data is corrupt",
-        "Wipe all user data to recover",
-        NULL
-    };
-    const char* const items[] = {
-        "Retry boot",
-        "Wipe user data",
-        NULL
-    };
-    for (;;) {
-        int chosen_item = get_menu_selection(headers, items, 1, 0, device);
-        if (chosen_item != 1) {
-            return true; // Just reboot, no wipe; not a failure, user asked for it
-        }
-        if (ask_to_wipe_data(device)) {
-            return wipe_data(device);
-        }
+  const char* const headers[] = {
+    "Boot halted, user data is corrupt",
+    "Wipe all user data to recover",
+    NULL
+  };
+  const char* const items[] = {
+    "Retry boot",
+    "Wipe user data",
+    NULL
+  };
+  for (;;) {
+    int chosen_item = get_menu_selection(headers, items, true, 0, device);
+    if (chosen_item != 1) {
+      return true;  // Just reboot, no wipe; not a failure, user asked for it
     }
+    if (ask_to_wipe_data(device)) {
+      return wipe_data(device);
+    }
+  }
 }
 
 // Return true on success.
@@ -922,98 +922,94 @@
 }
 
 static void choose_recovery_file(Device* device) {
-    // "Back" + KEEP_LOG_COUNT * 2 + terminating nullptr entry
-    char* entries[1 + KEEP_LOG_COUNT * 2 + 1];
-    memset(entries, 0, sizeof(entries));
-
-    unsigned int n = 0;
-
-    if (has_cache) {
-        // Add LAST_LOG_FILE + LAST_LOG_FILE.x
-        // Add LAST_KMSG_FILE + LAST_KMSG_FILE.x
-        for (int i = 0; i < KEEP_LOG_COUNT; i++) {
-            char* log_file;
-            int ret;
-            ret = (i == 0) ? asprintf(&log_file, "%s", LAST_LOG_FILE) :
-                    asprintf(&log_file, "%s.%d", LAST_LOG_FILE, i);
-            if (ret == -1) {
-                // memory allocation failure - return early. Should never happen.
-                return;
-            }
-            if ((ensure_path_mounted(log_file) != 0) || (access(log_file, R_OK) == -1)) {
-                free(log_file);
-            } else {
-                entries[n++] = log_file;
-            }
-
-            char* kmsg_file;
-            ret = (i == 0) ? asprintf(&kmsg_file, "%s", LAST_KMSG_FILE) :
-                    asprintf(&kmsg_file, "%s.%d", LAST_KMSG_FILE, i);
-            if (ret == -1) {
-                // memory allocation failure - return early. Should never happen.
-                return;
-            }
-            if ((ensure_path_mounted(kmsg_file) != 0) || (access(kmsg_file, R_OK) == -1)) {
-                free(kmsg_file);
-            } else {
-                entries[n++] = kmsg_file;
-            }
+  std::vector<std::string> entries;
+  if (has_cache) {
+    for (int i = 0; i < KEEP_LOG_COUNT; i++) {
+      auto add_to_entries = [&](const char* filename) {
+        std::string log_file(filename);
+        if (i > 0) {
+          log_file += "." + std::to_string(i);
         }
+
+        if (ensure_path_mounted(log_file.c_str()) == 0 && access(log_file.c_str(), R_OK) == 0) {
+          entries.push_back(std::move(log_file));
+        }
+      };
+
+      // Add LAST_LOG_FILE + LAST_LOG_FILE.x
+      add_to_entries(LAST_LOG_FILE);
+
+      // Add LAST_KMSG_FILE + LAST_KMSG_FILE.x
+      add_to_entries(LAST_KMSG_FILE);
+    }
+  } else {
+    // If cache partition is not found, view /tmp/recovery.log instead.
+    if (access(TEMPORARY_LOG_FILE, R_OK) == -1) {
+      return;
     } else {
-        // If cache partition is not found, view /tmp/recovery.log instead.
-        if (access(TEMPORARY_LOG_FILE, R_OK) == -1) {
-            return;
-        } else{
-            entries[n++] = strdup(TEMPORARY_LOG_FILE);
-        }
+      entries.push_back(TEMPORARY_LOG_FILE);
     }
+  }
 
-    entries[n++] = strdup("Back");
+  entries.push_back("Back");
 
-    const char* headers[] = { "Select file to view", nullptr };
+  std::vector<const char*> menu_entries(entries.size());
+  std::transform(entries.cbegin(), entries.cend(), menu_entries.begin(),
+                 [](const std::string& entry) { return entry.c_str(); });
+  menu_entries.push_back(nullptr);
 
-    int chosen_item = 0;
-    while (true) {
-        chosen_item = get_menu_selection(headers, entries, 1, chosen_item, device);
-        if (strcmp(entries[chosen_item], "Back") == 0) break;
+  const char* headers[] = { "Select file to view", nullptr };
 
-        ui->ShowFile(entries[chosen_item]);
-    }
+  int chosen_item = 0;
+  while (true) {
+    chosen_item = get_menu_selection(headers, menu_entries.data(), true, chosen_item, device);
+    if (entries[chosen_item] == "Back") break;
 
-    for (size_t i = 0; i < (sizeof(entries) / sizeof(*entries)); i++) {
-        free(entries[i]);
-    }
+    ui->ShowFile(entries[chosen_item].c_str());
+  }
 }
 
-static void run_graphics_test(Device* device) {
-    // Switch to graphics screen.
-    ui->ShowText(false);
+static void run_graphics_test() {
+  // Switch to graphics screen.
+  ui->ShowText(false);
 
-    ui->SetProgressType(RecoveryUI::INDETERMINATE);
-    ui->SetBackground(RecoveryUI::INSTALLING_UPDATE);
-    sleep(1);
+  ui->SetProgressType(RecoveryUI::INDETERMINATE);
+  ui->SetBackground(RecoveryUI::INSTALLING_UPDATE);
+  sleep(1);
 
-    ui->SetBackground(RecoveryUI::ERROR);
-    sleep(1);
+  ui->SetBackground(RecoveryUI::ERROR);
+  sleep(1);
 
-    ui->SetBackground(RecoveryUI::NO_COMMAND);
-    sleep(1);
+  ui->SetBackground(RecoveryUI::NO_COMMAND);
+  sleep(1);
 
-    ui->SetBackground(RecoveryUI::ERASING);
-    sleep(1);
+  ui->SetBackground(RecoveryUI::ERASING);
+  sleep(1);
 
-    ui->SetBackground(RecoveryUI::INSTALLING_UPDATE);
+  // Calling SetBackground() after SetStage() to trigger a redraw.
+  ui->SetStage(1, 3);
+  ui->SetBackground(RecoveryUI::INSTALLING_UPDATE);
+  sleep(1);
+  ui->SetStage(2, 3);
+  ui->SetBackground(RecoveryUI::INSTALLING_UPDATE);
+  sleep(1);
+  ui->SetStage(3, 3);
+  ui->SetBackground(RecoveryUI::INSTALLING_UPDATE);
+  sleep(1);
 
-    ui->SetProgressType(RecoveryUI::DETERMINATE);
-    ui->ShowProgress(1.0, 10.0);
-    float fraction = 0.0;
-    for (size_t i = 0; i < 100; ++i) {
-      fraction += .01;
-      ui->SetProgress(fraction);
-      usleep(100000);
-    }
+  ui->SetStage(-1, -1);
+  ui->SetBackground(RecoveryUI::INSTALLING_UPDATE);
 
-    ui->ShowText(true);
+  ui->SetProgressType(RecoveryUI::DETERMINATE);
+  ui->ShowProgress(1.0, 10.0);
+  float fraction = 0.0;
+  for (size_t i = 0; i < 100; ++i) {
+    fraction += .01;
+    ui->SetProgress(fraction);
+    usleep(100000);
+  }
+
+  ui->ShowText(true);
 }
 
 // How long (in seconds) we wait for the fuse-provided package file to
@@ -1115,7 +1111,7 @@
         }
         ui->SetProgressType(RecoveryUI::EMPTY);
 
-        int chosen_item = get_menu_selection(nullptr, device->GetMenuItems(), 0, 0, device);
+        int chosen_item = get_menu_selection(nullptr, device->GetMenuItems(), false, 0, device);
 
         // device-specific code may take some action here.  It may
         // return one of the core actions handled in the switch
@@ -1181,7 +1177,7 @@
                 break;
 
             case Device::RUN_GRAPHICS_TEST:
-                run_graphics_test(device);
+                run_graphics_test();
                 break;
 
             case Device::MOUNT_SYSTEM:
@@ -1422,14 +1418,6 @@
         case 'a': sideload = true; sideload_auto_reboot = true; break;
         case 'x': just_exit = true; break;
         case 'l': locale = optarg; break;
-        case 'g': {
-            if (stage == NULL || *stage == '\0') {
-                char buffer[20] = "1/";
-                strncat(buffer, optarg, sizeof(buffer)-3);
-                stage = strdup(buffer);
-            }
-            break;
-        }
         case 'p': shutdown_after = true; break;
         case 'r': reason = optarg; break;
         case 'e': security_update = true; break;
@@ -1461,7 +1449,7 @@
     }
 
     printf("locale is [%s]\n", locale.c_str());
-    printf("stage is [%s]\n", stage);
+    printf("stage is [%s]\n", stage.c_str());
     printf("reason is [%s]\n", reason);
 
     Device* device = make_device();
@@ -1476,7 +1464,7 @@
     ui->SetSystemUpdateText(security_update);
 
     int st_cur, st_max;
-    if (stage != NULL && sscanf(stage, "%d/%d", &st_cur, &st_max) == 2) {
+    if (!stage.empty() && sscanf(stage.c_str(), "%d/%d", &st_cur, &st_max) == 2) {
         ui->SetStage(st_cur, st_max);
     }