3 files changed, 130 insertions, 45 deletions

diff --git a/install.h b/install.h
index c6db1d1d9..20a2b23e3 100644
--- a/install.h
+++ b/install.h
@@ -21,6 +21,7 @@
 
 #include <map>
 #include <string>
+#include <vector>
 
 #include <ziparchive/zip_archive.h>
 
@@ -53,6 +54,9 @@ bool verify_package(const unsigned char* package_data, size_t package_size);
 // result to |metadata|. Return true if succeed, otherwise return false.
 bool ReadMetadataFromPackage(ZipArchiveHandle zip, std::map<std::string, std::string>* metadata);
 
+// Reads the "recovery.wipe" entry in the zip archive returns a list of partitions to wipe.
+std::vector<std::string> GetWipePartitionList(const std::string& wipe_package);
+
 // Verifies the compatibility info in a Treble-compatible package. Returns true directly if the
 // entry doesn't exist.
 bool verify_package_compatibility(ZipArchiveHandle package_zip);
diff --git a/recovery.cpp b/recovery.cpp
index 703923e81..90c84878b 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -497,45 +497,105 @@ static bool secure_wipe_partition(const std::string& partition) {
   return true;
 }
 
-// Check if the wipe package matches expectation:
+static std::string ReadWipePackage(size_t wipe_package_size) {
+  if (wipe_package_size == 0) {
+    LOG(ERROR) << "wipe_package_size is zero";
+    return "";
+  }
+
+  std::string wipe_package;
+  std::string err_str;
+  if (!read_wipe_package(&wipe_package, wipe_package_size, &err_str)) {
+    PLOG(ERROR) << "Failed to read wipe package" << err_str;
+    return "";
+  }
+  return wipe_package;
+}
+
+// Checks if the wipe package matches expectation. If the check passes, reads the list of
+// partitions to wipe from the package. Checks include
 // 1. verify the package.
 // 2. check metadata (ota-type, pre-device and serial number if having one).
