6 files changed, 197 insertions, 137 deletions

diff --git a/error_code.h b/error_code.h
index dfea0eb38..5dad6b263 100644
--- a/error_code.h
+++ b/error_code.h
@@ -18,52 +18,53 @@
 #define _ERROR_CODE_H_
 
 enum ErrorCode {
-    kNoError = -1,
-    kLowBattery = 20,
-    kZipVerificationFailure,
-    kZipOpenFailure,
-    kBootreasonInBlacklist
+  kNoError = -1,
+  kLowBattery = 20,
+  kZipVerificationFailure,
+  kZipOpenFailure,
+  kBootreasonInBlacklist
 };
 
 enum CauseCode {
-    kNoCause = -1,
-    kArgsParsingFailure = 100,
-    kStashCreationFailure,
-    kFileOpenFailure,
-    kLseekFailure,
-    kFreadFailure,
-    kFwriteFailure,
-    kFsyncFailure,
-    kLibfecFailure,
-    kFileGetPropFailure,
-    kFileRenameFailure,
-    kSymlinkFailure,
-    kSetMetadataFailure,
-    kTune2FsFailure,
-    kRebootFailure,
-    kVendorFailure = 200
+  kNoCause = -1,
+  kArgsParsingFailure = 100,
+  kStashCreationFailure,
+  kFileOpenFailure,
+  kLseekFailure,
+  kFreadFailure,
+  kFwriteFailure,
+  kFsyncFailure,
+  kLibfecFailure,
+  kFileGetPropFailure,
+  kFileRenameFailure,
+  kSymlinkFailure,
+  kSetMetadataFailure,
+  kTune2FsFailure,
+  kRebootFailure,
+  kPackageExtractFileFailure,
+  kVendorFailure = 200
 };
 
 enum UncryptErrorCode {
-    kUncryptNoError = -1,
-    kUncryptErrorPlaceholder = 50,
-    kUncryptTimeoutError = 100,
-    kUncryptFileRemoveError,
-    kUncryptFileOpenError,
-    kUncryptSocketOpenError,
-    kUncryptSocketWriteError,
-    kUncryptSocketListenError,
-    kUncryptSocketAcceptError,
-    kUncryptFstabReadError,
-    kUncryptFileStatError,
-    kUncryptBlockOpenError,
-    kUncryptIoctlError,
-    kUncryptReadError,
-    kUncryptWriteError,
-    kUncryptFileSyncError,
-    kUncryptFileCloseError,
-    kUncryptFileRenameError,
-    kUncryptPackageMissingError,
+  kUncryptNoError = -1,
+  kUncryptErrorPlaceholder = 50,
+  kUncryptTimeoutError = 100,
+  kUncryptFileRemoveError,
+  kUncryptFileOpenError,
+  kUncryptSocketOpenError,
+  kUncryptSocketWriteError,
+  kUncryptSocketListenError,
+  kUncryptSocketAcceptError,
+  kUncryptFstabReadError,
+  kUncryptFileStatError,
+  kUncryptBlockOpenError,
+  kUncryptIoctlError,
+  kUncryptReadError,
+  kUncryptWriteError,
+  kUncryptFileSyncError,
+  kUncryptFileCloseError,
+  kUncryptFileRenameError,
+  kUncryptPackageMissingError,
 };
 
-#endif
+#endif // _ERROR_CODE_H_
diff --git a/tests/Android.mk b/tests/Android.mk
index e87a22964..fdc947028 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -92,6 +92,8 @@ LOCAL_STATIC_LIBRARIES := \
     libcrypto \
     libcutils \
     libbz \
+    libziparchive \
+    libutils \
     libz \
     libbase \
     libtune2fs \
diff --git a/tests/common/test_constants.h b/tests/common/test_constants.h
index 97e74a3c2..93e4ab5b2 100644
--- a/tests/common/test_constants.h
+++ b/tests/common/test_constants.h
@@ -19,6 +19,15 @@
 
 #include <stdlib.h>
 
