Add zip hint generation support to signapk tool

Test: unzip -q -c myapp.apk.signed pinlist.meta | od --endian=big -w8 -tx4
Bug: 79259761
Bug: 65316207
Change-Id: I71c01ac24e93afe75f60697a9849e1dd35e1b49d

2 files changed, 132 insertions, 4 deletions

diff --git a/tools/signapk/src/com/android/signapk/CountingOutputStream.java b/tools/signapk/src/com/android/signapk/CountingOutputStream.java
new file mode 100644
index 0000000000..893a780440
--- /dev/null
+++ b/tools/signapk/src/com/android/signapk/CountingOutputStream.java
@@ -0,0 +1,59 @@
+/*
+ * Copyright (C) 2018 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.signapk;
+import java.io.OutputStream;
+import java.io.IOException;
+
+class CountingOutputStream extends OutputStream {
+    private final OutputStream mBase;
+    private long mWrittenBytes;
+
+    public CountingOutputStream(OutputStream base) {
+        mBase = base;
+    }
+
+    @Override
+    public void close() throws IOException {
+        mBase.close();
+    }
+
+    @Override
+    public void flush() throws IOException {
+        mBase.flush();
+    }
+
+    @Override
+    public void write(byte[] b) throws IOException {
+        mBase.write(b);
+        mWrittenBytes += b.length;
+    }
+
+    @Override
+    public void write(byte[] b, int off, int len) throws IOException {
+        mBase.write(b, off, len);
+        mWrittenBytes += len;
+    }
+
+    @Override
+    public void write(int b) throws IOException {
+        mBase.write(b);
+        mWrittenBytes += 1;
+    }
+
+    public long getWrittenBytes() {
+        return mWrittenBytes;
+    }
+}
diff --git a/tools/signapk/src/com/android/signapk/SignApk.java b/tools/signapk/src/com/android/signapk/SignApk.java
index fdf6283a03..57973ec044 100644
--- a/tools/signapk/src/com/android/signapk/SignApk.java
+++ b/tools/signapk/src/com/android/signapk/SignApk.java
@@ -36,6 +36,7 @@ import org.conscrypt.OpenSSLProvider;
 
 import com.android.apksig.ApkSignerEngine;
 import com.android.apksig.DefaultApkSignerEngine;
+import com.android.apksig.Hints;
 import com.android.apksig.apk.ApkUtils;
 import com.android.apksig.apk.MinSdkVersionException;
 import com.android.apksig.util.DataSink;
@@ -73,6 +74,7 @@ import java.security.spec.PKCS8EncodedKeySpec;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Enumeration;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.TimeZone;
@@ -80,6 +82,7 @@ import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
 import java.util.jar.JarOutputStream;
 import java.util.regex.Pattern;
+import java.util.zip.ZipEntry;
 
 import javax.crypto.Cipher;
 import javax.crypto.EncryptedPrivateKeyInfo;
@@ -372,11 +375,16 @@ class SignApk {
             Pattern ignoredFilenamePattern,
             ApkSignerEngine apkSigner,
             JarOutputStream out,
+            CountingOutputStream outCounter,
             long timestamp,
             int defaultAlignment) throws IOException {
         byte[] buffer = new byte[4096];
         int num;
 
+        List<Pattern> pinPatterns = extractPinPatterns(in);
+        ArrayList<Hints.ByteRange> pinByteRanges = pinPatterns == null ? null : new ArrayList<>();
+        HashSet<String> namesToPin = new HashSet<>();
+
         ArrayList<String> names = new ArrayList<String>();
         for (Enumeration<JarEntry> e = in.entries(); e.hasMoreElements();) {
             JarEntry entry = e.nextElement();
@@ -388,6 +396,16 @@ class SignApk {
                     && (ignoredFilenamePattern.matcher(entryName).matches())) {
                 continue;
             }
+            if (Hints.PIN_BYTE_RANGE_ZIP_ENTRY_NAME.equals(entryName)) {
+                continue;  // We regenerate it below.
+            }
+            if (pinPatterns != null) {
+                for (Pattern pinPattern : pinPatterns) {
+                    if (pinPattern.matcher(entryName).matches()) {
+                        namesToPin.add(entryName);
+                    }
+                }
+            }
             names.add(entryName);
         }
         Collections.sort(names);
@@ -460,6 +478,7 @@ class SignApk {
             outEntry.setExtra(extra);
             offset += extra.length;
 
+            long entryHeaderStart = outCounter.getWrittenBytes();
             out.putNextEntry(outEntry);
             ApkSignerEngine.InspectJarEntryRequest inspectEntryRequest =
                     (apkSigner != null) ? apkSigner.outputJarEntry(name) : null;
@@ -475,10 +494,18 @@ class SignApk {
                     offset += num;
                 }
             }
+            out.closeEntry();
             out.flush();
             if (inspectEntryRequest != null) {
                 inspectEntryRequest.done();
             }
+
+            if (namesToPin.contains(name)) {
+                pinByteRanges.add(
+                    new Hints.ByteRange(
+                        entryHeaderStart,
+                        outCounter.getWrittenBytes()));
+            }
         }
 
         // Copy all the non-STORED entries.  We don't attempt to
@@ -494,6 +521,7 @@ class SignApk {
             // Create a new entry so that the compressed len is recomputed.
             JarEntry outEntry = new JarEntry(name);
             outEntry.setTime(timestamp);
+            long entryHeaderStart = outCounter.getWrittenBytes();
             out.putNextEntry(outEntry);
             ApkSignerEngine.InspectJarEntryRequest inspectEntryRequest =
                     (apkSigner != null) ? apkSigner.outputJarEntry(name) : null;
@@ -507,11 +535,47 @@ class SignApk {
                     entryDataSink.consume(buffer, 0, num);
                 }
             }
+            out.closeEntry();
             out.flush();
             if (inspectEntryRequest != null) {
                 inspectEntryRequest.done();
             }
+
+            if (namesToPin.contains(name)) {
+                pinByteRanges.add(
+                    new Hints.ByteRange(
+                        entryHeaderStart,
+                        outCounter.getWrittenBytes()));
+            }
         }
+
+        if (pinByteRanges != null) {
+            // Cover central directory
+            pinByteRanges.add(
+                new Hints.ByteRange(outCounter.getWrittenBytes(),
+                                    Long.MAX_VALUE));
+            addPinByteRanges(out, pinByteRanges, timestamp);
+        }
+    }
+
+    private static List<Pattern> extractPinPatterns(JarFile in) throws IOException {
+        ZipEntry pinMetaEntry = in.getEntry(Hints.PIN_HINT_ASSET_ZIP_ENTRY_NAME);
+        if (pinMetaEntry == null) {
+            return null;
+        }
+        InputStream pinMetaStream = in.getInputStream(pinMetaEntry);
+        byte[] patternBlob = new byte[(int) pinMetaEntry.getSize()];
+        pinMetaStream.read(patternBlob);
+        return Hints.parsePinPatterns(patternBlob);
+    }
+
+    private static void addPinByteRanges(JarOutputStream outputJar,
+                                         ArrayList<Hints.ByteRange> pinByteRanges,
+                                         long timestamp) throws IOException {
+        JarEntry je = new JarEntry(Hints.PIN_BYTE_RANGE_ZIP_ENTRY_NAME);
+        je.setTime(timestamp);
+        outputJar.putNextEntry(je);
+        outputJar.write(Hints.encodeByteRangeList(pinByteRanges));
     }
 
     private static boolean shouldOutputApkEntry(
@@ -679,9 +743,11 @@ class SignApk {
         public void write(OutputStream out) throws IOException {
             try {
                 signer = new WholeFileSignerOutputStream(out, outputStream);
-                JarOutputStream outputJar = new JarOutputStream(signer);
+                CountingOutputStream outputJarCounter = new CountingOutputStream(signer);
+                JarOutputStream outputJar = new JarOutputStream(outputJarCounter);
 
-                copyFiles(inputJar, STRIP_PATTERN, null, outputJar, timestamp, 0);
+                copyFiles(inputJar, STRIP_PATTERN, null, outputJar,
+                          outputJarCounter, timestamp, 0);
                 addOtacert(outputJar, publicKeyFile, timestamp);
 
                 signer.notifyClosing();
@@ -1065,11 +1131,14 @@ class SignApk {
                     // Build the output APK in memory, by copying input APK's ZIP entries across
                     // and then signing the output APK.
                     ByteArrayOutputStream v1SignedApkBuf = new ByteArrayOutputStream();
-                    JarOutputStream outputJar = new JarOutputStream(v1SignedApkBuf);
+                    CountingOutputStream outputJarCounter =
+                            new CountingOutputStream(v1SignedApkBuf);
+                    JarOutputStream outputJar = new JarOutputStream(outputJarCounter);
                     // Use maximum compression for compressed entries because the APK lives forever
                     // on the system partition.
                     outputJar.setLevel(9);
-                    copyFiles(inputJar, null, apkSigner, outputJar, timestamp, alignment);
+                    copyFiles(inputJar, null, apkSigner, outputJar,
+                              outputJarCounter, timestamp, alignment);
                     ApkSignerEngine.OutputJarSignatureRequest addV1SignatureRequest =
                             apkSigner.outputJarEntries();
                     if (addV1SignatureRequest != null) {