3 files changed, 103 insertions, 33 deletions

diff --git a/install.cpp b/install.cpp
index 689f4a0c6..a1f2e4fbd 100644
--- a/install.cpp
+++ b/install.cpp
@@ -51,6 +51,7 @@
 #include "error_code.h"
 #include "otautil/SysUtil.h"
 #include "otautil/ThermalUtil.h"
+#include "private/install.h"
 #include "roots.h"
 #include "ui.h"
 #include "verifier.h"
@@ -125,12 +126,6 @@ static void read_source_target_build(ZipArchiveHandle zip, std::vector<std::stri
   }
 }
 
-// Extract the update binary from the open zip archive |zip| located at |path| and store into |cmd|
-// the command line that should be called. The |status_fd| is the file descriptor the child process
-// should use to report back the progress of the update.
-int update_binary_command(const std::string& path, ZipArchiveHandle zip, int retry_count,
-                          int status_fd, std::vector<std::string>* cmd);
-
 #ifdef AB_OTA_UPDATER
 
 // Parses the metadata of the OTA package in |zip| and checks whether we are
@@ -211,8 +206,9 @@ static int check_newer_ab_build(ZipArchiveHandle zip) {
   return 0;
 }
 
-int update_binary_command(const std::string& path, ZipArchiveHandle zip, int /* retry_count */,
-                          int status_fd, std::vector<std::string>* cmd) {
+int update_binary_command(const std::string& package, ZipArchiveHandle zip,
+                          const std::string& binary_path, int /* retry_count */, int status_fd,
+                          std::vector<std::string>* cmd) {
   CHECK(cmd != nullptr);
   int ret = check_newer_ab_build(zip);
   if (ret != 0) {
@@ -246,8 +242,8 @@ int update_binary_command(const std::string& path, ZipArchiveHandle zip, int /*
   }
   long payload_offset = payload_entry.offset;
   *cmd = {
-    "/sbin/update_engine_sideload",
-    "--payload=file://" + path,
+    binary_path,
+    "--payload=file://" + package,
     android::base::StringPrintf("--offset=%ld", payload_offset),
     "--headers=" + std::string(payload_properties.begin(), payload_properties.end()),
     android::base::StringPrintf("--status_fd=%d", status_fd),
@@ -257,8 +253,9 @@ int update_binary_command(const std::string& path, ZipArchiveHandle zip, int /*
 
 #else  // !AB_OTA_UPDATER
 
-int update_binary_command(const std::string& path, ZipArchiveHandle zip, int retry_count,
-                          int status_fd, std::vector<std::string>* cmd) {
+int update_binary_command(const std::string& package, ZipArchiveHandle zip,
+                          const std::string& binary_path, int retry_count, int status_fd,
+                          std::vector<std::string>* cmd) {
   CHECK(cmd != nullptr);
 
   // On traditional updates we extract the update binary from the package.
@@ -270,11 +267,10 @@ int update_binary_command(const std::string& path, ZipArchiveHandle zip, int ret
     return INSTALL_CORRUPT;
   }
 
-  const char* binary = "/tmp/update_binary";
-  unlink(binary);
-  int fd = creat(binary, 0755);
+  unlink(binary_path.c_str());
+  int fd = creat(binary_path.c_str(), 0755);
   if (fd == -1) {
-    PLOG(ERROR) << "Failed to create " << binary;
+    PLOG(ERROR) << "Failed to create " << binary_path;
     return INSTALL_ERROR;
   }
 
@@ -286,10 +282,10 @@ int update_binary_command(const std::string& path, ZipArchiveHandle zip, int ret
   }
 
   *cmd = {
-    binary,
+    binary_path,
     EXPAND(RECOVERY_API_VERSION),  // defined in Android.mk
     std::to_string(status_fd),
-    path,
+    package,
   };
   if (retry_count > 0) {
     cmd->push_back("retry");
@@ -308,7 +304,7 @@ static void log_max_temperature(int* max_temperature) {
 }
 
 // If the package contains an update binary, extract it and run it.
-static int try_update_binary(const std::string& path, ZipArchiveHandle zip, bool* wipe_cache,
+static int try_update_binary(const std::string& package, ZipArchiveHandle zip, bool* wipe_cache,
                              std::vector<std::string>* log_buffer, int retry_count,
                              int* max_temperature) {
   read_source_target_build(zip, log_buffer);
@@ -317,7 +313,13 @@ static int try_update_binary(const std::string& path, ZipArchiveHandle zip, bool
   pipe(pipefd);
 
   std::vector<std::string> args;
-  int ret = update_binary_command(path, zip, retry_count, pipefd[1], &args);
+#ifdef AB_OTA_UPDATER
+  int ret = update_binary_command(package, zip, "/sbin/update_engine_sideload", retry_count,
+                                  pipefd[1], &args);
+#else
+  int ret = update_binary_command(package, zip, "/tmp/update-binary", retry_count, pipefd[1],
+                                  &args);
+#endif
   if (ret) {
     close(pipefd[0]);
     close(pipefd[1]);
@@ -472,7 +474,7 @@ static int try_update_binary(const std::string& path, ZipArchiveHandle zip, bool
     return INSTALL_RETRY;
   }
   if (!WIFEXITED(status) || WEXITSTATUS(status) != 0) {
-    LOG(ERROR) << "Error in " << path << " (Status " << WEXITSTATUS(status) << ")";
+    LOG(ERROR) << "Error in " << package << " (Status " << WEXITSTATUS(status) << ")";
     return INSTALL_ERROR;
   }
 
diff --git a/private/install.h b/private/install.h
index 12d303b01..ef64bd41d 100644
--- a/private/install.h
+++ b/private/install.h
@@ -23,5 +23,9 @@
 
 #include <ziparchive/zip_archive.h>
 
-int update_binary_command(const std::string& path, ZipArchiveHandle zip, int retry_count,
-                          int status_fd, std::vector<std::string>* cmd);
+// Extract the update binary from the open zip archive |zip| located at |package| to |binary_path|.
+// Store the command line that should be called into |cmd|. The |status_fd| is the file descriptor
+// the child process should use to report back the progress of the update.
+int update_binary_command(const std::string& package, ZipArchiveHandle zip,
+                          const std::string& binary_path, int retry_count, int status_fd,
+                          std::vector<std::string>* cmd);
diff --git a/tests/component/install_test.cpp b/tests/component/install_test.cpp
index 40201d76f..968196fc0 100644
--- a/tests/component/install_test.cpp
+++ b/tests/component/install_test.cpp
@@ -15,6 +15,8 @@
  */
 
 #include <stdio.h>
+#include <sys/stat.h>
+#include <sys/types.h>
 #include <unistd.h>
 
 #include <string>
@@ -225,18 +227,62 @@ TEST(InstallTest, update_binary_command_smoke) {
 
   ZipArchiveHandle zip;
   ASSERT_EQ(0, OpenArchive(temp_file.path, &zip));
+  ZipString payload_name("payload.bin");
+  ZipEntry payload_entry;
+  ASSERT_EQ(0, FindEntry(zip, payload_name, &payload_entry));
   int status_fd = 10;
-  std::string path = "/path/to/update.zip";
+  std::string package = "/path/to/update.zip";
+  std::string binary_path = "/sbin/update_engine_sideload";
   std::vector<std::string> cmd;
-  ASSERT_EQ(0, update_binary_command(path, zip, 0, status_fd, &cmd));
-  ASSERT_EQ("/sbin/update_engine_sideload", cmd[0]);
-  ASSERT_EQ("--payload=file://" + path, cmd[1]);
+  ASSERT_EQ(0, update_binary_command(package, zip, binary_path, 0, status_fd, &cmd));
+  ASSERT_EQ(5U, cmd.size());
+  ASSERT_EQ(binary_path, cmd[0]);
+  ASSERT_EQ("--payload=file://" + package, cmd[1]);
+  ASSERT_EQ("--offset=" + std::to_string(payload_entry.offset), cmd[2]);
   ASSERT_EQ("--headers=" + properties, cmd[3]);
   ASSERT_EQ("--status_fd=" + std::to_string(status_fd), cmd[4]);
   CloseArchive(zip);
 #else
-  // Cannot test update_binary_command() because it tries to extract update-binary to /tmp.
-  GTEST_LOG_(INFO) << "Test skipped on non-A/B device.";
+  TemporaryFile temp_file;
+  FILE* zip_file = fdopen(temp_file.fd, "w");
+  ZipWriter writer(zip_file);
+  static constexpr const char* UPDATE_BINARY_NAME = "META-INF/com/google/android/update-binary";
+  ASSERT_EQ(0, writer.StartEntry(UPDATE_BINARY_NAME, kCompressStored));
+  ASSERT_EQ(0, writer.FinishEntry());
+  ASSERT_EQ(0, writer.Finish());
+  ASSERT_EQ(0, fclose(zip_file));
+
+  ZipArchiveHandle zip;
+  ASSERT_EQ(0, OpenArchive(temp_file.path, &zip));
+  int status_fd = 10;
+  std::string package = "/path/to/update.zip";
+  TemporaryDir td;
+  std::string binary_path = std::string(td.path) + "/update_binary";
+  std::vector<std::string> cmd;
+  ASSERT_EQ(0, update_binary_command(package, zip, binary_path, 0, status_fd, &cmd));
+  ASSERT_EQ(4U, cmd.size());
+  ASSERT_EQ(binary_path, cmd[0]);
+  ASSERT_EQ("3", cmd[1]);  // RECOVERY_API_VERSION
+  ASSERT_EQ(std::to_string(status_fd), cmd[2]);
+  ASSERT_EQ(package, cmd[3]);
+  struct stat sb;
+  ASSERT_EQ(0, stat(binary_path.c_str(), &sb));
+  ASSERT_EQ(static_cast<mode_t>(0755), sb.st_mode & (S_IRWXU | S_IRWXG | S_IRWXO));
+
+  // With non-zero retry count. update_binary will be removed automatically.
+  cmd.clear();
+  ASSERT_EQ(0, update_binary_command(package, zip, binary_path, 2, status_fd, &cmd));
+  ASSERT_EQ(5U, cmd.size());
+  ASSERT_EQ(binary_path, cmd[0]);
+  ASSERT_EQ("3", cmd[1]);  // RECOVERY_API_VERSION
+  ASSERT_EQ(std::to_string(status_fd), cmd[2]);
+  ASSERT_EQ(package, cmd[3]);
+  ASSERT_EQ("retry", cmd[4]);
+  sb = {};
+  ASSERT_EQ(0, stat(binary_path.c_str(), &sb));
+  ASSERT_EQ(static_cast<mode_t>(0755), sb.st_mode & (S_IRWXU | S_IRWXG | S_IRWXO));
+
+  CloseArchive(zip);
 #endif  // AB_OTA_UPDATER
 }
 
@@ -267,12 +313,30 @@ TEST(InstallTest, update_binary_command_invalid) {
   ZipArchiveHandle zip;
   ASSERT_EQ(0, OpenArchive(temp_file.path, &zip));
   int status_fd = 10;
-  std::string path = "/path/to/update.zip";
+  std::string package = "/path/to/update.zip";
+  std::string binary_path = "/sbin/update_engine_sideload";
   std::vector<std::string> cmd;
-  ASSERT_EQ(INSTALL_CORRUPT, update_binary_command(path, zip, 0, status_fd, &cmd));
+  ASSERT_EQ(INSTALL_CORRUPT, update_binary_command(package, zip, binary_path, 0, status_fd, &cmd));
   CloseArchive(zip);
 #else
-  // Cannot test update_binary_command() because it tries to extract update-binary to /tmp.
-  GTEST_LOG_(INFO) << "Test skipped on non-A/B device.";
+  TemporaryFile temp_file;
+  FILE* zip_file = fdopen(temp_file.fd, "w");
+  ZipWriter writer(zip_file);
+  // The archive must have something to be opened correctly.
+  ASSERT_EQ(0, writer.StartEntry("dummy_entry", 0));
+  ASSERT_EQ(0, writer.FinishEntry());
+  ASSERT_EQ(0, writer.Finish());
+  ASSERT_EQ(0, fclose(zip_file));
+
+  // Missing update binary.
+  ZipArchiveHandle zip;
+  ASSERT_EQ(0, OpenArchive(temp_file.path, &zip));
+  int status_fd = 10;
+  std::string package = "/path/to/update.zip";
+  TemporaryDir td;
+  std::string binary_path = std::string(td.path) + "/update_binary";
+  std::vector<std::string> cmd;
+  ASSERT_EQ(INSTALL_CORRUPT, update_binary_command(package, zip, binary_path, 0, status_fd, &cmd));
+  CloseArchive(zip);
 #endif  // AB_OTA_UPDATER
 }