+// Zip entries in ziptest_valid.zip.
+static const std::string kATxtContents("abcdefghabcdefgh\n");
+static const std::string kBTxtContents("abcdefgh\n");
+
+// echo -n -e "abcdefghabcdefgh\n" | sha1sum
+static const std::string kATxtSha1Sum("32c96a03dc8cd20097940f351bca6261ee5a1643");
+// echo -n -e "abcdefgh\n" | sha1sum
+static const std::string kBTxtSha1Sum("e414af7161c9554089f4106d6f1797ef14a73666");
+
 static const char* data_root = getenv("ANDROID_DATA");
 
 static std::string from_testdata_base(const std::string& fname) {
diff --git a/tests/component/updater_test.cpp b/tests/component/updater_test.cpp
index 973c19de3..a029cf449 100644
--- a/tests/component/updater_test.cpp
+++ b/tests/component/updater_test.cpp
@@ -24,35 +24,40 @@
 #include <android-base/properties.h>
 #include <android-base/test_utils.h>
 #include <gtest/gtest.h>
+#include <ziparchive/zip_archive.h>
 
+#include "common/test_constants.h"
 #include "edify/expr.h"
 #include "error_code.h"
 #include "updater/install.h"
+#include "updater/updater.h"
 
 struct selabel_handle *sehandle = nullptr;
 
-static void expect(const char* expected, const char* expr_str, CauseCode cause_code) {
-    Expr* e;
-    int error_count;
-    EXPECT_EQ(parse_string(expr_str, &e, &error_count), 0);
+static void expect(const char* expected, const char* expr_str, CauseCode cause_code,
+                   UpdaterInfo* info = nullptr) {
+  Expr* e;
+  int error_count = 0;
+  ASSERT_EQ(0, parse_string(expr_str, &e, &error_count));
+  ASSERT_EQ(0, error_count);
 
-    State state(expr_str, nullptr);
+  State state(expr_str, info);
 
-    std::string result;
-    bool status = Evaluate(&state, e, &result);
+  std::string result;
+  bool status = Evaluate(&state, e, &result);
 
-    if (expected == nullptr) {
-        EXPECT_FALSE(status);
-    } else {
-        EXPECT_STREQ(expected, result.c_str());
-    }
-
-    // Error code is set in updater/updater.cpp only, by parsing State.errmsg.
-    EXPECT_EQ(kNoError, state.error_code);
+  if (expected == nullptr) {
+    ASSERT_FALSE(status);
+  } else {
+    ASSERT_TRUE(status);
+    ASSERT_STREQ(expected, result.c_str());
+  }
 
-    // Cause code should always be available.
-    EXPECT_EQ(cause_code, state.cause_code);
+  // Error code is set in updater/updater.cpp only, by parsing State.errmsg.
+  ASSERT_EQ(kNoError, state.error_code);
 
+  // Cause code should always be available.
+  ASSERT_EQ(cause_code, state.cause_code);
 }
 
 class UpdaterTest : public ::testing::Test {
@@ -264,3 +269,56 @@ TEST_F(UpdaterTest, symlink) {
     ASSERT_EQ(0, unlink(src1.c_str()));
     ASSERT_EQ(0, unlink(src2.c_str()));
 }
+
+// TODO: Test extracting to block device.
+TEST_F(UpdaterTest, package_extract_file) {
+  // package_extract_file expects 1 or 2 arguments.
+  expect(nullptr, "package_extract_file()", kArgsParsingFailure);
+  expect(nullptr, "package_extract_file(\"arg1\", \"arg2\", \"arg3\")", kArgsParsingFailure);
+
+  std::string zip_path = from_testdata_base("ziptest_valid.zip");
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchive(zip_path.c_str(), &handle));
+
+  // Need to set up the ziphandle.
+  UpdaterInfo updater_info;
+  updater_info.package_zip = handle;
+
+  // Two-argument version.
+  TemporaryFile temp_file1;
+  std::string script("package_extract_file(\"a.txt\", \"" + std::string(temp_file1.path) + "\")");
+  expect("t", script.c_str(), kNoCause, &updater_info);
+
+  // Verify the extracted entry.
+  std::string data;
+  ASSERT_TRUE(android::base::ReadFileToString(temp_file1.path, &data));
+  ASSERT_EQ(kATxtContents, data);
+
+  // Now extract another entry to the same location, which should overwrite.
+  script = "package_extract_file(\"b.txt\", \"" + std::string(temp_file1.path) + "\")";
+  expect("t", script.c_str(), kNoCause, &updater_info);
+
+  ASSERT_TRUE(android::base::ReadFileToString(temp_file1.path, &data));
+  ASSERT_EQ(kBTxtContents, data);
+
+  // Missing zip entry. The two-argument version doesn't abort.
+  script = "package_extract_file(\"doesntexist\", \"" + std::string(temp_file1.path) + "\")";
+  expect("", script.c_str(), kNoCause, &updater_info);
+
+  // Extract to /dev/full should fail.
+  script = "package_extract_file(\"a.txt\", \"/dev/full\")";
+  expect("", script.c_str(), kNoCause, &updater_info);
+
+  // One-argument version.
+  script = "sha1_check(package_extract_file(\"a.txt\"))";
+  expect(kATxtSha1Sum.c_str(), script.c_str(), kNoCause, &updater_info);
+
+  script = "sha1_check(package_extract_file(\"b.txt\"))";
+  expect(kBTxtSha1Sum.c_str(), script.c_str(), kNoCause, &updater_info);
+
+  // Missing entry. The one-argument version aborts the evaluation.
+  script = "package_extract_file(\"doesntexist\")";
+  expect(nullptr, script.c_str(), kPackageExtractFileFailure, &updater_info);
+
+  CloseArchive(handle);
+}
diff --git a/tests/unit/zip_test.cpp b/tests/unit/zip_test.cpp
index 49729467d..ef0ee4c1d 100644
--- a/tests/unit/zip_test.cpp
+++ b/tests/unit/zip_test.cpp
@@ -30,9 +30,6 @@
 
 #include "common/test_constants.h"
 
-static const std::string kATxtContents("abcdefghabcdefgh\n");
-static const std::string kBTxtContents("abcdefgh\n");
-
 TEST(ZipTest, ExtractPackageRecursive) {
   std::string zip_path = from_testdata_base("ziptest_valid.zip");
   ZipArchiveHandle handle;
diff --git a/updater/install.cpp b/updater/install.cpp
index 59c54dd4c..b885f864e 100644
--- a/updater/install.cpp
+++ b/updater/install.cpp
@@ -477,93 +477,86 @@ Value* PackageExtractDirFn(const char* name, State* state,
     return StringValue(success ? "t" : "");
 }
 
+// package_extract_file(package_file[, dest_file])
+//   Extracts a single package_file from the update package and writes it to dest_file,
+//   overwriting existing files if necessary. Without the dest_file argument, returns the
+//   contents of the package file as a binary blob.
+Value* PackageExtractFileFn(const char* name, State* state, int argc, Expr* argv[]) {
+  if (argc < 1 || argc > 2) {
+    return ErrorAbort(state, kArgsParsingFailure, "%s() expects 1 or 2 args, got %d", name, argc);
+  }
+
+  if (argc == 2) {
+    // The two-argument version extracts to a file.
 
-// package_extract_file(package_path, destination_path)
-//   or
-// package_extract_file(package_path)
-//   to return the entire contents of the file as the result of this
-//   function (the char* returned is actually a FileContents*).
-Value* PackageExtractFileFn(const char* name, State* state,
-                           int argc, Expr* argv[]) {
-    if (argc < 1 || argc > 2) {
-        return ErrorAbort(state, kArgsParsingFailure, "%s() expects 1 or 2 args, got %d",
-                          name, argc);
+    std::vector<std::string> args;
+    if (!ReadArgs(state, 2, argv, &args)) {
+      return ErrorAbort(state, kArgsParsingFailure, "%s() Failed to parse %d args", name, argc);
     }
-    bool success = false;
-
-    if (argc == 2) {
-        // The two-argument version extracts to a file.
-
-        ZipArchiveHandle za = ((UpdaterInfo*)(state->cookie))->package_zip;
-
-        std::vector<std::string> args;
-        if (!ReadArgs(state, 2, argv, &args)) {
-            return ErrorAbort(state, kArgsParsingFailure, "%s() Failed to parse %d args", name,
-                              argc);
-        }
-        const std::string& zip_path = args[0];
-        const std::string& dest_path = args[1];
+    const std::string& zip_path = args[0];
+    const std::string& dest_path = args[1];
 
-        ZipString zip_string_path(zip_path.c_str());
-        ZipEntry entry;
-        if (FindEntry(za, zip_string_path, &entry) != 0) {
-            printf("%s: no %s in package\n", name, zip_path.c_str());
-            return StringValue("");
-        }
+    ZipArchiveHandle za = static_cast<UpdaterInfo*>(state->cookie)->package_zip;
+    ZipString zip_string_path(zip_path.c_str());
+    ZipEntry entry;
+    if (FindEntry(za, zip_string_path, &entry) != 0) {
+      printf("%s: no %s in package\n", name, zip_path.c_str());
+      return StringValue("");
+    }
 
-        int fd = TEMP_FAILURE_RETRY(ota_open(dest_path.c_str(), O_WRONLY | O_CREAT | O_TRUNC,
-              S_IRUSR | S_IWUSR));
-        if (fd == -1) {
-            printf("%s: can't open %s for write: %s\n", name, dest_path.c_str(), strerror(errno));
-            return StringValue("");
-        }
-        success = ExtractEntryToFile(za, &entry, fd);
-        if (ota_fsync(fd) == -1) {
-            printf("fsync of \"%s\" failed: %s\n", dest_path.c_str(), strerror(errno));
-            success = false;
-        }
-        if (ota_close(fd) == -1) {
-            printf("close of \"%s\" failed: %s\n", dest_path.c_str(), strerror(errno));
-            success = false;
-        }
+    int fd = TEMP_FAILURE_RETRY(
+        ota_open(dest_path.c_str(), O_WRONLY | O_CREAT | O_TRUNC, S_IRUSR | S_IWUSR));
+    if (fd == -1) {
+      printf("%s: can't open %s for write: %s\n", name, dest_path.c_str(), strerror(errno));
+      return StringValue("");
+    }
 
-        return StringValue(success ? "t" : "");
-    } else {
-        // The one-argument version returns the contents of the file
-        // as the result.
+    bool success = true;
+    int32_t ret = ExtractEntryToFile(za, &entry, fd);
+    if (ret != 0) {
+      printf("%s: Failed to extract entry \"%s\" (%u bytes) to \"%s\": %s\n", name,
+             zip_path.c_str(), entry.uncompressed_length, dest_path.c_str(), ErrorCodeString(ret));
+      success = false;
+    }
+    if (ota_fsync(fd) == -1) {
+      printf("fsync of \"%s\" failed: %s\n", dest_path.c_str(), strerror(errno));
+      success = false;
+    }
+    if (ota_close(fd) == -1) {
+      printf("close of \"%s\" failed: %s\n", dest_path.c_str(), strerror(errno));
+      success = false;
+    }
 
-        std::vector<std::string> args;
-        if (!ReadArgs(state, 1, argv, &args)) {
-            return ErrorAbort(state, kArgsParsingFailure, "%s() Failed to parse %d args", name,
-                              argc);
-        }
-        const std::string& zip_path = args[0];
+    return StringValue(success ? "t" : "");
+  } else {
+    // The one-argument version returns the contents of the file as the result.
 
-        Value* v = new Value(VAL_INVALID, "");
+    std::vector<std::string> args;
+    if (!ReadArgs(state, 1, argv, &args)) {
+      return ErrorAbort(state, kArgsParsingFailure, "%s() Failed to parse %d args", name, argc);
+    }
+    const std::string& zip_path = args[0];
 
-        ZipArchiveHandle za = ((UpdaterInfo*)(state->cookie))->package_zip;
-        ZipString zip_string_path(zip_path.c_str());
-        ZipEntry entry;
-        if (FindEntry(za, zip_string_path, &entry) != 0) {
-            printf("%s: no %s in package\n", name, zip_path.c_str());
-            return v;
-        }
+    ZipArchiveHandle za = static_cast<UpdaterInfo*>(state->cookie)->package_zip;
+    ZipString zip_string_path(zip_path.c_str());
+    ZipEntry entry;
+    if (FindEntry(za, zip_string_path, &entry) != 0) {
+      return ErrorAbort(state, kPackageExtractFileFailure, "%s(): no %s in package", name,
+                        zip_path.c_str());
+    }
 
-        v->data.resize(entry.uncompressed_length);
-        if (ExtractToMemory(za, &entry, reinterpret_cast<uint8_t*>(&v->data[0]),
-                            v->data.size()) != 0) {
-            printf("%s: faled to extract %zu bytes to memory\n", name, v->data.size());
-        } else {
-            success = true;
-        }
+    std::string buffer;
+    buffer.resize(entry.uncompressed_length);
 
-        if (!success) {
-            v->data.clear();
-        } else {
-            v->type = VAL_BLOB;
-        }
-        return v;
+    int32_t ret = ExtractToMemory(za, &entry, reinterpret_cast<uint8_t*>(&buffer[0]), buffer.size());
+    if (ret != 0) {
+      return ErrorAbort(state, kPackageExtractFileFailure,
+                        "%s: Failed to extract entry \"%s\" (%zu bytes) to memory: %s", name,
+                        zip_path.c_str(), buffer.size(), ErrorCodeString(ret));
     }
+
+    return new Value(VAL_BLOB, buffer);
+  }
 }
 
 // symlink(target, [src1, src2, ...])