zipalign: Allow specifiying the target page size

Allow apps to specify the target page size for aligning their
uncompressed shared libraries.

This allows apps that want to support larger page sizes to do
so by specifiying the -P <pagesize_kb> flag.

However, apps built for 4k-only devices are unaffected as they
can continue to use -p flag for 4kB page alignment of uncompressed
shared libraries.

Bug: 276963821
Test: atest -c zipalign_tests
Change-Id: I890db067b8f898045f73e86788662f94a48af772
Signed-off-by: Kalesh Singh <kaleshsingh@google.com>
diff --git a/tools/zipalign/Android.bp b/tools/zipalign/Android.bp
index 0e1d58e..8be7e25 100644
--- a/tools/zipalign/Android.bp
+++ b/tools/zipalign/Android.bp
@@ -70,6 +70,7 @@
         "libgmock",
     ],
     data: [
+         "tests/data/apkWithUncompressedSharedLibs.zip",
          "tests/data/archiveWithOneDirectoryEntry.zip",
          "tests/data/diffOrders.zip",
          "tests/data/holes.zip",
diff --git a/tools/zipalign/ZipAlign.cpp b/tools/zipalign/ZipAlign.cpp
index 23840e3..f32f90b 100644
--- a/tools/zipalign/ZipAlign.cpp
+++ b/tools/zipalign/ZipAlign.cpp
@@ -17,6 +17,7 @@
 #include "ZipFile.h"
 
 #include <stdio.h>
+#include <string.h>
 #include <stdlib.h>
 #include <unistd.h>
 
@@ -36,17 +37,14 @@
 }
 
 static int getAlignment(bool pageAlignSharedLibs, int defaultAlignment,
-    ZipEntry* pEntry) {
-
-    static const int kPageAlignment = 4096;
-
+    ZipEntry* pEntry, int pageSize) {
     if (!pageAlignSharedLibs) {
         return defaultAlignment;
     }
 
     const char* ext = strrchr(pEntry->getFileName(), '.');
     if (ext && strcmp(ext, ".so") == 0) {
-        return kPageAlignment;
+        return pageSize;
     }
 
     return defaultAlignment;
@@ -56,7 +54,7 @@
  * Copy all entries from "pZin" to "pZout", aligning as needed.
  */
 static int copyAndAlign(ZipFile* pZin, ZipFile* pZout, int alignment, bool zopfli,
-    bool pageAlignSharedLibs)
+    bool pageAlignSharedLibs, int pageSize)
 {
     int numEntries = pZin->getNumEntries();
     ZipEntry* pEntry;
@@ -84,7 +82,8 @@
                 status = pZout->add(pZin, pEntry, padding, &pNewEntry);
             }
         } else {
-            const int alignTo = getAlignment(pageAlignSharedLibs, alignment, pEntry);
+            const int alignTo = getAlignment(pageAlignSharedLibs, alignment, pEntry,
+                                             pageSize);
 
             //printf("--- %s: orig at %ld(+%d) len=%ld, adding pad=%d\n",
             //    pEntry->getFileName(), (long) pEntry->getFileOffset(),
@@ -107,7 +106,7 @@
  * output file exists and "force" wasn't specified.
  */
 int process(const char* inFileName, const char* outFileName,
-    int alignment, bool force, bool zopfli, bool pageAlignSharedLibs)
+    int alignment, bool force, bool zopfli, bool pageAlignSharedLibs, int pageSize)
 {
     ZipFile zin, zout;
 
@@ -127,7 +126,7 @@
     }
 
     if (zin.open(inFileName, ZipFile::kOpenReadOnly) != OK) {
-        fprintf(stderr, "Unable to open '%s' as zip archive\n", inFileName);
+        fprintf(stderr, "Unable to open '%s' as zip archive: %s\n", inFileName, strerror(errno));
         return 1;
     }
     if (zout.open(outFileName,
@@ -138,7 +137,8 @@
         return 1;
     }
 
-    int result = copyAndAlign(&zin, &zout, alignment, zopfli, pageAlignSharedLibs);
+    int result = copyAndAlign(&zin, &zout, alignment, zopfli, pageAlignSharedLibs,
+                              pageSize);
     if (result != 0) {
         printf("zipalign: failed rewriting '%s' to '%s'\n",
             inFileName, outFileName);
@@ -150,7 +150,7 @@
  * Verify the alignment of a zip archive.
  */
 int verify(const char* fileName, int alignment, bool verbose,
-    bool pageAlignSharedLibs)
+    bool pageAlignSharedLibs, int pageSize)
 {
     ZipFile zipFile;
     bool foundBad = false;
@@ -181,7 +181,8 @@
             continue;
        } else {
             off_t offset = pEntry->getFileOffset();
-            const int alignTo = getAlignment(pageAlignSharedLibs, alignment, pEntry);
+            const int alignTo = getAlignment(pageAlignSharedLibs, alignment, pEntry,
+                                             pageSize);
             if ((offset % alignTo) != 0) {
                 if (verbose) {
                     printf("%8jd %s (BAD - %jd)\n",
diff --git a/tools/zipalign/ZipAlignMain.cpp b/tools/zipalign/ZipAlignMain.cpp
index 53fc8d4..2f24403 100644
--- a/tools/zipalign/ZipAlignMain.cpp
+++ b/tools/zipalign/ZipAlignMain.cpp
@@ -34,15 +34,18 @@
     fprintf(stderr, "Zip alignment utility\n");
     fprintf(stderr, "Copyright (C) 2009 The Android Open Source Project\n\n");
     fprintf(stderr,
-        "Usage: zipalign [-f] [-p] [-v] [-z] <align> infile.zip outfile.zip\n"
-        "       zipalign -c [-p] [-v] <align> infile.zip\n\n" );
+        "Usage: zipalign [-f] [-p] [-P <pagesize_kb>] [-v] [-z] <align> infile.zip outfile.zip\n"
+        "       zipalign -c [-p] [-P <pagesize_kb>] [-v] <align> infile.zip\n\n" );
     fprintf(stderr,
         "  <align>: alignment in bytes, e.g. '4' provides 32-bit alignment\n");
     fprintf(stderr, "  -c: check alignment only (does not modify file)\n");
     fprintf(stderr, "  -f: overwrite existing outfile.zip\n");
-    fprintf(stderr, "  -p: page-align uncompressed .so files\n");
+    fprintf(stderr, "  -p: 4kb page-align uncompressed .so files\n");
     fprintf(stderr, "  -v: verbose output\n");
     fprintf(stderr, "  -z: recompress using Zopfli\n");
+    fprintf(stderr, "  -P <pagesize_kb>: Align uncompressed .so files to the specified\n");
+    fprintf(stderr, "                    page size. Valid values for <pagesize_kb> are 4, 16\n");
+    fprintf(stderr, "                    and 64. '-P' cannot be used in combination with '-p'.\n");
 }
 
 
@@ -57,12 +60,16 @@
     bool verbose = false;
     bool zopfli = false;
     bool pageAlignSharedLibs = false;
+    int pageSize = 4096;
+    bool legacyPageAlignmentFlag = false;   // -p
+    bool pageAlignmentFlag = false;         // -P <pagesize_kb>
     int result = 1;
     int alignment;
     char* endp;
 
     int opt;
-    while ((opt = getopt(argc, argv, "fcpvz")) != -1) {
+
+    while ((opt = getopt(argc, argv, "fcpvzP:")) != -1) {
         switch (opt) {
         case 'c':
             check = true;
@@ -77,7 +84,29 @@
             zopfli = true;
             break;
         case 'p':
+            legacyPageAlignmentFlag = true;
             pageAlignSharedLibs = true;
+            pageSize = 4096;
+            break;
+        case 'P':
+            pageAlignmentFlag = true;
+            pageAlignSharedLibs = true;
+
+            if (!optarg) {
+                fprintf(stderr, "ERROR: -P requires an argument\n");
+                wantUsage = true;
+                goto bail;
+            }
+
+            pageSize = atoi(optarg);
+            if (pageSize != 4 && pageSize != 16 && pageSize != 64) {
+                fprintf(stderr, "ERROR: Invalid argument for -P: %s\n", optarg);
+                wantUsage = true;
+                goto bail;
+            }
+
+            pageSize *= 1024;  // Convert from kB to bytes.
+
             break;
         default:
             fprintf(stderr, "ERROR: unknown flag -%c\n", opt);
@@ -86,6 +115,13 @@
         }
     }
 
+    if (legacyPageAlignmentFlag && pageAlignmentFlag) {
+            fprintf(stderr, "ERROR: Invalid options: '-P <pagesize_kb>' and '-p'"
+                            "cannot be used in combination.\n");
+            wantUsage = true;
+            goto bail;
+    }
+
     if (!((check && (argc - optind) == 2) || (!check && (argc - optind) == 3))) {
         wantUsage = true;
         goto bail;
@@ -100,14 +136,15 @@
 
     if (check) {
         /* check existing archive for correct alignment */
-        result = verify(argv[optind + 1], alignment, verbose, pageAlignSharedLibs);
+        result = verify(argv[optind + 1], alignment, verbose, pageAlignSharedLibs, pageSize);
     } else {
         /* create the new archive */
-        result = process(argv[optind + 1], argv[optind + 2], alignment, force, zopfli, pageAlignSharedLibs);
+        result = process(argv[optind + 1], argv[optind + 2], alignment, force, zopfli,
+                         pageAlignSharedLibs, pageSize);
 
         /* trust, but verify */
         if (result == 0) {
-            result = verify(argv[optind + 2], alignment, verbose, pageAlignSharedLibs);
+            result = verify(argv[optind + 2], alignment, verbose, pageAlignSharedLibs, pageSize);
         }
     }
 
diff --git a/tools/zipalign/include/ZipAlign.h b/tools/zipalign/include/ZipAlign.h
index ab36086..85dda14 100644
--- a/tools/zipalign/include/ZipAlign.h
+++ b/tools/zipalign/include/ZipAlign.h
@@ -25,24 +25,28 @@
  * - force  : Overwrite output if it exists, fail otherwise.
  * - zopfli : Recompress compressed entries with more efficient algorithm.
  *            Copy compressed entries as-is, and unaligned, otherwise.
- * - pageAlignSharedLibs: Align .so files to 4096 and other files to
+ * - pageAlignSharedLibs: Align .so files to @pageSize and other files to
  *   alignTo, or all files to alignTo if false..
+ * - pageSize: Specifies the page size of the target device. This is used
+ *             to correctly page-align shared libraries.
  *
  * Returns 0 on success.
  */
 int process(const char* input, const char* output, int alignTo, bool force,
-    bool zopfli, bool pageAlignSharedLibs);
+    bool zopfli, bool pageAlignSharedLibs, int pageSize);
 
 /*
  * Verify the alignment of a zip archive.
  * - alignTo: Alignment (in bytes) for uncompressed entries.
- * - pageAlignSharedLibs: Align .so files to 4096 and other files to
+ * - pageAlignSharedLibs: Align .so files to @pageSize and other files to
  *   alignTo, or all files to alignTo if false..
+ * - pageSize: Specifies the page size of the target device. This is used
+ *             to correctly page-align shared libraries.
  *
  * Returns 0 on success.
  */
 int verify(const char* fileName, int alignTo, bool verbose,
-    bool pageAlignSharedLibs);
+    bool pageAlignSharedLibs, int pageSize);
 
 } // namespace android
 
diff --git a/tools/zipalign/tests/data/apkWithUncompressedSharedLibs.zip b/tools/zipalign/tests/data/apkWithUncompressedSharedLibs.zip
new file mode 100644
index 0000000..930e3b5
--- /dev/null
+++ b/tools/zipalign/tests/data/apkWithUncompressedSharedLibs.zip
Binary files differ
diff --git a/tools/zipalign/tests/src/align_test.cpp b/tools/zipalign/tests/src/align_test.cpp
index 09be445..07ad7cc 100644
--- a/tools/zipalign/tests/src/align_test.cpp
+++ b/tools/zipalign/tests/src/align_test.cpp
@@ -48,11 +48,12 @@
 TEST(Align, Unaligned) {
   const std::string src = GetTestPath("unaligned.zip");
   const std::string dst = GetTempPath("unaligned_out.zip");
+  int pageSize = 4096;
 
-  int processed = process(src.c_str(), dst.c_str(), 4, true, false, false);
+  int processed = process(src.c_str(), dst.c_str(), 4, true, false, false, pageSize);
   ASSERT_EQ(0, processed);
 
-  int verified = verify(dst.c_str(), 4, true, false);
+  int verified = verify(dst.c_str(), 4, true, false, pageSize);
   ASSERT_EQ(0, verified);
 }
 
@@ -60,18 +61,19 @@
   const std::string src = GetTestPath("unaligned.zip");
   const std::string tmp = GetTempPath("da_aligned.zip");
   const std::string dst = GetTempPath("da_d_aligner.zip");
+  int pageSize = 4096;
 
-  int processed = process(src.c_str(), tmp.c_str(), 4, true, false, false);
+  int processed = process(src.c_str(), tmp.c_str(), 4, true, false, false, pageSize);
   ASSERT_EQ(0, processed);
 
-  int verified = verify(tmp.c_str(), 4, true, false);
+  int verified = verify(tmp.c_str(), 4, true, false, pageSize);
   ASSERT_EQ(0, verified);
 
   // Align the result of the previous run. Essentially double aligning.
-  processed = process(tmp.c_str(), dst.c_str(), 4, true, false, false);
+  processed = process(tmp.c_str(), dst.c_str(), 4, true, false, false, pageSize);
   ASSERT_EQ(0, processed);
 
-  verified = verify(dst.c_str(), 4, true, false);
+  verified = verify(dst.c_str(), 4, true, false, pageSize);
   ASSERT_EQ(0, verified);
 
   // Nothing should have changed between tmp and dst.
@@ -90,11 +92,12 @@
 TEST(Align, Holes) {
   const std::string src = GetTestPath("holes.zip");
   const std::string dst = GetTempPath("holes_out.zip");
+  int pageSize = 4096;
 
-  int processed = process(src.c_str(), dst.c_str(), 4, true, false, true);
+  int processed = process(src.c_str(), dst.c_str(), 4, true, false, true, pageSize);
   ASSERT_EQ(0, processed);
 
-  int verified = verify(dst.c_str(), 4, false, true);
+  int verified = verify(dst.c_str(), 4, false, true, pageSize);
   ASSERT_EQ(0, verified);
 }
 
@@ -102,28 +105,85 @@
 TEST(Align, DifferenteOrders) {
   const std::string src = GetTestPath("diffOrders.zip");
   const std::string dst = GetTempPath("diffOrders_out.zip");
+  int pageSize = 4096;
 
-  int processed = process(src.c_str(), dst.c_str(), 4, true, false, true);
+  int processed = process(src.c_str(), dst.c_str(), 4, true, false, true, pageSize);
   ASSERT_EQ(0, processed);
 
-  int verified = verify(dst.c_str(), 4, false, true);
+  int verified = verify(dst.c_str(), 4, false, true, pageSize);
   ASSERT_EQ(0, verified);
 }
 
 TEST(Align, DirectoryEntryDoNotRequireAlignment) {
   const std::string src = GetTestPath("archiveWithOneDirectoryEntry.zip");
-  int verified = verify(src.c_str(), 4, false, true);
+  int pageSize = 4096;
+  int verified = verify(src.c_str(), 4, false, true, pageSize);
   ASSERT_EQ(0, verified);
 }
 
 TEST(Align, DirectoryEntry) {
   const std::string src = GetTestPath("archiveWithOneDirectoryEntry.zip");
   const std::string dst = GetTempPath("archiveWithOneDirectoryEntry_out.zip");
+  int pageSize = 4096;
 
-  int processed = process(src.c_str(), dst.c_str(), 4, true, false, true);
+  int processed = process(src.c_str(), dst.c_str(), 4, true, false, true, pageSize);
   ASSERT_EQ(0, processed);
   ASSERT_EQ(true, sameContent(src, dst));
 
-  int verified = verify(dst.c_str(), 4, false, true);
+  int verified = verify(dst.c_str(), 4, false, true, pageSize);
+  ASSERT_EQ(0, verified);
+}
+
+class UncompressedSharedLibsTest : public ::testing::Test {
+  protected:
+    static void SetUpTestSuite() {
+      src = GetTestPath("apkWithUncompressedSharedLibs.zip");
+      dst = GetTempPath("apkWithUncompressedSharedLibs_out.zip");
+    }
+
+    static std::string src;
+    static std::string dst;
+};
+
+std::string UncompressedSharedLibsTest::src;
+std::string UncompressedSharedLibsTest::dst;
+
+TEST_F(UncompressedSharedLibsTest, Unaligned) {
+  int pageSize = 4096;
+
+  int processed = process(src.c_str(), dst.c_str(), 4, true, false, false, pageSize);
+  ASSERT_EQ(0, processed);
+
+  int verified = verify(dst.c_str(), 4, true, true, pageSize);
+  ASSERT_NE(0, verified); // .so's not page-aligned
+}
+
+TEST_F(UncompressedSharedLibsTest, AlignedPageSize4kB) {
+  int pageSize = 4096;
+
+  int processed = process(src.c_str(), dst.c_str(), 4, true, false, true, pageSize);
+  ASSERT_EQ(0, processed);
+
+  int verified = verify(dst.c_str(), 4, true, true, pageSize);
+  ASSERT_EQ(0, verified);
+}
+
+TEST_F(UncompressedSharedLibsTest, AlignedPageSize16kB) {
+  int pageSize = 16384;
+
+  int processed = process(src.c_str(), dst.c_str(), 4, true, false, true, pageSize);
+  ASSERT_EQ(0, processed);
+
+  int verified = verify(dst.c_str(), 4, true, true, pageSize);
+  ASSERT_EQ(0, verified);
+}
+
+TEST_F(UncompressedSharedLibsTest, AlignedPageSize64kB) {
+  int pageSize = 65536;
+
+  int processed = process(src.c_str(), dst.c_str(), 4, true, false, true, pageSize);
+  ASSERT_EQ(0, processed);
+
+  int verified = verify(dst.c_str(), 4, true, true, pageSize);
   ASSERT_EQ(0, verified);
 }