Add "align file size" option to ApkSigner

The CLI tool apksigner can pass the option with "--align-file-size" flag
to ApkSigner.

This corresponds to the change made to signapk.
(https://android-review.googlesource.com/c/platform/build/+/1762580)

Bug: 194771859
Test: apksigner sign --key key --cert cert --in in.apk --out out.apk
   out.apk should be sized as multiples of 4K
Test: ./gradlew test
Change-Id: Ib5ec9865aa9d92a37fc00e4ac9d1358f2414b07b
diff --git a/src/apksigner/java/com/android/apksigner/ApkSignerTool.java b/src/apksigner/java/com/android/apksigner/ApkSignerTool.java
index 9fd0c34..d21515c 100644
--- a/src/apksigner/java/com/android/apksigner/ApkSignerTool.java
+++ b/src/apksigner/java/com/android/apksigner/ApkSignerTool.java
@@ -142,6 +142,7 @@
         boolean v3SigningEnabled = true;
         boolean v4SigningEnabled = true;
         boolean forceSourceStampOverwrite = false;
+        boolean alignFileSize = false;
         boolean verityEnabled = false;
         boolean debuggableApkPermitted = true;
         int minSdkVersion = 1;
@@ -186,6 +187,8 @@
                 v4SigningFlagFound = true;
             } else if ("force-stamp-overwrite".equals(optionName)) {
                 forceSourceStampOverwrite = optionsParser.getOptionalBooleanValue(true);
+            } else if ("align-file-size".equals(optionName)) {
+                alignFileSize = true;
             } else if ("verity-enabled".equals(optionName)) {
                 verityEnabled = optionsParser.getOptionalBooleanValue(true);
             } else if ("debuggable-apk-permitted".equals(optionName)) {
@@ -359,6 +362,7 @@
                         .setV3SigningEnabled(v3SigningEnabled)
                         .setV4SigningEnabled(v4SigningEnabled)
                         .setForceSourceStampOverwrite(forceSourceStampOverwrite)
+                        .setAlignFileSize(alignFileSize)
                         .setVerityEnabled(verityEnabled)
                         .setV4ErrorReportingEnabled(v4SigningEnabled && v4SigningFlagFound)
                         .setDebuggableApkPermitted(debuggableApkPermitted)
diff --git a/src/apksigner/java/com/android/apksigner/help_sign.txt b/src/apksigner/java/com/android/apksigner/help_sign.txt
index d66b7a3..5266ad9 100644
--- a/src/apksigner/java/com/android/apksigner/help_sign.txt
+++ b/src/apksigner/java/com/android/apksigner/help_sign.txt
@@ -50,6 +50,8 @@
                       APK, if found. By default, it is set to false. It has no
                       effect if no source stamp signer config is provided.
 
+--align-file-size     Produces APK file sized as multiples of 4K bytes.
+
 --verity-enabled      Whether to enable the verity signature algorithm for the
                       v2 and v3 signature schemes.
 
diff --git a/src/main/java/com/android/apksig/ApkSigner.java b/src/main/java/com/android/apksig/ApkSigner.java
index ca792c4..d6c7799 100644
--- a/src/main/java/com/android/apksig/ApkSigner.java
+++ b/src/main/java/com/android/apksig/ApkSigner.java
@@ -81,6 +81,8 @@
 
     private static final short ANDROID_COMMON_PAGE_ALIGNMENT_BYTES = 4096;
 
+    private static final short ANDROID_FILE_ALIGNMENT_BYTES = 4096;
+
     /** Name of the Android manifest ZIP entry in APKs. */
     private static final String ANDROID_MANIFEST_ZIP_ENTRY_NAME = "AndroidManifest.xml";
 
@@ -93,6 +95,7 @@
     private final boolean mV2SigningEnabled;
     private final boolean mV3SigningEnabled;
     private final boolean mV4SigningEnabled;
+    private final boolean mAlignFileSize;
     private final boolean mVerityEnabled;
     private final boolean mV4ErrorReportingEnabled;
     private final boolean mDebuggableApkPermitted;
@@ -122,6 +125,7 @@
             boolean v2SigningEnabled,
             boolean v3SigningEnabled,
             boolean v4SigningEnabled,
+            boolean alignFileSize,
             boolean verityEnabled,
             boolean v4ErrorReportingEnabled,
             boolean debuggableApkPermitted,
@@ -145,6 +149,7 @@
         mV2SigningEnabled = v2SigningEnabled;
         mV3SigningEnabled = v3SigningEnabled;
         mV4SigningEnabled = v4SigningEnabled;
+        mAlignFileSize = alignFileSize;
         mVerityEnabled = verityEnabled;
         mV4ErrorReportingEnabled = v4ErrorReportingEnabled;
         mDebuggableApkPermitted = debuggableApkPermitted;
@@ -579,6 +584,9 @@
         int outputCentralDirRecordCount = outputCdRecords.size();
 
         // Step 10. Construct output ZIP End of Central Directory record in an in-memory buffer
+        // because it can be adjusted in Step 11 due to signing block.
+        //   - CD offset (it's shifted by signing block)
+        //   - Comments (when the output file needs to be sized 4k-aligned)
         ByteBuffer outputEocd =
                 EocdRecord.createWithModifiedCentralDirectoryInfo(
                         inputZipSections.getZipEndOfCentralDirectory(),
@@ -597,13 +605,39 @@
 
         if (outputApkSigningBlockRequest != null) {
             int padding = outputApkSigningBlockRequest.getPaddingSizeBeforeApkSigningBlock();
-            outputApkOut.consume(ByteBuffer.allocate(padding));
             byte[] outputApkSigningBlock = outputApkSigningBlockRequest.getApkSigningBlock();
+            outputApkSigningBlockRequest.done();
+
+            long fileSize =
+                    outputCentralDirStartOffset
+                            + outputCentralDirDataSource.size()
+                            + padding
+                            + outputApkSigningBlock.length
+                            + outputEocd.remaining();
+            if (mAlignFileSize && (fileSize % ANDROID_FILE_ALIGNMENT_BYTES != 0)) {
+                int eocdPadding =
+                        (int)
+                                (ANDROID_FILE_ALIGNMENT_BYTES
+                                        - fileSize % ANDROID_FILE_ALIGNMENT_BYTES);
+                // Replace EOCD with padding one so that output file size can be the multiples of
+                // alignment.
+                outputEocd = EocdRecord.createWithPaddedComment(outputEocd, eocdPadding);
+
+                // Since EoCD has changed, we need to regenerate signing block as well.
+                outputApkSigningBlockRequest =
+                        signerEngine.outputZipSections2(
+                                outputApkIn,
+                                new ByteBufferDataSource(outputCentralDir),
+                                DataSources.asDataSource(outputEocd));
+                outputApkSigningBlock = outputApkSigningBlockRequest.getApkSigningBlock();
+                outputApkSigningBlockRequest.done();
+            }
+
+            outputApkOut.consume(ByteBuffer.allocate(padding));
             outputApkOut.consume(outputApkSigningBlock, 0, outputApkSigningBlock.length);
             ZipUtils.setZipEocdCentralDirectoryOffset(
                     outputEocd,
                     outputCentralDirStartOffset + padding + outputApkSigningBlock.length);
-            outputApkSigningBlockRequest.done();
         }
 
         // Step 12. Output ZIP Central Directory and ZIP End of Central Directory
@@ -1087,6 +1121,7 @@
         private boolean mV2SigningEnabled = true;
         private boolean mV3SigningEnabled = true;
         private boolean mV4SigningEnabled = true;
+        private boolean mAlignFileSize = false;
         private boolean mVerityEnabled = false;
         private boolean mV4ErrorReportingEnabled = false;
         private boolean mDebuggableApkPermitted = true;
@@ -1411,6 +1446,21 @@
             return this;
         }
 
+       /**
+         * Sets whether the output APK files should be sized as multiples of 4K.
+         *
+         * <p><em>Note:</em> This method may only be invoked when this builder is not initialized
+         * with an {@link ApkSignerEngine}.
+         *
+         * @throws IllegalStateException if this builder was initialized with an {@link
+         *     ApkSignerEngine}
+         */
+        public Builder setAlignFileSize(boolean alignFileSize) {
+            checkInitializedWithoutEngine();
+            mAlignFileSize = alignFileSize;
+            return this;
+        }
+
         /**
          * Sets whether to enable the verity signature algorithm for the v2 and v3 signature
          * schemes.
@@ -1542,6 +1592,7 @@
                     mV2SigningEnabled,
                     mV3SigningEnabled,
                     mV4SigningEnabled,
+                    mAlignFileSize,
                     mVerityEnabled,
                     mV4ErrorReportingEnabled,
                     mDebuggableApkPermitted,
diff --git a/src/main/java/com/android/apksig/internal/zip/EocdRecord.java b/src/main/java/com/android/apksig/internal/zip/EocdRecord.java
index 9c531f4..d2000b4 100644
--- a/src/main/java/com/android/apksig/internal/zip/EocdRecord.java
+++ b/src/main/java/com/android/apksig/internal/zip/EocdRecord.java
@@ -45,4 +45,13 @@
         ZipUtils.setUnsignedInt32(result, CD_OFFSET_OFFSET, centralDirectoryOffset);
         return result;
     }
+
+    public static ByteBuffer createWithPaddedComment(ByteBuffer original, int padding) {
+        ByteBuffer result = ByteBuffer.allocate((int) original.remaining() + padding);
+        result.order(ByteOrder.LITTLE_ENDIAN);
+        result.put(original.slice());
+        result.rewind();
+        ZipUtils.updateZipEocdCommentLen(result);
+        return result;
+    }
 }
diff --git a/src/test/java/com/android/apksig/ApkSignerTest.java b/src/test/java/com/android/apksig/ApkSignerTest.java
index d799201..69b8d4e 100644
--- a/src/test/java/com/android/apksig/ApkSignerTest.java
+++ b/src/test/java/com/android/apksig/ApkSignerTest.java
@@ -382,7 +382,11 @@
                         .setV2SigningEnabled(true)
                         .setV3SigningEnabled(true)
                         .setVerityEnabled(true));
-
+        signGolden(
+                "original.apk",
+                new File(outDir, "golden-file-size-aligned.apk"),
+                new ApkSigner.Builder(rsa2048SignerConfig)
+                        .setAlignFileSize(true));
         signGolden(
                 "pinsapp-unsigned.apk",
                 new File(outDir, "golden-pinsapp-signed.apk"),
@@ -695,6 +699,19 @@
     }
 
     @Test
+    public void testAlignFileSize_Golden() throws Exception {
+        List<ApkSigner.SignerConfig> rsaSignerConfig =
+                Collections.singletonList(
+                        getDefaultSignerConfigFromResources(FIRST_RSA_2048_SIGNER_RESOURCE_NAME));
+        String goldenOutput = "golden-file-size-aligned.apk";
+        assertGolden(
+                "original.apk",
+                goldenOutput,
+                new ApkSigner.Builder(rsaSignerConfig).setAlignFileSize(true));
+        assertTrue(Resources.toByteArray(getClass(), goldenOutput).length % 4096 == 0);
+    }
+
+    @Test
     public void testRsaSignedVerifies() throws Exception {
         List<ApkSigner.SignerConfig> signers =
                 Collections.singletonList(
diff --git a/src/test/resources/com/android/apksig/golden-file-size-aligned.apk b/src/test/resources/com/android/apksig/golden-file-size-aligned.apk
new file mode 100644
index 0000000..8dd95fc
--- /dev/null
+++ b/src/test/resources/com/android/apksig/golden-file-size-aligned.apk
Binary files differ