Add |optional_suffix| to StartIteration()

so that PackageManagerService can iterate over files with a
specific file extension like ".so".

Bug: 21957428
Change-Id: I76ed9560d4d1e00d297a97d518ec357166be1981

3 files changed, 144 insertions, 15 deletions

diff --git a/include/ziparchive/zip_archive.h b/include/ziparchive/zip_archive.h
index 386a390..3b00683 100644
--- a/include/ziparchive/zip_archive.h
+++ b/include/ziparchive/zip_archive.h
@@ -153,7 +153,9 @@ int32_t FindEntry(const ZipArchiveHandle handle, const ZipEntryName& entryName,
  * Returns 0 on success and negative values on failure.
  */
 int32_t StartIteration(ZipArchiveHandle handle, void** cookie_ptr,
-                       const ZipEntryName* optional_prefix);
+                       const ZipEntryName* optional_prefix,
+                       // TODO: Remove the default parameter.
+                       const ZipEntryName* optional_suffix = NULL);
 
 /*
  * Advance to the next element in the zipfile in iteration order.
diff --git a/libziparchive/zip_archive.cc b/libziparchive/zip_archive.cc
index b2a9f88..cc39aa5 100644
--- a/libziparchive/zip_archive.cc
+++ b/libziparchive/zip_archive.cc
@@ -852,25 +852,38 @@ struct IterationHandle {
   // We're not using vector here because this code is used in the Windows SDK
   // where the STL is not available.
   const uint8_t* prefix;
-  uint16_t prefix_len;
+  const uint16_t prefix_len;
+  const uint8_t* suffix;
+  const uint16_t suffix_len;
   ZipArchive* archive;
 
-  IterationHandle() : prefix(NULL), prefix_len(0) {}
-
-  IterationHandle(const ZipEntryName& prefix_name)
-      : prefix_len(prefix_name.name_length) {
-    uint8_t* prefix_copy = new uint8_t[prefix_len];
-    memcpy(prefix_copy, prefix_name.name, prefix_len);
-    prefix = prefix_copy;
+  IterationHandle(const ZipEntryName* prefix_name,
+                  const ZipEntryName* suffix_name)
+    : prefix(NULL),
+      prefix_len(prefix_name ? prefix_name->name_length : 0),
+      suffix(NULL),
+      suffix_len(suffix_name ? suffix_name->name_length : 0) {
+    if (prefix_name) {
+      uint8_t* prefix_copy = new uint8_t[prefix_len];
+      memcpy(prefix_copy, prefix_name->name, prefix_len);
+      prefix = prefix_copy;
+    }
+    if (suffix_name) {
+      uint8_t* suffix_copy = new uint8_t[suffix_len];
+      memcpy(suffix_copy, suffix_name->name, suffix_len);
+      suffix = suffix_copy;
+    }
   }
 
   ~IterationHandle() {
     delete[] prefix;
+    delete[] suffix;
   }
 };
 
 int32_t StartIteration(ZipArchiveHandle handle, void** cookie_ptr,
-                       const ZipEntryName* optional_prefix) {
+                       const ZipEntryName* optional_prefix,
+                       const ZipEntryName* optional_suffix) {
   ZipArchive* archive = reinterpret_cast<ZipArchive*>(handle);
 
   if (archive == NULL || archive->hash_table == NULL) {
@@ -878,8 +891,7 @@ int32_t StartIteration(ZipArchiveHandle handle, void** cookie_ptr,
     return kInvalidHandle;
   }
 
-  IterationHandle* cookie =
-      optional_prefix != NULL ? new IterationHandle(*optional_prefix) : new IterationHandle();
+  IterationHandle* cookie = new IterationHandle(optional_prefix, optional_suffix);
   cookie->position = 0;
   cookie->archive = archive;
 
@@ -929,7 +941,13 @@ int32_t Next(void* cookie, ZipEntry* data, ZipEntryName* name) {
   for (uint32_t i = currentOffset; i < hash_table_length; ++i) {
     if (hash_table[i].name != NULL &&
         (handle->prefix_len == 0 ||
-         (memcmp(handle->prefix, hash_table[i].name, handle->prefix_len) == 0))) {
+         (hash_table[i].name_length >= handle->prefix_len &&
+          memcmp(handle->prefix, hash_table[i].name, handle->prefix_len) == 0)) &&
+        (handle->suffix_len == 0 ||
+         (hash_table[i].name_length >= handle->suffix_len &&
+          memcmp(handle->suffix,
+                 hash_table[i].name + hash_table[i].name_length - handle->suffix_len,
+                 handle->suffix_len) == 0))) {
       handle->position = (i + 1);
       const int error = FindEntry(archive, i, data);
       if (!error) {
@@ -1265,4 +1283,3 @@ const char* ErrorCodeString(int32_t error_code) {
 int GetFileDescriptor(const ZipArchiveHandle handle) {
   return reinterpret_cast<ZipArchive*>(handle)->fd;
 }
-
diff --git a/libziparchive/zip_archive_test.cc b/libziparchive/zip_archive_test.cc
index f8952ce..c799869 100644
--- a/libziparchive/zip_archive_test.cc
+++ b/libziparchive/zip_archive_test.cc
@@ -115,7 +115,7 @@ TEST(ziparchive, Iteration) {
   ASSERT_EQ(0, OpenArchiveWrapper(kValidZip, &handle));
 
   void* iteration_cookie;
-  ASSERT_EQ(0, StartIteration(handle, &iteration_cookie, NULL));
+  ASSERT_EQ(0, StartIteration(handle, &iteration_cookie, NULL, NULL));
 
   ZipEntry data;
   ZipEntryName name;
@@ -146,6 +146,116 @@ TEST(ziparchive, Iteration) {
   CloseArchive(handle);
 }
 
+TEST(ziparchive, IterationWithPrefix) {
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchiveWrapper(kValidZip, &handle));
+
+  void* iteration_cookie;
+  ZipEntryName prefix("b/");
+  ASSERT_EQ(0, StartIteration(handle, &iteration_cookie, &prefix, NULL));
+
+  ZipEntry data;
+  ZipEntryName name;
+
+  // b/c.txt
+  ASSERT_EQ(0, Next(iteration_cookie, &data, &name));
+  AssertNameEquals("b/c.txt", name);
+
+  // b/d.txt
+  ASSERT_EQ(0, Next(iteration_cookie, &data, &name));
+  AssertNameEquals("b/d.txt", name);
+
+  // b/
+  ASSERT_EQ(0, Next(iteration_cookie, &data, &name));
+  AssertNameEquals("b/", name);
+
+  // End of iteration.
+  ASSERT_EQ(-1, Next(iteration_cookie, &data, &name));
+
+  CloseArchive(handle);
+}
+
+TEST(ziparchive, IterationWithSuffix) {
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchiveWrapper(kValidZip, &handle));
+
+  void* iteration_cookie;
+  ZipEntryName suffix(".txt");
+  ASSERT_EQ(0, StartIteration(handle, &iteration_cookie, NULL, &suffix));
+
+  ZipEntry data;
+  ZipEntryName name;
+
+  // b/c.txt
+  ASSERT_EQ(0, Next(iteration_cookie, &data, &name));
+  AssertNameEquals("b/c.txt", name);
+
+  // b/d.txt
+  ASSERT_EQ(0, Next(iteration_cookie, &data, &name));
+  AssertNameEquals("b/d.txt", name);
+
+  // a.txt
+  ASSERT_EQ(0, Next(iteration_cookie, &data, &name));
+  AssertNameEquals("a.txt", name);
+
+  // b.txt
+  ASSERT_EQ(0, Next(iteration_cookie, &data, &name));
+  AssertNameEquals("b.txt", name);
+
+  // End of iteration.
+  ASSERT_EQ(-1, Next(iteration_cookie, &data, &name));
+
+  CloseArchive(handle);
+}
+
+TEST(ziparchive, IterationWithPrefixAndSuffix) {
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchiveWrapper(kValidZip, &handle));
+
+  void* iteration_cookie;
+  ZipEntryName prefix("b");
+  ZipEntryName suffix(".txt");
+  ASSERT_EQ(0, StartIteration(handle, &iteration_cookie, &prefix, &suffix));
+
+  ZipEntry data;
+  ZipEntryName name;
+
+  // b/c.txt
+  ASSERT_EQ(0, Next(iteration_cookie, &data, &name));
+  AssertNameEquals("b/c.txt", name);
+
+  // b/d.txt
+  ASSERT_EQ(0, Next(iteration_cookie, &data, &name));
+  AssertNameEquals("b/d.txt", name);
+
+  // b.txt
+  ASSERT_EQ(0, Next(iteration_cookie, &data, &name));
+  AssertNameEquals("b.txt", name);
+
+  // End of iteration.
+  ASSERT_EQ(-1, Next(iteration_cookie, &data, &name));
+
+  CloseArchive(handle);
+}
+
+TEST(ziparchive, IterationWithBadPrefixAndSuffix) {
+  ZipArchiveHandle handle;
+  ASSERT_EQ(0, OpenArchiveWrapper(kValidZip, &handle));
+
+  void* iteration_cookie;
+  ZipEntryName prefix("x");
+  ZipEntryName suffix("y");
+  ASSERT_EQ(0, StartIteration(handle, &iteration_cookie, &prefix, &suffix));
+
+  ZipEntry data;
+  ZipEntryName name;
+
+  // End of iteration.
+  ASSERT_EQ(-1, Next(iteration_cookie, &data, &name));
+
+  CloseArchive(handle);
+}
+
 TEST(ziparchive, FindEntry) {
   ZipArchiveHandle handle;
   ASSERT_EQ(0, OpenArchiveWrapper(kValidZip, &handle));