-static bool check_wipe_package(size_t wipe_package_size) {
-    if (wipe_package_size == 0) {
-        LOG(ERROR) << "wipe_package_size is zero";
-        return false;
-    }
-    std::string wipe_package;
-    std::string err_str;
-    if (!read_wipe_package(&wipe_package, wipe_package_size, &err_str)) {
-        PLOG(ERROR) << "Failed to read wipe package";
-        return false;
-    }
-    if (!verify_package(reinterpret_cast<const unsigned char*>(wipe_package.data()),
-                        wipe_package.size())) {
-        LOG(ERROR) << "Failed to verify package";
-        return false;
-    }
+static bool CheckWipePackage(const std::string& wipe_package) {
+  if (!verify_package(reinterpret_cast<const unsigned char*>(wipe_package.data()),
+                      wipe_package.size())) {
+    LOG(ERROR) << "Failed to verify package";
+    return false;
+  }
 
-    // Extract metadata
-    ZipArchiveHandle zip;
-    int err = OpenArchiveFromMemory(static_cast<void*>(&wipe_package[0]), wipe_package.size(),
-                                    "wipe_package", &zip);
-    if (err != 0) {
-        LOG(ERROR) << "Can't open wipe package : " << ErrorCodeString(err);
-        return false;
-    }
+  // Extract metadata
+  ZipArchiveHandle zip;
+  if (auto err =
+          OpenArchiveFromMemory(const_cast<void*>(static_cast<const void*>(&wipe_package[0])),
+                                wipe_package.size(), "wipe_package", &zip);
+      err != 0) {
+    LOG(ERROR) << "Can't open wipe package : " << ErrorCodeString(err);
+    return false;
+  }
 
-    std::map<std::string, std::string> metadata;
-    if (!ReadMetadataFromPackage(zip, &metadata)) {
-      LOG(ERROR) << "Failed to parse metadata in the zip file";
-      return false;
+  std::map<std::string, std::string> metadata;
+  if (!ReadMetadataFromPackage(zip, &metadata)) {
+    LOG(ERROR) << "Failed to parse metadata in the zip file";
+    return false;
+  }
+
+  int result = CheckPackageMetadata(metadata, OtaType::BRICK);
+  CloseArchive(zip);
+
+  return result == 0;
+}
+
+std::vector<std::string> GetWipePartitionList(const std::string& wipe_package) {
+  ZipArchiveHandle zip;
+  if (auto err =
+          OpenArchiveFromMemory(const_cast<void*>(static_cast<const void*>(&wipe_package[0])),
+                                wipe_package.size(), "wipe_package", &zip);
+      err != 0) {
+    LOG(ERROR) << "Can't open wipe package : " << ErrorCodeString(err);
+    return {};
+  }
+
+  static constexpr const char* RECOVERY_WIPE_ENTRY_NAME = "recovery.wipe";
+
+  std::string partition_list_content;
+  ZipString path(RECOVERY_WIPE_ENTRY_NAME);
+  ZipEntry entry;
+  if (FindEntry(zip, path, &entry) == 0) {
+    uint32_t length = entry.uncompressed_length;
+    partition_list_content = std::string(length, '\0');
+    if (auto err = ExtractToMemory(
+            zip, &entry, reinterpret_cast<uint8_t*>(partition_list_content.data()), length);
+        err != 0) {
+      LOG(ERROR) << "Failed to extract " << RECOVERY_WIPE_ENTRY_NAME << ": "
+                 << ErrorCodeString(err);
+      CloseArchive(zip);
+      return {};
     }
+  } else {
+    LOG(INFO) << "Failed to find " << RECOVERY_WIPE_ENTRY_NAME
+              << ", falling back to use the partition list on device.";
+
+    static constexpr const char* RECOVERY_WIPE_ON_DEVICE = "/etc/recovery.wipe";
+    if (!android::base::ReadFileToString(RECOVERY_WIPE_ON_DEVICE, &partition_list_content)) {
+      PLOG(ERROR) << "failed to read \"" << RECOVERY_WIPE_ON_DEVICE << "\"";
+      CloseArchive(zip);
+      return {};
+    }
+  }
 
-    int result = CheckPackageMetadata(metadata, OtaType::BRICK);
-    CloseArchive(zip);
+  std::vector<std::string> result;
+  std::vector<std::string> lines = android::base::Split(partition_list_content, "\n");
+  for (const std::string& line : lines) {
+    std::string partition = android::base::Trim(line);
+    // Ignore '#' comment or empty lines.
+    if (android::base::StartsWith(partition, "#") || partition.empty()) {
+      continue;
+    }
+    result.push_back(line);
+  }
 
-    return result == 0;
+  CloseArchive(zip);
+  return result;
 }
 
 // Wipes the current A/B device, with a secure wipe of all the partitions in RECOVERY_WIPE.
@@ -543,25 +603,23 @@ static bool wipe_ab_device(size_t wipe_package_size) {
   ui->SetBackground(RecoveryUI::ERASING);
   ui->SetProgressType(RecoveryUI::INDETERMINATE);
 
-  if (!check_wipe_package(wipe_package_size)) {
-    LOG(ERROR) << "Failed to verify wipe package";
+  std::string wipe_package = ReadWipePackage(wipe_package_size);
+  if (wipe_package.empty()) {
     return false;
   }
-  static constexpr const char* RECOVERY_WIPE = "/etc/recovery.wipe";
-  std::string partition_list;
-  if (!android::base::ReadFileToString(RECOVERY_WIPE, &partition_list)) {
-    LOG(ERROR) << "failed to read \"" << RECOVERY_WIPE << "\"";
+
+  if (!CheckWipePackage(wipe_package)) {
+    LOG(ERROR) << "Failed to verify wipe package";
     return false;
   }
 
-  std::vector<std::string> lines = android::base::Split(partition_list, "\n");
-  for (const std::string& line : lines) {
-    std::string partition = android::base::Trim(line);
-    // Ignore '#' comment or empty lines.
-    if (android::base::StartsWith(partition, "#") || partition.empty()) {
-      continue;
-    }
+  std::vector<std::string> partition_list = GetWipePartitionList(wipe_package);
+  if (partition_list.empty()) {
+    LOG(ERROR) << "Empty wipe ab partition list";
+    return false;
+  }
 
+  for (const auto& partition : partition_list) {
     // Proceed anyway even if it fails to wipe some partition.
     secure_wipe_partition(partition);
   }
diff --git a/tests/component/install_test.cpp b/tests/component/install_test.cpp
index 47a54714c..117813694 100644
--- a/tests/component/install_test.cpp
+++ b/tests/component/install_test.cpp
@@ -107,6 +107,29 @@ TEST(InstallTest, read_metadata_from_package_no_entry) {
   CloseArchive(zip);
 }
 
+TEST(InstallTest, read_wipe_ab_partition_list) {
+  std::vector<std::string> partition_list = {
+    "/dev/block/bootdevice/by-name/system_a", "/dev/block/bootdevice/by-name/system_b",
+    "/dev/block/bootdevice/by-name/vendor_a", "/dev/block/bootdevice/by-name/vendor_b",
+    "/dev/block/bootdevice/by-name/userdata", "# Wipe the boot partitions last",
+    "/dev/block/bootdevice/by-name/boot_a",   "/dev/block/bootdevice/by-name/boot_b",
+  };
+  TemporaryFile temp_file;
+  BuildZipArchive({ { "recovery.wipe", android::base::Join(partition_list, '\n') } },
+                  temp_file.release(), kCompressDeflated);
+  std::string wipe_package;
+  ASSERT_TRUE(android::base::ReadFileToString(temp_file.path, &wipe_package));
+
+  std::vector<std::string> read_partition_list = GetWipePartitionList(wipe_package);
+  std::vector<std::string> expected = {
+    "/dev/block/bootdevice/by-name/system_a", "/dev/block/bootdevice/by-name/system_b",
+    "/dev/block/bootdevice/by-name/vendor_a", "/dev/block/bootdevice/by-name/vendor_b",
+    "/dev/block/bootdevice/by-name/userdata", "/dev/block/bootdevice/by-name/boot_a",
+    "/dev/block/bootdevice/by-name/boot_b",
+  };
+  ASSERT_EQ(expected, read_partition_list);
+}
+
 TEST(InstallTest, verify_package_compatibility_with_libvintf_malformed_xml) {
   TemporaryFile compatibility_zip_file;
   std::string malformed_xml = "malformed";