[libziparchive] add an option to start iteration with functor

To reduce the seeks for local file headers in large APK files, we can
specify entry prefix/suffix when we call StartIteration(). However,
some use cases need additional name matches that is outside the
prefix/suffix matches.

Adding a new option to StartIteration, which allows additional functor
that restricts the iteration to customized name matching schemes.

Test: atest ziparchive-tests
BUG: 151676293
Change-Id: Iff45e083b334602f183c05cb39ba521e7070252c
diff --git a/libziparchive/include/ziparchive/zip_archive.h b/libziparchive/include/ziparchive/zip_archive.h
index a17127d..435bfb6 100644
--- a/libziparchive/include/ziparchive/zip_archive.h
+++ b/libziparchive/include/ziparchive/zip_archive.h
@@ -188,6 +188,15 @@
                        const std::string_view optional_suffix = "");
 
 /*
+ * Start iterating over all entries of a zip file. Use the matcher functor to
+ * restrict iteration to entry names that make the functor return true.
+ *
+ * Returns 0 on success and negative values on failure.
+ */
+int32_t StartIteration(ZipArchiveHandle archive, void** cookie_ptr,
+                       std::function<bool(std::string_view entry_name)> matcher);
+
+/*
  * Advance to the next element in the zipfile in iteration order.
  *
  * Returns 0 on success, -1 if there are no more elements in this
diff --git a/libziparchive/zip_archive.cc b/libziparchive/zip_archive.cc
index afbc5d8..6b6502b 100644
--- a/libziparchive/zip_archive.cc
+++ b/libziparchive/zip_archive.cc
@@ -32,6 +32,7 @@
 #include <unistd.h>
 
 #include <memory>
+#include <optional>
 #include <vector>
 
 #if defined(__APPLE__)
@@ -675,31 +676,40 @@
 struct IterationHandle {
   ZipArchive* archive;
 
-  std::string prefix;
-  std::string suffix;
+  std::function<bool(std::string_view)> matcher;
 
   uint32_t position = 0;
 
-  IterationHandle(ZipArchive* archive, std::string_view in_prefix, std::string_view in_suffix)
-      : archive(archive), prefix(in_prefix), suffix(in_suffix) {}
+  IterationHandle(ZipArchive* archive, std::function<bool(std::string_view)> in_matcher)
+      : archive(archive), matcher(std::move(in_matcher)) {}
+
+  bool Match(std::string_view entry_name) const { return matcher(entry_name); }
 };
 
 int32_t StartIteration(ZipArchiveHandle archive, void** cookie_ptr,
                        const std::string_view optional_prefix,
                        const std::string_view optional_suffix) {
-  if (archive == nullptr || archive->cd_entry_map == nullptr) {
-    ALOGW("Zip: Invalid ZipArchiveHandle");
-    return kInvalidHandle;
-  }
-
   if (optional_prefix.size() > static_cast<size_t>(UINT16_MAX) ||
       optional_suffix.size() > static_cast<size_t>(UINT16_MAX)) {
     ALOGW("Zip: prefix/suffix too long");
     return kInvalidEntryName;
   }
+  auto matcher = [prefix = std::string(optional_prefix),
+                  suffix = std::string(optional_suffix)](std::string_view name) mutable {
+    return android::base::StartsWith(name, prefix) && android::base::EndsWith(name, suffix);
+  };
+  return StartIteration(archive, cookie_ptr, std::move(matcher));
+}
+
+int32_t StartIteration(ZipArchiveHandle archive, void** cookie_ptr,
+                       std::function<bool(std::string_view)> matcher) {
+  if (archive == nullptr || archive->cd_entry_map == nullptr) {
+    ALOGW("Zip: Invalid ZipArchiveHandle");
+    return kInvalidHandle;
+  }
 
   archive->cd_entry_map->ResetIteration();
-  *cookie_ptr = new IterationHandle(archive, optional_prefix, optional_suffix);
+  *cookie_ptr = new IterationHandle(archive, matcher);
   return 0;
 }
 
@@ -749,8 +759,7 @@
   auto entry = archive->cd_entry_map->Next(archive->central_directory.GetBasePtr());
   while (entry != std::pair<std::string_view, uint64_t>()) {
     const auto [entry_name, offset] = entry;
-    if (android::base::StartsWith(entry_name, handle->prefix) &&
-        android::base::EndsWith(entry_name, handle->suffix)) {
+    if (handle->Match(entry_name)) {
       const int error = FindEntry(archive, entry_name, offset, data);
       if (!error && name) {
         *name = entry_name;
diff --git a/libziparchive/zip_archive_test.cc b/libziparchive/zip_archive_test.cc
index 5caca8a..6c1a9a1 100644
--- a/libziparchive/zip_archive_test.cc
+++ b/libziparchive/zip_archive_test.cc
@@ -227,6 +227,22 @@
   CloseArchive(handle);
 }
 
+static void AssertIterationNames(void* iteration_cookie,
+                                 const std::vector<std::string>& expected_names_sorted) {
+  ZipEntry data;
+  std::vector<std::string> names;
+  std::string name;
+  for (size_t i = 0; i < expected_names_sorted.size(); ++i) {
+    ASSERT_EQ(0, Next(iteration_cookie, &data, &name));
+    names.push_back(name);
+  }
+  // End of iteration.
+  ASSERT_EQ(-1, Next(iteration_cookie, &data, &name));
+  // Assert that the names are as expected.
+  std::sort(names.begin(), names.end());
+  ASSERT_EQ(expected_names_sorted, names);
+}
+
 static void AssertIterationOrder(const std::string_view prefix, const std::string_view suffix,
                                  const std::vector<std::string>& expected_names_sorted) {
   ZipArchiveHandle handle;
@@ -234,23 +250,19 @@
 
   void* iteration_cookie;
   ASSERT_EQ(0, StartIteration(handle, &iteration_cookie, prefix, suffix));
-
-  ZipEntry data;
-  std::vector<std::string> names;
-
-  std::string name;
-  for (size_t i = 0; i < expected_names_sorted.size(); ++i) {
-    ASSERT_EQ(0, Next(iteration_cookie, &data, &name));
-    names.push_back(name);
-  }
-
-  // End of iteration.
-  ASSERT_EQ(-1, Next(iteration_cookie, &data, &name));
+  AssertIterationNames(iteration_cookie, expected_names_sorted);
   CloseArchive(handle);
+}
 
-  // Assert that the names are as expected.
-  std::sort(names.begin(), names.end());
-  ASSERT_EQ(expected_names_sorted, names);
+static void AssertIterationOrderWithMatcher(std::function<bool(std::string_view)> matcher,
+                                            const std::vector<std::string>& expected_names_sorted) {
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchiveWrapper(kValidZip, &handle));
+
+  void* iteration_cookie;
+  ASSERT_EQ(0, StartIteration(handle, &iteration_cookie, matcher));
+  AssertIterationNames(iteration_cookie, expected_names_sorted);
+  CloseArchive(handle);
 }
 
 TEST(ziparchive, Iteration) {
@@ -279,6 +291,30 @@
   AssertIterationOrder("b", ".txt", kExpectedMatchesSorted);
 }
 
+TEST(ziparchive, IterationWithAdditionalMatchesExactly) {
+  static const std::vector<std::string> kExpectedMatchesSorted = {"a.txt"};
+  auto matcher = [](std::string_view name) { return name == "a.txt"; };
+  AssertIterationOrderWithMatcher(matcher, kExpectedMatchesSorted);
+}
+
+TEST(ziparchive, IterationWithAdditionalMatchesWithSuffix) {
+  static const std::vector<std::string> kExpectedMatchesSorted = {"a.txt", "b.txt", "b/c.txt",
+                                                                  "b/d.txt"};
+  auto matcher = [](std::string_view name) {
+    return name == "a.txt" || android::base::EndsWith(name, ".txt");
+  };
+  AssertIterationOrderWithMatcher(matcher, kExpectedMatchesSorted);
+}
+
+TEST(ziparchive, IterationWithAdditionalMatchesWithPrefixAndSuffix) {
+  static const std::vector<std::string> kExpectedMatchesSorted = {"a.txt", "b/c.txt", "b/d.txt"};
+  auto matcher = [](std::string_view name) {
+    return name == "a.txt" ||
+           (android::base::EndsWith(name, ".txt") && android::base::StartsWith(name, "b/"));
+  };
+  AssertIterationOrderWithMatcher(matcher, kExpectedMatchesSorted);
+}
+
 TEST(ziparchive, IterationWithBadPrefixAndSuffix) {
   ZipArchiveHandle handle;
   ASSERT_EQ(0, OpenArchiveWrapper(kValidZip, &handle));