Fix alignment when merging zips.

When merging a zip into an apk, apkzlib used to copy all
the non-ignored files from the zip as-is. This was done to
improve performance as it avoid recompressing files that
were already compressed.

However, if the files need to be uncompressed (for example,
native libraries in M+ devices), this would not ensure that
the libraries were uncompressed: if the libraries were
compressed in the original zip file, they were copied
as-is, i.e., compressed, to the apk breaking alignment.

This CL fixes this by ensuring that files that *need* to
be uncompressed (and only those) are actually added to
the apk following compression and alignment rules.

Test: included
Bug: http://b/37926537
Change-Id: I69848b37ef33b4f0af07dde8c778c7823443d55b

2 files changed, 252 insertions, 4 deletions

diff --git a/src/main/java/com/android/apkzlib/zfile/ApkZFileCreator.java b/src/main/java/com/android/apkzlib/zfile/ApkZFileCreator.java
index e56f99b..c85ad44 100644
--- a/src/main/java/com/android/apkzlib/zfile/ApkZFileCreator.java
+++ b/src/main/java/com/android/apkzlib/zfile/ApkZFileCreator.java
@@ -26,6 +26,7 @@ import com.google.common.io.Closer;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import javax.annotation.Nonnull;
@@ -114,14 +115,30 @@ class ApkZFileCreator implements ApkCreator {
         try {
             ZFile toMerge = closer.register(new ZFile(zip));
 
-            Predicate<String> predicate;
+            Predicate<String> ignorePredicate;
             if (isIgnored == null) {
-                predicate = s -> false;
+                ignorePredicate = s -> false;
             } else {
-                predicate = isIgnored;
+                ignorePredicate = isIgnored;
             }
 
-            this.zip.mergeFrom(toMerge, predicate);
+            // Files that *must* be uncompressed in the result should not be merged and should be
+            // added after. This is just very slightly less efficient than ignoring just the ones
+            // that were compressed and must be uncompressed, but it is a lot simpler :)
+            Predicate<String> noMergePredicate = ignorePredicate.or(noCompressPredicate);
+
+            this.zip.mergeFrom(toMerge, noMergePredicate);
+
+            for (StoredEntry toMergeEntry : toMerge.entries()) {
+                String path = toMergeEntry.getCentralDirectoryHeader().getName();
+                if (noCompressPredicate.test(path) && !ignorePredicate.test(path)) {
+                    // This entry *must* be uncompressed so it was ignored in the merge and should
+                    // now be added to the apk.
+                    try (InputStream ignoredData = toMergeEntry.open()) {
+                        this.zip.add(path, ignoredData, false);
+                    }
+                }
+            }
         } catch (Throwable t) {
             throw closer.rethrow(t);
         } finally {
diff --git a/src/test/java/com/android/apkzlib/zfile/ApkAlignmentTest.java b/src/test/java/com/android/apkzlib/zfile/ApkAlignmentTest.java
new file mode 100644
index 0000000..1731ba9
--- /dev/null
+++ b/src/test/java/com/android/apkzlib/zfile/ApkAlignmentTest.java
@@ -0,0 +1,231 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.apkzlib.zfile;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import com.android.apkzlib.zip.CompressionMethod;
+import com.android.apkzlib.zip.StoredEntry;
+import com.android.apkzlib.zip.ZFile;
+import com.android.apkzlib.zip.ZFileOptions;
+import java.io.ByteArrayInputStream;
+import java.io.File;
+import java.nio.file.Files;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+public class ApkAlignmentTest {
+    @Rule public TemporaryFolder mTemporaryFolder = new TemporaryFolder();
+
+    @Test
+    public void soFilesUncompressedAndAligned() throws Exception {
+        File apk = new File(mTemporaryFolder.getRoot(), "a.apk");
+
+        File soFile = new File(mTemporaryFolder.getRoot(), "doesnt_work.so");
+        Files.write(soFile.toPath(), new byte[500]);
+
+        ApkZFileCreatorFactory cf = new ApkZFileCreatorFactory(new ZFileOptions());
+        ApkCreatorFactory.CreationData creationData =
+                new ApkCreatorFactory.CreationData(
+                        apk,
+                        null,
+                        null,
+                        false,
+                        false,
+                        null,
+                        null,
+                        20,
+                        NativeLibrariesPackagingMode.UNCOMPRESSED_AND_ALIGNED,
+                        path -> false);
+
+        ApkCreator creator = cf.make(creationData);
+
+        creator.writeFile(soFile, "/doesnt_work.so");
+        creator.close();
+
+        try (ZFile zf = new ZFile(apk)) {
+            StoredEntry soEntry = zf.get("/doesnt_work.so");
+            assertNotNull(soEntry);
+            assertEquals(
+                    CompressionMethod.STORE,
+                    soEntry.getCentralDirectoryHeader().getCompressionInfoWithWait().getMethod());
+            long offset =
+                    soEntry.getCentralDirectoryHeader().getOffset() + soEntry.getLocalHeaderSize();
+            assertTrue(offset % 4096 == 0);
+        }
+    }
+
+    @Test
+    public void soFilesMergedFromZipsCanBeUncompressedAndAligned() throws Exception {
+
+        // Create a zip file with a compressed, unaligned so file.
+        File zipToMerge = new File(mTemporaryFolder.getRoot(), "a.zip");
+        try (ZFile zf = new ZFile(zipToMerge)) {
+            zf.add("/zero.so", new ByteArrayInputStream(new byte[500]));
+        }
+
+        try (ZFile zf = new ZFile(zipToMerge)) {
+            StoredEntry zeroSo = zf.get("/zero.so");
+            assertNotNull(zeroSo);
+            assertEquals(
+                    CompressionMethod.DEFLATE,
+                    zeroSo.getCentralDirectoryHeader().getCompressionInfoWithWait().getMethod());
+            long offset =
+                    zeroSo.getCentralDirectoryHeader().getOffset() + zeroSo.getLocalHeaderSize();
+            assertFalse(offset % 4096 == 0);
+        }
+
+        // Create an APK and merge the zip file.
+        File apk = new File(mTemporaryFolder.getRoot(), "b.apk");
+        ApkZFileCreatorFactory cf = new ApkZFileCreatorFactory(new ZFileOptions());
+        ApkCreatorFactory.CreationData creationData =
+                new ApkCreatorFactory.CreationData(
+                        apk,
+                        null,
+                        null,
+                        false,
+                        false,
+                        null,
+                        null,
+                        20,
+                        NativeLibrariesPackagingMode.UNCOMPRESSED_AND_ALIGNED,
+                        path -> false);
+
+        try (ApkCreator creator = cf.make(creationData)) {
+            creator.writeZip(zipToMerge, null, null);
+        }
+
+        // Make sure the file is uncompressed and aligned.
+        try (ZFile zf = new ZFile(apk)) {
+            StoredEntry soEntry = zf.get("/zero.so");
+            assertNotNull(soEntry);
+            assertEquals(
+                    CompressionMethod.STORE,
+                    soEntry.getCentralDirectoryHeader().getCompressionInfoWithWait().getMethod());
+            long offset =
+                    soEntry.getCentralDirectoryHeader().getOffset() + soEntry.getLocalHeaderSize();
+            assertTrue(offset % 4096 == 0);
+
+            byte[] data = soEntry.read();
+            assertEquals(500, data.length);
+            for (int i = 0; i < data.length; i++) {
+                assertEquals(0, data[i]);
+            }
+        }
+    }
+
+    @Test
+    public void soFilesUncompressedAndNotAligned() throws Exception {
+        File apk = new File(mTemporaryFolder.getRoot(), "a.apk");
+
+        File soFile = new File(mTemporaryFolder.getRoot(), "doesnt_work.so");
+        Files.write(soFile.toPath(), new byte[500]);
+
+        ApkZFileCreatorFactory cf = new ApkZFileCreatorFactory(new ZFileOptions());
+        ApkCreatorFactory.CreationData creationData =
+                new ApkCreatorFactory.CreationData(
+                        apk,
+                        null,
+                        null,
+                        false,
+                        false,
+                        null,
+                        null,
+                        20,
+                        NativeLibrariesPackagingMode.COMPRESSED,
+                        path -> false);
+
+        ApkCreator creator = cf.make(creationData);
+
+        creator.writeFile(soFile, "/doesnt_work.so");
+        creator.close();
+
+        try (ZFile zf = new ZFile(apk)) {
+            StoredEntry soEntry = zf.get("/doesnt_work.so");
+            assertNotNull(soEntry);
+            assertEquals(
+                    CompressionMethod.DEFLATE,
+                    soEntry.getCentralDirectoryHeader().getCompressionInfoWithWait().getMethod());
+            long offset =
+                    soEntry.getCentralDirectoryHeader().getOffset() + soEntry.getLocalHeaderSize();
+            assertTrue(offset % 4096 != 0);
+        }
+    }
+
+    @Test
+    public void soFilesMergedFromZipsCanBeUncompressedAndNotAligned() throws Exception {
+
+        // Create a zip file with a compressed, unaligned so file.
+        File zipToMerge = new File(mTemporaryFolder.getRoot(), "a.zip");
+        try (ZFile zf = new ZFile(zipToMerge)) {
+            zf.add("/zero.so", new ByteArrayInputStream(new byte[500]));
+        }
+
+        try (ZFile zf = new ZFile(zipToMerge)) {
+            StoredEntry zeroSo = zf.get("/zero.so");
+            assertNotNull(zeroSo);
+            assertEquals(
+                    CompressionMethod.DEFLATE,
+                    zeroSo.getCentralDirectoryHeader().getCompressionInfoWithWait().getMethod());
+            long offset =
+                    zeroSo.getCentralDirectoryHeader().getOffset() + zeroSo.getLocalHeaderSize();
+            assertFalse(offset % 4096 == 0);
+        }
+
+        // Create an APK and merge the zip file.
+        File apk = new File(mTemporaryFolder.getRoot(), "b.apk");
+        ApkZFileCreatorFactory cf = new ApkZFileCreatorFactory(new ZFileOptions());
+        ApkCreatorFactory.CreationData creationData =
+                new ApkCreatorFactory.CreationData(
+                        apk,
+                        null,
+                        null,
+                        false,
+                        false,
+                        null,
+                        null,
+                        20,
+                        NativeLibrariesPackagingMode.COMPRESSED,
+                        path -> false);
+
+        try (ApkCreator creator = cf.make(creationData)) {
+            creator.writeZip(zipToMerge, null, null);
+        }
+
+        // Make sure the file is uncompressed and aligned.
+        try (ZFile zf = new ZFile(apk)) {
+            StoredEntry soEntry = zf.get("/zero.so");
+            assertNotNull(soEntry);
+            assertEquals(
+                    CompressionMethod.DEFLATE,
+                    soEntry.getCentralDirectoryHeader().getCompressionInfoWithWait().getMethod());
+            long offset =
+                    soEntry.getCentralDirectoryHeader().getOffset() + soEntry.getLocalHeaderSize();
+            assertTrue(offset % 4096 != 0);
+
+            byte[] data = soEntry.read();
+            assertEquals(500, data.length);
+            for (int i = 0; i < data.length; i++) {
+                assertEquals(0, data[i]);
+            }
+        }
+    }
+}