5 files changed, 40 insertions, 7 deletions

diff --git a/tests/unit/host/update_simulator_test.cpp b/tests/unit/host/update_simulator_test.cpp
index bf89b7859..fb1217877 100644
--- a/tests/unit/host/update_simulator_test.cpp
+++ b/tests/unit/host/update_simulator_test.cpp
@@ -93,7 +93,7 @@ static void RunSimulation(std::string_view src_tf, std::string_view ota_package,
 
   // 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 @@ TEST_F(UpdateSimulatorTest, BuildInfo_ParseTargetFile) {
   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 3072aab54..f168008ec 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 @@ bool BuildInfo::ParseTargetFile(const std::string_view target_file_path, bool ex
         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 22299579f..0073bfa4a 100644
--- a/updater/include/updater/build_info.h
+++ b/updater/include/updater/build_info.h
@@ -43,7 +43,8 @@ class FakeBlockDevice {
 // 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 @@ class BuildInfo {
 
   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 c82351ec4..be0ceb06c 100644
--- a/updater/install.cpp
+++ b/updater/install.cpp
@@ -113,7 +113,7 @@ Value* PackageExtractFileFn(const char* name, State* state,
                         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 @@ Value* PackageExtractFileFn(const char* name, State* state,
       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 278a44057..6c6989bac 100644
--- a/updater/update_simulator_main.cpp
+++ b/updater/update_simulator_main.cpp
@@ -58,12 +58,16 @@ int main(int argc, char** argv) {
   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 @@ int main(int argc, char** argv) {
       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 @@ int main(int argc, char** argv) {
 
   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;