diff --git a/bootloader_message/bootloader_message.cpp b/bootloader_message/bootloader_message.cpp
index e684abb..b15a9b9 100644
--- a/bootloader_message/bootloader_message.cpp
+++ b/bootloader_message/bootloader_message.cpp
@@ -20,6 +20,7 @@
 #include <fcntl.h>
 #include <string.h>
 
+#include <optional>
 #include <string>
 #include <string_view>
 #include <vector>
@@ -37,7 +38,7 @@
 using android::fs_mgr::Fstab;
 using android::fs_mgr::ReadDefaultFstab;
 
-static std::string g_misc_device_for_test;
+static std::optional<std::string> g_misc_device_for_test;
 
 // Exposed for test purpose.
 void SetMiscBlockDeviceForTest(std::string_view misc_device) {
@@ -45,8 +46,8 @@
 }
 
 static std::string get_misc_blk_device(std::string* err) {
-  if (!g_misc_device_for_test.empty()) {
-    return g_misc_device_for_test;
+  if (g_misc_device_for_test.has_value() && !g_misc_device_for_test->empty()) {
+    return *g_misc_device_for_test;
   }
   Fstab fstab;
   if (!ReadDefaultFstab(&fstab)) {
diff --git a/common.h b/common.h
index a524a41..128a69d 100644
--- a/common.h
+++ b/common.h
@@ -18,21 +18,8 @@
 
 #include <string>
 
-// Not using the command-line defined macro here because this header could be included by
-// device-specific recovery libraries. We static assert the value consistency in recovery.cpp.
-static constexpr int kRecoveryApiVersion = 3;
-
-class RecoveryUI;
-struct selabel_handle;
-
-extern struct selabel_handle* sehandle;
-extern RecoveryUI* ui;
-extern bool has_cache;
-
 // The current stage, e.g. "1/2".
 extern std::string stage;
 
 // The reason argument provided in "--reason=".
 extern const char* reason;
-
-bool is_ro_debuggable();
diff --git a/install/install.cpp b/install/install.cpp
index 09b8839..9d67b01 100644
--- a/install/install.cpp
+++ b/install/install.cpp
@@ -60,7 +60,8 @@
 using namespace std::chrono_literals;
 
 static constexpr int kRecoveryApiVersion = 3;
-// Assert the version defined in code and in Android.mk are consistent.
+// We define RECOVERY_API_VERSION in Android.mk, which will be picked up by build system and packed
+// 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.");
 
 // Default allocation of progress bar segments to operations
@@ -621,6 +622,8 @@
   InstallResult result;
   std::vector<std::string> log_buffer;
 
+  ui->Print("Supported API: %d\n", kRecoveryApiVersion);
+
   ui->Print("Finding update package...\n");
   LOG(INFO) << "Update package id: " << package_id;
   if (!package) {
diff --git a/otautil/include/otautil/logging.h b/otautil/include/otautil/logging.h
index 6083497..4462eca 100644
--- a/otautil/include/otautil/logging.h
+++ b/otautil/include/otautil/logging.h
@@ -53,7 +53,7 @@
 void check_and_fclose(FILE* fp, const std::string& name);
 
 void copy_log_file_to_pmsg(const std::string& source, const std::string& destination);
-void copy_logs(bool save_current_log, bool has_cache, const selabel_handle* sehandle);
+void copy_logs(bool save_current_log);
 void reset_tmplog_offset();
 
 void save_kernel_log(const char* destination);
diff --git a/otautil/include/otautil/roots.h b/otautil/include/otautil/roots.h
index 2ab3f45..92ee756 100644
--- a/otautil/include/otautil/roots.h
+++ b/otautil/include/otautil/roots.h
@@ -53,3 +53,6 @@
 // Ensure that all and only the volumes that packages expect to find
 // mounted (/tmp and /cache) are mounted.  Returns 0 on success.
 int setup_install_mounts();
+
+// Returns true if there is /cache in the volumes.
+bool HasCache();
diff --git a/otautil/logging.cpp b/otautil/logging.cpp
index 484f115..3db0e8a 100644
--- a/otautil/logging.cpp
+++ b/otautil/logging.cpp
@@ -178,9 +178,8 @@
   tmplog_offset = 0;
 }
 
-static void copy_log_file(const std::string& source, const std::string& destination, bool append,
-                          const selabel_handle* sehandle) {
-  FILE* dest_fp = fopen_path(destination, append ? "ae" : "we", sehandle);
+static void copy_log_file(const std::string& source, const std::string& destination, bool append) {
+  FILE* dest_fp = fopen_path(destination, append ? "ae" : "we", logging_sehandle);
   if (dest_fp == nullptr) {
     PLOG(ERROR) << "Can't open " << destination;
   } else {
@@ -203,7 +202,7 @@
   }
 }
 
-void copy_logs(bool save_current_log, bool has_cache, const selabel_handle* sehandle) {
+void copy_logs(bool save_current_log) {
   // We only rotate and record the log of the current session if explicitly requested. This usually
   // happens after wipes, installation from BCB or menu selections. This is to avoid unnecessary
   // rotation (and possible deletion) of log files, if it does not do anything loggable.
@@ -216,7 +215,7 @@
   copy_log_file_to_pmsg(Paths::Get().temporary_install_file(), LAST_INSTALL_FILE);
 
   // We can do nothing for now if there's no /cache partition.
-  if (!has_cache) {
+  if (!HasCache()) {
     return;
   }
 
@@ -225,9 +224,9 @@
   rotate_logs(LAST_LOG_FILE, LAST_KMSG_FILE);
 
   // Copy logs to cache so the system can find out what happened.
-  copy_log_file(Paths::Get().temporary_log_file(), LOG_FILE, true, sehandle);
-  copy_log_file(Paths::Get().temporary_log_file(), LAST_LOG_FILE, false, sehandle);
-  copy_log_file(Paths::Get().temporary_install_file(), LAST_INSTALL_FILE, false, sehandle);
+  copy_log_file(Paths::Get().temporary_log_file(), LOG_FILE, true);
+  copy_log_file(Paths::Get().temporary_log_file(), LAST_LOG_FILE, false);
+  copy_log_file(Paths::Get().temporary_install_file(), LAST_INSTALL_FILE, false);
   save_kernel_log(LAST_KMSG_FILE);
   chmod(LOG_FILE, 0600);
   chown(LOG_FILE, AID_SYSTEM, AID_SYSTEM);
@@ -319,7 +318,7 @@
   // Reset the pointer so we copy from the beginning of the temp
   // log.
   reset_tmplog_offset();
-  copy_logs(true /* save_current_log */, true /* has_cache */, logging_sehandle);
+  copy_logs(true /* save_current_log */);
 
   return true;
 }
diff --git a/otautil/roots.cpp b/otautil/roots.cpp
index a778e05..4315517 100644
--- a/otautil/roots.cpp
+++ b/otautil/roots.cpp
@@ -51,6 +51,8 @@
 
 static Fstab fstab;
 
+constexpr const char* CACHE_ROOT = "/cache";
+
 void load_volume_table() {
   if (!ReadDefaultFstab(&fstab)) {
     LOG(ERROR) << "Failed to read default fstab";
@@ -275,3 +277,9 @@
   }
   return 0;
 }
+
+bool HasCache() {
+  CHECK(!fstab.empty());
+  static bool has_cache = volume_for_mount_point(CACHE_ROOT) != nullptr;
+  return has_cache;
+}
diff --git a/recovery.cpp b/recovery.cpp
index 97ca0a5..b989b24 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -69,10 +69,6 @@
 
 static constexpr const char* CACHE_ROOT = "/cache";
 
-// We define RECOVERY_API_VERSION in Android.mk, which will be picked up by build system and packed
-// 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 save_current_log = false;
 std::string stage;
 const char* reason = nullptr;
@@ -106,7 +102,7 @@
  *    -- after this, rebooting will restart the erase --
  * 5. erase_volume() reformats /data
  * 6. erase_volume() reformats /cache
- * 7. finish_recovery() erases BCB
+ * 7. FinishRecovery() erases BCB
  *    -- after this, rebooting will restart the main system --
  * 8. main() calls reboot() to boot main system
  *
@@ -118,25 +114,25 @@
  *    -- after this, rebooting will attempt to reinstall the update --
  * 5. InstallPackage() attempts to install the update
  *    NOTE: the package install must itself be restartable from any point
- * 6. finish_recovery() erases BCB
+ * 6. FinishRecovery() erases BCB
  *    -- after this, rebooting will (try to) restart the main system --
  * 7. ** if install failed **
  *    7a. PromptAndWait() shows an error icon and waits for the user
  *    7b. the user reboots (pulling the battery, etc) into the main system
  */
 
-bool is_ro_debuggable() {
-    return android::base::GetBoolProperty("ro.debuggable", false);
+static bool IsRoDebuggable() {
+  return android::base::GetBoolProperty("ro.debuggable", false);
 }
 
 // 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 FinishRecovery(RecoveryUI* ui) {
   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) {
+  if (!locale.empty() && HasCache()) {
     LOG(INFO) << "Saving locale \"" << locale << "\"";
     if (ensure_path_mounted(LOCALE_FILE) != 0) {
       LOG(ERROR) << "Failed to mount " << LOCALE_FILE;
@@ -145,7 +141,7 @@
     }
   }
 
-  copy_logs(save_current_log, has_cache, sehandle);
+  copy_logs(save_current_log);
 
   // Reset to normal system boot so recovery won't cycle indefinitely.
   std::string err;
@@ -154,7 +150,7 @@
   }
 
   // Remove the command file, so recovery won't repeat indefinitely.
-  if (has_cache) {
+  if (HasCache()) {
     if (ensure_path_mounted(COMMAND_FILE) != 0 || (unlink(COMMAND_FILE) && errno != ENOENT)) {
       LOG(WARNING) << "Can't unlink " << COMMAND_FILE;
     }
@@ -168,7 +164,7 @@
   std::vector<std::string> headers{ question1, question2 };
   std::vector<std::string> items{ " No", " Yes" };
 
-  size_t chosen_item = ui->ShowMenu(
+  size_t chosen_item = device->GetUI()->ShowMenu(
       headers, items, 0, true,
       std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
   return (chosen_item == 1);
@@ -178,7 +174,7 @@
   std::vector<std::string> headers{ "Wipe all user data?", "  THIS CAN NOT BE UNDONE!" };
   std::vector<std::string> items{ " Cancel", " Factory data reset" };
 
-  size_t chosen_item = ui->ShowPromptWipeDataConfirmationMenu(
+  size_t chosen_item = device->GetUI()->ShowPromptWipeDataConfirmationMenu(
       headers, items,
       std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
 
@@ -200,7 +196,7 @@
   };
   // clang-format on
   for (;;) {
-    size_t chosen_item = ui->ShowPromptWipeDataMenu(
+    size_t chosen_item = device->GetUI()->ShowPromptWipeDataMenu(
         wipe_data_menu_headers, wipe_data_menu_items,
         std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
     // If ShowMenu() returned RecoveryUI::KeyError::INTERRUPTED, WaitKey() was interrupted.
@@ -224,7 +220,7 @@
 
 static void choose_recovery_file(Device* device) {
   std::vector<std::string> entries;
-  if (has_cache) {
+  if (HasCache()) {
     for (int i = 0; i < KEEP_LOG_COUNT; i++) {
       auto add_to_entries = [&](const char* filename) {
         std::string log_file(filename);
@@ -258,7 +254,7 @@
 
   size_t chosen_item = 0;
   while (true) {
-    chosen_item = ui->ShowMenu(
+    chosen_item = device->GetUI()->ShowMenu(
         headers, entries, chosen_item, true,
         std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
 
@@ -268,11 +264,11 @@
     }
     if (entries[chosen_item] == "Back") break;
 
-    ui->ShowFile(entries[chosen_item]);
+    device->GetUI()->ShowFile(entries[chosen_item]);
   }
 }
 
-static void run_graphics_test() {
+static void run_graphics_test(RecoveryUI* ui) {
   // Switch to graphics screen.
   ui->ShowText(false);
 
@@ -319,8 +315,9 @@
 // as REBOOT, SHUTDOWN, or REBOOT_BOOTLOADER. Returning NO_ACTION means to take the default, which
 // is to reboot or shutdown depending on if the --shutdown_after flag was passed to recovery.
 static Device::BuiltinAction PromptAndWait(Device* device, InstallResult status) {
+  auto ui = device->GetUI();
   for (;;) {
-    finish_recovery();
+    FinishRecovery(ui);
     switch (status) {
       case INSTALL_SUCCESS:
       case INSTALL_NONE:
@@ -421,7 +418,7 @@
         if (status != INSTALL_SUCCESS) {
           ui->SetBackground(RecoveryUI::ERROR);
           ui->Print("Installation aborted.\n");
-          copy_logs(save_current_log, has_cache, sehandle);
+          copy_logs(save_current_log);
         } else if (!ui->IsTextVisible()) {
           return Device::NO_ACTION;  // reboot if logs aren't visible
         }
@@ -433,7 +430,7 @@
         break;
 
       case Device::RUN_GRAPHICS_TEST:
-        run_graphics_test();
+        run_graphics_test(ui);
         break;
 
       case Device::RUN_LOCALE_TEST: {
@@ -680,6 +677,8 @@
   printf("stage is [%s]\n", stage.c_str());
   printf("reason is [%s]\n", reason);
 
+  auto ui = device->GetUI();
+
   // Set background string to "installing security update" for security update,
   // otherwise set it to "installing system update".
   ui->SetSystemUpdateText(security_update);
@@ -706,8 +705,6 @@
   property_list(print_property, nullptr);
   printf("\n");
 
-  ui->Print("Supported API: %d\n", kRecoveryApiVersion);
-
   InstallResult status = INSTALL_SUCCESS;
   // next_action indicates the next target to reboot into upon finishing the install. It could be
   // overridden to a different reboot target per user request.
@@ -768,7 +765,7 @@
         // RETRY_LIMIT times before we abandon this OTA update.
         static constexpr int RETRY_LIMIT = 4;
         if (status == INSTALL_RETRY && retry_count < RETRY_LIMIT) {
-          copy_logs(save_current_log, has_cache, sehandle);
+          copy_logs(save_current_log);
           retry_count += 1;
           set_retry_bootloader_message(retry_count, args);
           // Print retry count on screen.
@@ -786,7 +783,7 @@
         // If this is an eng or userdebug build, then automatically
         // turn the text display on if the script fails so the error
         // message is visible.
-        if (is_ro_debuggable()) {
+        if (IsRoDebuggable()) {
           ui->ShowText(true);
         }
       }
@@ -843,7 +840,7 @@
     // If this is an eng or userdebug build, automatically turn on the text display if no command
     // is specified. Note that this should be called before setting the background to avoid
     // flickering the background image.
-    if (is_ro_debuggable()) {
+    if (IsRoDebuggable()) {
       ui->ShowText(true);
     }
     status = INSTALL_NONE;  // No command specified
@@ -876,7 +873,7 @@
   }
 
   // Save logs and clean up before rebooting or shutting down.
-  finish_recovery();
+  FinishRecovery(ui);
 
   return next_action;
 }
diff --git a/recovery_main.cpp b/recovery_main.cpp
index 7fbdf9a..f0d75ee 100644
--- a/recovery_main.cpp
+++ b/recovery_main.cpp
@@ -63,12 +63,11 @@
 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";
+static RecoveryUI* ui = nullptr;
 
-bool has_cache = false;
-
-RecoveryUI* ui = nullptr;
-struct selabel_handle* sehandle;
+static bool IsRoDebuggable() {
+  return android::base::GetBoolProperty("ro.debuggable", false);
+}
 
 static void UiLogger(android::base::LogId /* id */, android::base::LogSeverity severity,
                      const char* /* tag */, const char* /* file */, unsigned int /* line */,
@@ -131,7 +130,7 @@
   }
 
   // --- if that doesn't work, try the command file (if we have /cache).
-  if (args.size() == 1 && has_cache) {
+  if (args.size() == 1 && HasCache()) {
     std::string content;
     if (ensure_path_mounted(COMMAND_FILE) == 0 &&
         android::base::ReadFileToString(COMMAND_FILE, &content)) {
@@ -148,7 +147,7 @@
 
   // 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.
+  // finish the pending work, until FinishRecovery() 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;
@@ -331,7 +330,6 @@
   redirect_stdio(Paths::Get().temporary_log_file().c_str());
 
   load_volume_table();
-  has_cache = volume_for_mount_point(CACHE_ROOT) != nullptr;
 
   std::vector<std::string> args = get_args(argc, argv);
   auto args_to_parse = StringVectorToNullTerminatedArray(args);
@@ -370,7 +368,7 @@
   optind = 1;
 
   if (locale.empty()) {
-    if (has_cache) {
+    if (HasCache()) {
       locale = load_locale_from_cache();
     }
 
@@ -416,7 +414,7 @@
   }
   ui = device->GetUI();
 
-  if (!has_cache) {
+  if (!HasCache()) {
     device->RemoveMenuItemForAction(Device::WIPE_CACHE);
   }
 
@@ -424,7 +422,7 @@
     device->RemoveMenuItemForAction(Device::ENTER_FASTBOOT);
   }
 
-  if (!is_ro_debuggable()) {
+  if (!IsRoDebuggable()) {
     device->RemoveMenuItemForAction(Device::ENTER_RESCUE);
   }
 
@@ -434,7 +432,7 @@
   LOG(INFO) << "Starting recovery (pid " << getpid() << ") on " << ctime(&start);
   LOG(INFO) << "locale is [" << locale << "]";
 
-  sehandle = selinux_android_file_context_handle();
+  auto sehandle = selinux_android_file_context_handle();
   selinux_android_set_sehandle(sehandle);
   if (!sehandle) {
     ui->Print("Warning: No file_contexts\n");
@@ -447,7 +445,7 @@
   listener_thread.detach();
 
   while (true) {
-    std::string usb_config = fastboot ? "fastboot" : is_ro_debuggable() ? "adb" : "none";
+    std::string usb_config = fastboot ? "fastboot" : IsRoDebuggable() ? "adb" : "none";
     std::string usb_state = android::base::GetProperty("sys.usb.state", "none");
     if (usb_config != usb_state) {
       if (!SetUsbConfig("none")) {
diff --git a/recovery_ui/include/recovery_ui/screen_ui.h b/recovery_ui/include/recovery_ui/screen_ui.h
index 5cda2a2..92b3c25 100644
--- a/recovery_ui/include/recovery_ui/screen_ui.h
+++ b/recovery_ui/include/recovery_ui/screen_ui.h
@@ -286,6 +286,9 @@
   // selected.
   virtual int SelectMenu(int sel);
 
+  // Returns the help message displayed on top of the menu.
+  virtual std::vector<std::string> GetMenuHelpMessage() const;
+
   virtual void draw_background_locked();
   virtual void draw_foreground_locked();
   virtual void draw_screen_locked();
diff --git a/recovery_ui/include/recovery_ui/ui.h b/recovery_ui/include/recovery_ui/ui.h
index a95f935..08ec1d7 100644
--- a/recovery_ui/include/recovery_ui/ui.h
+++ b/recovery_ui/include/recovery_ui/ui.h
@@ -118,7 +118,7 @@
 
   // Returns true if you have the volume up/down and power trio typical of phones and tablets, false
   // otherwise.
-  virtual bool HasThreeButtons();
+  virtual bool HasThreeButtons() const;
 
   // Returns true if it has a power key.
   virtual bool HasPowerKey() const;
diff --git a/recovery_ui/screen_ui.cpp b/recovery_ui/screen_ui.cpp
index f2af66c..087fc0e 100644
--- a/recovery_ui/screen_ui.cpp
+++ b/recovery_ui/screen_ui.cpp
@@ -673,6 +673,19 @@
   title_lines_ = lines;
 }
 
+std::vector<std::string> ScreenRecoveryUI::GetMenuHelpMessage() const {
+  // clang-format off
+  static std::vector<std::string> REGULAR_HELP{
+    "Use volume up/down and power.",
+  };
+  static std::vector<std::string> LONG_PRESS_HELP{
+    "Any button cycles highlight.",
+    "Long-press activates.",
+  };
+  // clang-format on
+  return HasThreeButtons() ? REGULAR_HELP : LONG_PRESS_HELP;
+}
+
 // Redraws everything on the screen. Does not flip pages. Should only be called with updateMutex
 // locked.
 void ScreenRecoveryUI::draw_screen_locked() {
@@ -685,16 +698,7 @@
   gr_color(0, 0, 0, 255);
   gr_clear();
 
-  // clang-format off
-  static std::vector<std::string> REGULAR_HELP{
-    "Use volume up/down and power.",
-  };
-  static std::vector<std::string> LONG_PRESS_HELP{
-    "Any button cycles highlight.",
-    "Long-press activates.",
-  };
-  // clang-format on
-  draw_menu_and_text_buffer_locked(HasThreeButtons() ? REGULAR_HELP : LONG_PRESS_HELP);
+  draw_menu_and_text_buffer_locked(GetMenuHelpMessage());
 }
 
 // Draws the menu and text buffer on the screen. Should only be called with updateMutex locked.
diff --git a/recovery_ui/ui.cpp b/recovery_ui/ui.cpp
index 98c654d..6f5cbbc 100644
--- a/recovery_ui/ui.cpp
+++ b/recovery_ui/ui.cpp
@@ -528,7 +528,7 @@
   return result;
 }
 
-bool RecoveryUI::HasThreeButtons() {
+bool RecoveryUI::HasThreeButtons() const {
   return has_power_key && has_up_key && has_down_key;
 }
 
diff --git a/tests/unit/host/update_simulator_test.cpp b/tests/unit/host/update_simulator_test.cpp
index bf89b78..fb12178 100644
--- a/tests/unit/host/update_simulator_test.cpp
+++ b/tests/unit/host/update_simulator_test.cpp
@@ -93,7 +93,7 @@
 
   // Run the update simulation and check the result.
   TemporaryDir work_dir;
-  BuildInfo build_info(work_dir.path);
+  BuildInfo build_info(work_dir.path, false);
   ASSERT_TRUE(build_info.ParseTargetFile(src_tf, false));
   Updater updater(std::make_unique<SimulatorRuntime>(&build_info));
   ASSERT_TRUE(updater.Init(cmd_pipe.release(), ota_package, false));
@@ -211,7 +211,7 @@
   AddZipEntries(zip_file.release(), entries);
 
   TemporaryDir temp_dir;
-  BuildInfo build_info(temp_dir.path);
+  BuildInfo build_info(temp_dir.path, false);
   ASSERT_TRUE(build_info.ParseTargetFile(zip_file.path, false));
 
   std::map<string, string> expected_result = {
diff --git a/updater/build_info.cpp b/updater/build_info.cpp
index 3072aab..f168008 100644
--- a/updater/build_info.cpp
+++ b/updater/build_info.cpp
@@ -16,6 +16,8 @@
 
 #include "updater/build_info.h"
 
+#include <stdio.h>
+
 #include <set>
 #include <vector>
 
@@ -55,12 +57,23 @@
         return false;
       }
 
+      std::string mapped_path = image_file.path;
+      // Rename the images to more readable ones if we want to keep the image.
+      if (keep_images_) {
+        mapped_path = work_dir_ + fstab_info.mount_point + ".img";
+        image_file.release();
+        if (rename(image_file.path, mapped_path.c_str()) != 0) {
+          PLOG(ERROR) << "Failed to rename " << image_file.path << " to " << mapped_path;
+          return false;
+        }
+      }
+
       LOG(INFO) << "Mounted " << fstab_info.mount_point << "\nMapping: " << fstab_info.blockdev_name
-                << " to " << image_file.path;
+                << " to " << mapped_path;
 
       blockdev_map_.emplace(
           fstab_info.blockdev_name,
-          FakeBlockDevice(fstab_info.blockdev_name, fstab_info.mount_point, image_file.path));
+          FakeBlockDevice(fstab_info.blockdev_name, fstab_info.mount_point, mapped_path));
       break;
     }
   }
diff --git a/updater/include/updater/build_info.h b/updater/include/updater/build_info.h
index 2229957..0073bfa 100644
--- a/updater/include/updater/build_info.h
+++ b/updater/include/updater/build_info.h
@@ -43,7 +43,8 @@
 // query the information and run the update on host.
 class BuildInfo {
  public:
-  explicit BuildInfo(const std::string_view work_dir) : work_dir_(work_dir) {}
+  BuildInfo(const std::string_view work_dir, bool keep_images)
+      : work_dir_(work_dir), keep_images_(keep_images) {}
   // Returns the value of the build properties.
   std::string GetProperty(const std::string_view key, const std::string_view default_value) const;
   // Returns the path to the mock block device.
@@ -69,4 +70,5 @@
 
   std::list<TemporaryFile> temp_files_;
   std::string work_dir_;  // A temporary directory to store the extracted image files
+  bool keep_images_;
 };
diff --git a/updater/install.cpp b/updater/install.cpp
index c82351e..be0ceb0 100644
--- a/updater/install.cpp
+++ b/updater/install.cpp
@@ -113,7 +113,7 @@
                         argv.size());
     }
     const std::string& zip_path = args[0];
-    const std::string& dest_path = args[1];
+    std::string dest_path = args[1];
 
     ZipArchiveHandle za = state->updater->GetPackageHandle();
     ZipEntry entry;
@@ -122,6 +122,13 @@
       return StringValue("");
     }
 
+    // Update the destination of package_extract_file if it's a block device. During simulation the
+    // destination will map to a fake file.
+    if (std::string block_device_name = state->updater->FindBlockDeviceName(dest_path);
+        !block_device_name.empty()) {
+      dest_path = block_device_name;
+    }
+
     android::base::unique_fd fd(TEMP_FAILURE_RETRY(
         open(dest_path.c_str(), O_WRONLY | O_CREAT | O_TRUNC, S_IRUSR | S_IWUSR)));
     if (fd == -1) {
diff --git a/updater/update_simulator_main.cpp b/updater/update_simulator_main.cpp
index 278a440..6c6989b 100644
--- a/updater/update_simulator_main.cpp
+++ b/updater/update_simulator_main.cpp
@@ -58,12 +58,16 @@
   std::string skip_function_file;
   std::string source_target_file;
   std::string package_name;
+  std::string work_dir;
+  bool keep_images = false;
 
   constexpr struct option OPTIONS[] = {
+    { "keep_images", no_argument, nullptr, 0 },
     { "oem_settings", required_argument, nullptr, 0 },
     { "ota_package", required_argument, nullptr, 0 },
     { "skip_functions", required_argument, nullptr, 0 },
     { "source", required_argument, nullptr, 0 },
+    { "work_dir", required_argument, nullptr, 0 },
     { nullptr, 0, nullptr, 0 },
   };
 
@@ -86,6 +90,10 @@
       source_target_file = optarg;
     } else if (option_name == "ota_package"s) {
       package_name = optarg;
+    } else if (option_name == "keep_images"s) {
+      keep_images = true;
+    } else if (option_name == "work_dir"s) {
+      work_dir = optarg;
     } else {
       Usage(argv[0]);
       return EXIT_FAILURE;
@@ -129,8 +137,11 @@
 
   TemporaryFile cmd_pipe;
   TemporaryDir source_temp_dir;
+  if (work_dir.empty()) {
+    work_dir = source_temp_dir.path;
+  }
 
-  BuildInfo source_build_info(source_temp_dir.path);
+  BuildInfo source_build_info(work_dir, keep_images);
   if (!source_build_info.ParseTargetFile(source_target_file, false)) {
     LOG(ERROR) << "Failed to parse the target file " << source_target_file;
     return EXIT_FAILURE;