4 files changed, 199 insertions, 82 deletions

diff --git a/install.cpp b/install.cpp
index 6dcd3565e..4fb809996 100644
--- a/install.cpp
+++ b/install.cpp
@@ -57,9 +57,6 @@
 
 using namespace std::chrono_literals;
 
-#define ASSUMED_UPDATE_BINARY_NAME  "META-INF/com/google/android/update-binary"
-static constexpr const char* AB_OTA_PAYLOAD_PROPERTIES = "payload_properties.txt";
-static constexpr const char* AB_OTA_PAYLOAD = "payload.bin";
 #define PUBLIC_KEYS_FILE "/res/keys"
 static constexpr const char* METADATA_PATH = "META-INF/com/android/metadata";
 static constexpr const char* UNCRYPT_STATUS = "/cache/recovery/uncrypt_status";
@@ -136,13 +133,11 @@ 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.
-static int
-update_binary_command(const char* 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 |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
 
@@ -224,87 +219,90 @@ static int check_newer_ab_build(ZipArchiveHandle zip) {
   return 0;
 }
 
-static int
-update_binary_command(const char* path, ZipArchiveHandle zip, int retry_count,
-                      int status_fd, std::vector<std::string>* cmd)
-{
-    int ret = check_newer_ab_build(zip);
-    if (ret) {
-        return ret;
-    }
+int update_binary_command(const std::string& path, ZipArchiveHandle zip, int retry_count,
+                          int status_fd, std::vector<std::string>* cmd) {
+  CHECK(cmd != nullptr);
+  int ret = check_newer_ab_build(zip);
+  if (ret != 0) {
+    return ret;
+  }
 
-    // For A/B updates we extract the payload properties to a buffer and obtain
-    // the RAW payload offset in the zip file.
-    ZipString property_name(AB_OTA_PAYLOAD_PROPERTIES);
-    ZipEntry properties_entry;
-    if (FindEntry(zip, property_name, &properties_entry) != 0) {
-        LOG(ERROR) << "Can't find " << AB_OTA_PAYLOAD_PROPERTIES;
-        return INSTALL_CORRUPT;
-    }
-    std::vector<uint8_t> payload_properties(
-            properties_entry.uncompressed_length);
-    if (ExtractToMemory(zip, &properties_entry, payload_properties.data(),
-                        properties_entry.uncompressed_length) != 0) {
-        LOG(ERROR) << "Can't extract " << AB_OTA_PAYLOAD_PROPERTIES;
-        return INSTALL_CORRUPT;
-    }
+  // For A/B updates we extract the payload properties to a buffer and obtain the RAW payload offset
+  // in the zip file.
+  static constexpr const char* AB_OTA_PAYLOAD_PROPERTIES = "payload_properties.txt";
+  ZipString property_name(AB_OTA_PAYLOAD_PROPERTIES);
+  ZipEntry properties_entry;
+  if (FindEntry(zip, property_name, &properties_entry) != 0) {
+    LOG(ERROR) << "Failed to find " << AB_OTA_PAYLOAD_PROPERTIES;
+    return INSTALL_CORRUPT;
+  }
+  uint32_t properties_entry_length = properties_entry.uncompressed_length;
+  std::vector<uint8_t> payload_properties(properties_entry_length);
+  int32_t err =
+      ExtractToMemory(zip, &properties_entry, payload_properties.data(), properties_entry_length);
+  if (err != 0) {
+    LOG(ERROR) << "Failed to extract " << AB_OTA_PAYLOAD_PROPERTIES << ": " << ErrorCodeString(err);
+    return INSTALL_CORRUPT;
+  }
 
-    ZipString payload_name(AB_OTA_PAYLOAD);
-    ZipEntry payload_entry;
-    if (FindEntry(zip, payload_name, &payload_entry) != 0) {
-        LOG(ERROR) << "Can't find " << AB_OTA_PAYLOAD;
-        return INSTALL_CORRUPT;
-    }
-    long payload_offset = payload_entry.offset;
-    *cmd = {
-        "/sbin/update_engine_sideload",
-        android::base::StringPrintf("--payload=file://%s", path),
-        android::base::StringPrintf("--offset=%ld", payload_offset),
-        "--headers=" + std::string(payload_properties.begin(),
-                                   payload_properties.end()),
-        android::base::StringPrintf("--status_fd=%d", status_fd),
-    };
-    return 0;
+  static constexpr const char* AB_OTA_PAYLOAD = "payload.bin";
+  ZipString payload_name(AB_OTA_PAYLOAD);
+  ZipEntry payload_entry;
+  if (FindEntry(zip, payload_name, &payload_entry) != 0) {
+    LOG(ERROR) << "Failed to find " << AB_OTA_PAYLOAD;
+    return INSTALL_CORRUPT;
+  }
+  long payload_offset = payload_entry.offset;
+  *cmd = {
+    "/sbin/update_engine_sideload",
+    "--payload=file://" + path,
+    android::base::StringPrintf("--offset=%ld", payload_offset),
+    "--headers=" + std::string(payload_properties.begin(), payload_properties.end()),
+    android::base::StringPrintf("--status_fd=%d", status_fd),
+  };
+  return 0;
 }
 
 #else  // !AB_OTA_UPDATER
 
-static int
-update_binary_command(const char* path, ZipArchiveHandle zip, int retry_count,
-                      int status_fd, std::vector<std::string>* cmd)
-{
-    // On traditional updates we extract the update binary from the package.
-    ZipString binary_name(ASSUMED_UPDATE_BINARY_NAME);
-    ZipEntry binary_entry;
-    if (FindEntry(zip, binary_name, &binary_entry) != 0) {
-        return INSTALL_CORRUPT;
-    }
+int update_binary_command(const std::string& path, ZipArchiveHandle zip, int retry_count,
+                          int status_fd, std::vector<std::string>* cmd) {
+  CHECK(cmd != nullptr);
 
-    const char* binary = "/tmp/update_binary";
-    unlink(binary);
-    int fd = creat(binary, 0755);
-    if (fd < 0) {
-        PLOG(ERROR) << "Can't make " << binary;
-        return INSTALL_ERROR;
-    }
-    int error = ExtractEntryToFile(zip, &binary_entry, fd);
-    close(fd);
+  // On traditional updates we extract the update binary from the package.
+  static constexpr const char* UPDATE_BINARY_NAME = "META-INF/com/google/android/update-binary";
+  ZipString binary_name(UPDATE_BINARY_NAME);
+  ZipEntry binary_entry;
+  if (FindEntry(zip, binary_name, &binary_entry) != 0) {
+    LOG(ERROR) << "Failed to find update binary " << UPDATE_BINARY_NAME;
+    return INSTALL_CORRUPT;
+  }
 
-    if (error != 0) {
-        LOG(ERROR) << "Can't copy " << ASSUMED_UPDATE_BINARY_NAME
-                   << " : " << ErrorCodeString(error);
-        return INSTALL_ERROR;
-    }
+  const char* binary = "/tmp/update_binary";
+  unlink(binary);
+  int fd = creat(binary, 0755);
+  if (fd == -1) {
+    PLOG(ERROR) << "Failed to create " << binary;
+    return INSTALL_ERROR;
+  }
 
-    *cmd = {
-        binary,
-        EXPAND(RECOVERY_API_VERSION),   // defined in Android.mk
-        std::to_string(status_fd),
-        path,
-    };
-    if (retry_count > 0)
-        cmd->push_back("retry");
-    return 0;
+  int32_t error = ExtractEntryToFile(zip, &binary_entry, fd);
+  close(fd);
+  if (error != 0) {
+    LOG(ERROR) << "Failed to extract " << UPDATE_BINARY_NAME << ": " << ErrorCodeString(error);
+    return INSTALL_ERROR;
+  }
+
+  *cmd = {
+    binary,
+    EXPAND(RECOVERY_API_VERSION),  // defined in Android.mk
+    std::to_string(status_fd),
+    path,
+  };
+  if (retry_count > 0) {
+    cmd->push_back("retry");
+  }
+  return 0;
 }
 #endif  // !AB_OTA_UPDATER
 
diff --git a/private/install.h b/private/install.h
new file mode 100644
index 000000000..12d303b01
--- /dev/null
+++ b/private/install.h
@@ -0,0 +1,27 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// Private headers exposed for testing purpose only.
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#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);
diff --git a/tests/Android.mk b/tests/Android.mk
index 1e433a76e..22a8efcbb 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -86,6 +86,11 @@ LOCAL_CFLAGS := \
     -D_FILE_OFFSET_BITS=64
 
 LOCAL_ADDITIONAL_DEPENDENCIES := $(LOCAL_PATH)/Android.mk
+
+ifeq ($(AB_OTA_UPDATER),true)
+LOCAL_CFLAGS += -DAB_OTA_UPDATER=1
+endif
+
 LOCAL_MODULE := recovery_component_test
 LOCAL_C_INCLUDES := bootable/recovery
 LOCAL_SRC_FILES := \
diff --git a/tests/component/install_test.cpp b/tests/component/install_test.cpp
index 3b6fbc301..fd3b28b07 100644
--- a/tests/component/install_test.cpp
+++ b/tests/component/install_test.cpp
@@ -16,12 +16,18 @@
 
 #include <stdio.h>
 
+#include <string>
+#include <vector>
+
+#include <android-base/properties.h>
+#include <android-base/strings.h>
 #include <android-base/test_utils.h>
 #include <gtest/gtest.h>
 #include <ziparchive/zip_archive.h>
 #include <ziparchive/zip_writer.h>
 
 #include "install.h"
+#include "private/install.h"
 
 TEST(InstallTest, verify_package_compatibility_no_entry) {
   TemporaryFile temp_file;
@@ -55,3 +61,84 @@ TEST(InstallTest, verify_package_compatibility_invalid_entry) {
   ASSERT_FALSE(verify_package_compatibility(zip));
   CloseArchive(zip);
 }
+
+TEST(InstallTest, update_binary_command_smoke) {
+#ifdef AB_OTA_UPDATER
+  TemporaryFile temp_file;
+  FILE* zip_file = fdopen(temp_file.fd, "w");
+  ZipWriter writer(zip_file);
+  ASSERT_EQ(0, writer.StartEntry("payload.bin", kCompressStored));
+  ASSERT_EQ(0, writer.FinishEntry());
+  ASSERT_EQ(0, writer.StartEntry("payload_properties.txt", kCompressStored));
+  const std::string properties = "some_properties";
+  ASSERT_EQ(0, writer.WriteBytes(properties.data(), properties.size()));
+  ASSERT_EQ(0, writer.FinishEntry());
+  // A metadata entry is mandatory.
+  ASSERT_EQ(0, writer.StartEntry("META-INF/com/android/metadata", kCompressStored));
+  std::string device = android::base::GetProperty("ro.product.device", "");
+  ASSERT_NE("", device);
+  std::string timestamp = android::base::GetProperty("ro.build.date.utc", "");
+  ASSERT_NE("", timestamp);
+  std::string metadata = android::base::Join(
+      std::vector<std::string>{
+          "ota-type=AB", "pre-device=" + device, "post-timestamp=" + timestamp,
+      },
+      "\n");
+  ASSERT_EQ(0, writer.WriteBytes(metadata.data(), metadata.size()));
+  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 path = "/path/to/update.zip";
+  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("--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.";
+#endif  // AB_OTA_UPDATER
+}
+
+TEST(InstallTest, update_binary_command_invalid) {
+#ifdef AB_OTA_UPDATER
+  TemporaryFile temp_file;
+  FILE* zip_file = fdopen(temp_file.fd, "w");
+  ZipWriter writer(zip_file);
+  // Missing payload_properties.txt.
+  ASSERT_EQ(0, writer.StartEntry("payload.bin", kCompressStored));
+  ASSERT_EQ(0, writer.FinishEntry());
+  // A metadata entry is mandatory.
+  ASSERT_EQ(0, writer.StartEntry("META-INF/com/android/metadata", kCompressStored));
+  std::string device = android::base::GetProperty("ro.product.device", "");
+  ASSERT_NE("", device);
+  std::string timestamp = android::base::GetProperty("ro.build.date.utc", "");
+  ASSERT_NE("", timestamp);
+  std::string metadata = android::base::Join(
+      std::vector<std::string>{
+          "ota-type=AB", "pre-device=" + device, "post-timestamp=" + timestamp,
+      },
+      "\n");
+  ASSERT_EQ(0, writer.WriteBytes(metadata.data(), metadata.size()));
+  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 path = "/path/to/update.zip";
+  std::vector<std::string> cmd;
+  ASSERT_EQ(INSTALL_CORRUPT, update_binary_command(path, zip, 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.";
+#endif  // AB_OTA_UPDATER
+}