Remove explicit concurrency from apksig implementation am: b6b4a0a9a6 am: e43306cafcandroid-mainline-10.0.0_r9android-mainline-10.0.0_r7android-mainline-10.0.0_r5android-mainline-10.0.0_r4android-mainline-10.0.0_r10android-10.0.0_r9android-10.0.0_r8android-10.0.0_r7android-10.0.0_r45android-10.0.0_r44android-10.0.0_r43android-10.0.0_r42android-10.0.0_r41android-10.0.0_r40android-10.0.0_r39android-10.0.0_r38android-10.0.0_r37android-10.0.0_r36android-10.0.0_r35android-10.0.0_r34android-10.0.0_r33android-10.0.0_r32android-10.0.0_r31android-10.0.0_r30android-10.0.0_r29android-10.0.0_r28android-10.0.0_r27android-10.0.0_r26android-10.0.0_r25android-10.0.0_r24android-10.0.0_r23android-10.0.0_r22android-10.0.0_r21android-10.0.0_r20android-10.0.0_r19android-10.0.0_r18android-10.0.0_r16android-10.0.0_r15android-10.0.0_r14android-10.0.0_r13android-10.0.0_r12android10-qpr3-s1-releaseandroid10-qpr3-releaseandroid10-qpr2-s4-releaseandroid10-qpr2-s3-releaseandroid10-qpr2-s2-releaseandroid10-qpr2-s1-releaseandroid10-qpr2-releaseandroid10-qpr1-releaseandroid10-qpr1-mainline-releaseandroid10-qpr1-d-releaseandroid10-qpr1-c-s1-releaseandroid10-qpr1-c-releaseandroid10-qpr1-b-s1-releaseandroid10-qpr1-b-releaseandroid10-mainline-media-releaseandroid10-devandroid10-d4-s1-releaseandroid10-d4-releaseandroid10-c2f2-s2-releaseandroid10-c2f2-s1-releaseandroid10-c2f2-release

am: 73ad7f1098

Change-Id: I298c6be71608d63825c2c7ea9ce429a6d84bf729

11 files changed, 195 insertions, 83 deletions

diff --git a/src/main/java/com/android/apksig/ApkSignerEngine.java b/src/main/java/com/android/apksig/ApkSignerEngine.java
index 9f77eb1..138bc38 100644
--- a/src/main/java/com/android/apksig/ApkSignerEngine.java
+++ b/src/main/java/com/android/apksig/ApkSignerEngine.java
@@ -19,6 +19,7 @@ package com.android.apksig;
 import com.android.apksig.apk.ApkFormatException;
 import com.android.apksig.util.DataSink;
 import com.android.apksig.util.DataSource;
+import com.android.apksig.util.RunnablesExecutor;
 import java.io.Closeable;
 import java.io.IOException;
 import java.lang.UnsupportedOperationException;
@@ -116,6 +117,10 @@ import java.util.Set;
  */
 public interface ApkSignerEngine extends Closeable {
 
+    default void setExecutor(RunnablesExecutor executor) {
+        throw new UnsupportedOperationException("setExecutor method is not implemented");
+    }
+
     /**
      * Initializes the signer engine with the data already present in the apk (if any). There
      * might already be data that can be reused if the entries has not been changed.
diff --git a/src/main/java/com/android/apksig/ApkVerifier.java b/src/main/java/com/android/apksig/ApkVerifier.java
index 6e0a520..3e1e7da 100644
--- a/src/main/java/com/android/apksig/ApkVerifier.java
+++ b/src/main/java/com/android/apksig/ApkVerifier.java
@@ -29,6 +29,7 @@ import com.android.apksig.internal.util.AndroidSdkVersion;
 import com.android.apksig.internal.zip.CentralDirectoryRecord;
 import com.android.apksig.util.DataSource;
 import com.android.apksig.util.DataSources;
+import com.android.apksig.util.RunnablesExecutor;
 import com.android.apksig.zip.ZipFormatException;
 import java.io.Closeable;
 import java.io.File;
@@ -211,12 +212,13 @@ public class ApkVerifier {
         // verification. If the signature is found but does not verify, the APK is rejected.
         Set<Integer> foundApkSigSchemeIds = new HashSet<>(2);
         if (maxSdkVersion >= AndroidSdkVersion.N) {
-
+            RunnablesExecutor executor = RunnablesExecutor.SINGLE_THREADED;
             // Android P and newer attempts to verify APKs using APK Signature Scheme v3
             if (maxSdkVersion >= AndroidSdkVersion.P) {
                 try {
                     ApkSigningBlockUtils.Result v3Result =
                             V3SchemeVerifier.verify(
+                                    executor,
                                     apk,
                                     zipSections,
                                     Math.max(minSdkVersion, AndroidSdkVersion.P),
@@ -239,6 +241,7 @@ public class ApkVerifier {
                 try {
                     ApkSigningBlockUtils.Result v2Result =
                             V2SchemeVerifier.verify(
+                                    executor,
                                     apk,
                                     zipSections,
                                     supportedSchemeNames,
diff --git a/src/main/java/com/android/apksig/DefaultApkSignerEngine.java b/src/main/java/com/android/apksig/DefaultApkSignerEngine.java
index 778154e..c88239e 100644
--- a/src/main/java/com/android/apksig/DefaultApkSignerEngine.java
+++ b/src/main/java/com/android/apksig/DefaultApkSignerEngine.java
@@ -33,6 +33,7 @@ import com.android.apksig.util.DataSink;
 import com.android.apksig.util.DataSinks;
 import com.android.apksig.util.DataSource;
 
+import com.android.apksig.util.RunnablesExecutor;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -141,6 +142,9 @@ public class DefaultApkSignerEngine implements ApkSignerEngine {
      */
     private OutputApkSigningBlockRequestImpl mAddSigningBlockRequest;
 
+
+    private RunnablesExecutor mExecutor = RunnablesExecutor.SINGLE_THREADED;
+
     private DefaultApkSignerEngine(
             List<SignerConfig> signerConfigs,
             int minSdkVersion,
@@ -447,6 +451,11 @@ public class DefaultApkSignerEngine implements ApkSignerEngine {
     }
 
     @Override
+    public void setExecutor(RunnablesExecutor executor) {
+        mExecutor = executor;
+    }
+
+    @Override
     public void inputApkSigningBlock(DataSource apkSigningBlock) {
         checkNotClosed();
 
@@ -774,16 +783,25 @@ public class DefaultApkSignerEngine implements ApkSignerEngine {
             List<ApkSigningBlockUtils.SignerConfig> v2SignerConfigs =
                     createV2SignerConfigs(apkSigningBlockPaddingSupported);
             signingSchemeBlocks.add(
-                    V2SchemeSigner.generateApkSignatureSchemeV2Block(beforeCentralDir,
-                            zipCentralDirectory, eocd, v2SignerConfigs, mV3SigningEnabled));
+                    V2SchemeSigner.generateApkSignatureSchemeV2Block(
+                            mExecutor,
+                            beforeCentralDir,
+                            zipCentralDirectory,
+                            eocd,
+                            v2SignerConfigs,
+                            mV3SigningEnabled));
         }
         if (mV3SigningEnabled) {
             invalidateV3Signature();
             List<ApkSigningBlockUtils.SignerConfig> v3SignerConfigs =
                     createV3SignerConfigs(apkSigningBlockPaddingSupported);
             signingSchemeBlocks.add(
-                    V3SchemeSigner.generateApkSignatureSchemeV3Block(beforeCentralDir,
-                            zipCentralDirectory, eocd, v3SignerConfigs));
+                    V3SchemeSigner.generateApkSignatureSchemeV3Block(
+                            mExecutor,
+                            beforeCentralDir,
+                            zipCentralDirectory,
+                            eocd,
+                            v3SignerConfigs));
         }
 
         // create APK Signing Block with v2 and/or v3 blocks
diff --git a/src/main/java/com/android/apksig/internal/apk/ApkSigningBlockUtils.java b/src/main/java/com/android/apksig/internal/apk/ApkSigningBlockUtils.java
index 9444702..cc69af3 100644
--- a/src/main/java/com/android/apksig/internal/apk/ApkSigningBlockUtils.java
+++ b/src/main/java/com/android/apksig/internal/apk/ApkSigningBlockUtils.java
@@ -16,8 +16,8 @@
 
 package com.android.apksig.internal.apk;
 
-import com.android.apksig.SigningCertificateLineage;
 import com.android.apksig.ApkVerifier;
+import com.android.apksig.SigningCertificateLineage;
 import com.android.apksig.apk.ApkFormatException;
 import com.android.apksig.apk.ApkSigningBlockNotFoundException;
 import com.android.apksig.apk.ApkUtils;
@@ -31,6 +31,7 @@ import com.android.apksig.util.DataSinks;
 import com.android.apksig.util.DataSource;
 import com.android.apksig.util.DataSources;
 
+import com.android.apksig.util.RunnablesExecutor;
 import java.io.IOException;
 import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
@@ -153,6 +154,7 @@ public class ApkSigningBlockUtils {
      * exhibit the same behavior on all Android platform versions.
      */
     public static void verifyIntegrity(
+            RunnablesExecutor executor,
             DataSource beforeApkSigningBlock,
             DataSource centralDir,
             ByteBuffer eocd,
@@ -180,6 +182,7 @@ public class ApkSigningBlockUtils {
         try {
             actualContentDigests =
                     computeContentDigests(
+                            executor,
                             contentDigestAlgorithms,
                             beforeApkSigningBlock,
                             centralDir,
@@ -406,6 +409,7 @@ public class ApkSigningBlockUtils {
     }
 
     public static Map<ContentDigestAlgorithm, byte[]> computeContentDigests(
+            RunnablesExecutor executor,
             Set<ContentDigestAlgorithm> digestAlgorithms,
             DataSource beforeCentralDir,
             DataSource centralDir,
@@ -415,7 +419,8 @@ public class ApkSigningBlockUtils {
                 .filter(a -> a == ContentDigestAlgorithm.CHUNKED_SHA256 ||
                              a == ContentDigestAlgorithm.CHUNKED_SHA512)
                 .collect(Collectors.toSet());
-        computeOneMbChunkContentDigestsMultithread(
+        computeOneMbChunkContentDigests(
+                executor,
                 oneMbChunkBasedAlgorithm,
                 new DataSource[] { beforeCentralDir, centralDir, eocd },
                 contentDigests);
@@ -529,28 +534,12 @@ public class ApkSigningBlockUtils {
         }
     }
 
-    static void computeOneMbChunkContentDigestsMultithread(
+    static void computeOneMbChunkContentDigests(
+            RunnablesExecutor executor,
             Set<ContentDigestAlgorithm> digestAlgorithms,
             DataSource[] contents,
             Map<ContentDigestAlgorithm, byte[]> outputContentDigests)
             throws NoSuchAlgorithmException, DigestException {
-        ForkJoinPool forkJoinPool = ForkJoinPool.commonPool();
-        computeOneMbChunkContentDigestsMultithread(
-                digestAlgorithms,
-                contents,
-                outputContentDigests,
-                forkJoinPool::submit,
-                forkJoinPool.getParallelism());
-        forkJoinPool.shutdown();
-    }
-
-    private static void computeOneMbChunkContentDigestsMultithread(
-            Set<ContentDigestAlgorithm> digestAlgorithms,
-            DataSource[] contents,
-            Map<ContentDigestAlgorithm, byte[]> outputContentDigests,
-            Function<Runnable, Future<?>> jobRunner,
-            int jobCount)
-            throws NoSuchAlgorithmException, DigestException {
         long chunkCountLong = 0;
         for (DataSource input : contents) {
             chunkCountLong +=
@@ -566,23 +555,8 @@ public class ApkSigningBlockUtils {
             chunkDigestsList.add(new ChunkDigests(algorithms, chunkCount));
         }
 
-        List<Future<?>> jobs = new ArrayList<>(jobCount);
         ChunkSupplier chunkSupplier = new ChunkSupplier(contents);
-        for (int i = 0; i < jobCount; i++) {
-            jobs.add(jobRunner.apply(new ChunkDigester(chunkSupplier, chunkDigestsList)));
-        }
-
-        try {
-            for (Future<?> future : jobs) {
-                future.get();
-            }
-        }
-        catch (InterruptedException e) {
-            Thread.currentThread().interrupt();
-            throw new RuntimeException(e);
-        } catch (ExecutionException e) {
-            throw new RuntimeException(e);
-        }
+        executor.execute(() -> new ChunkDigester(chunkSupplier, chunkDigestsList));
 
         // Compute and write out final digest for each algorithm.
         for (ChunkDigests chunkDigests : chunkDigestsList) {
@@ -594,9 +568,9 @@ public class ApkSigningBlockUtils {
     }
 
     private static class ChunkDigests {
-        private ContentDigestAlgorithm algorithm;
-        private int digestOutputSize;
-        private byte[] concatOfDigestsOfChunks;
+        private final ContentDigestAlgorithm algorithm;
+        private final int digestOutputSize;
+        private final byte[] concatOfDigestsOfChunks;
 
         private ChunkDigests(ContentDigestAlgorithm algorithm, int chunkCount) {
             this.algorithm = algorithm;
@@ -627,13 +601,16 @@ public class ApkSigningBlockUtils {
         private final List<MessageDigest> messageDigests;
         private final DataSink mdSink;
 
-        private ChunkDigester(ChunkSupplier dataSupplier, List<ChunkDigests> chunkDigests)
-                throws NoSuchAlgorithmException {
+        private ChunkDigester(ChunkSupplier dataSupplier, List<ChunkDigests> chunkDigests) {
             this.dataSupplier = dataSupplier;
             this.chunkDigests = chunkDigests;
             messageDigests = new ArrayList<>(chunkDigests.size());
             for (ChunkDigests chunkDigest : chunkDigests) {
-                messageDigests.add(chunkDigest.createMessageDigest());
+                try {
+                    messageDigests.add(chunkDigest.createMessageDigest());
+                } catch (NoSuchAlgorithmException ex) {
+                    throw new RuntimeException(ex);
+                }
             }
             mdSink = DataSinks.asDataSink(messageDigests.toArray(new MessageDigest[0]));
         }
@@ -782,7 +759,7 @@ public class ApkSigningBlockUtils {
         outputContentDigests.put(ContentDigestAlgorithm.VERITY_CHUNKED_SHA256, encoded.array());
     }
 
-    private static final long getChunkCount(long inputSize, long chunkSize) {
+    private static long getChunkCount(long inputSize, long chunkSize) {
         return (inputSize + chunkSize - 1) / chunkSize;
     }
 
@@ -1032,6 +1009,7 @@ public class ApkSigningBlockUtils {
      */
     public static Pair<List<SignerConfig>, Map<ContentDigestAlgorithm, byte[]>>
             computeContentDigests(
+                    RunnablesExecutor executor,
                     DataSource beforeCentralDir,
                     DataSource centralDir,
                     DataSource eocd,
@@ -1055,6 +1033,7 @@ public class ApkSigningBlockUtils {
         try {
             contentDigests =
                     computeContentDigests(
+                            executor,
                             contentDigestAlgorithms,
                             beforeCentralDir,
                             centralDir,
diff --git a/src/main/java/com/android/apksig/internal/apk/v2/V2SchemeSigner.java b/src/main/java/com/android/apksig/internal/apk/v2/V2SchemeSigner.java
index e9f710b..d8e4723 100644
--- a/src/main/java/com/android/apksig/internal/apk/v2/V2SchemeSigner.java
+++ b/src/main/java/com/android/apksig/internal/apk/v2/V2SchemeSigner.java
@@ -27,7 +27,7 @@ import com.android.apksig.internal.apk.ContentDigestAlgorithm;
 import com.android.apksig.internal.apk.SignatureAlgorithm;
 import com.android.apksig.internal.util.Pair;
 import com.android.apksig.util.DataSource;
-
+import com.android.apksig.util.RunnablesExecutor;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -139,6 +139,7 @@ public abstract class V2SchemeSigner {
     }
 
     public static Pair<byte[], Integer> generateApkSignatureSchemeV2Block(
+            RunnablesExecutor executor,
             DataSource beforeCentralDir,
             DataSource centralDir,
             DataSource eocd,
@@ -148,8 +149,8 @@ public abstract class V2SchemeSigner {
                             SignatureException {
         Pair<List<SignerConfig>,
                 Map<ContentDigestAlgorithm, byte[]>> digestInfo =
-                ApkSigningBlockUtils.computeContentDigests(beforeCentralDir, centralDir, eocd,
-                        signerConfigs);
+                ApkSigningBlockUtils.computeContentDigests(
+                        executor, beforeCentralDir, centralDir, eocd, signerConfigs);
         return generateApkSignatureSchemeV2Block(
                 digestInfo.getFirst(), digestInfo.getSecond(),v3SigningEnabled);
     }
diff --git a/src/main/java/com/android/apksig/internal/apk/v2/V2SchemeVerifier.java b/src/main/java/com/android/apksig/internal/apk/v2/V2SchemeVerifier.java
index bbef027..51c40bd 100644
--- a/src/main/java/com/android/apksig/internal/apk/v2/V2SchemeVerifier.java
+++ b/src/main/java/com/android/apksig/internal/apk/v2/V2SchemeVerifier.java
@@ -27,8 +27,7 @@ import com.android.apksig.internal.util.ByteBufferUtils;
 import com.android.apksig.internal.util.X509CertificateUtils;
 import com.android.apksig.internal.util.GuaranteedEncodedFormX509Certificate;
 import com.android.apksig.util.DataSource;
-
-import java.io.ByteArrayInputStream;
+import com.android.apksig.util.RunnablesExecutor;
 import java.io.IOException;
 import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
@@ -89,6 +88,7 @@ public abstract class V2SchemeVerifier {
      * @throws IOException if an I/O error occurs when reading the APK
      */
     public static ApkSigningBlockUtils.Result verify(
+            RunnablesExecutor executor,
             DataSource apk,
             ApkUtils.ZipSections zipSections,
             Map<Integer, String> supportedApkSigSchemeNames,
@@ -110,7 +110,8 @@ public abstract class V2SchemeVerifier {
                         signatureInfo.eocdOffset - signatureInfo.centralDirOffset);
         ByteBuffer eocd = signatureInfo.eocd;
 
-        verify(beforeApkSigningBlock,
+        verify(executor,
+                beforeApkSigningBlock,
                 signatureInfo.signatureBlock,
                 centralDir,
                 eocd,
@@ -125,13 +126,14 @@ public abstract class V2SchemeVerifier {
     /**
      * Verifies the provided APK's v2 signatures and outputs the results into the provided
      * {@code result}. APK is considered verified only if there are no errors reported in the
-     * {@code result}. See {@link #verify(DataSource, ApkUtils.ZipSections, Map, Set, int, int)} for
-     * more information about the contract of this method.
+     * {@code result}. See {@link #verify(RunnablesExecutor, DataSource, ApkUtils.ZipSections, Map,
+     * Set, int, int)} for more information about the contract of this method.
      *
      * @param result result populated by this method with interesting information about the APK,
      *        such as information about signers, and verification errors and warnings.
      */
     private static void verify(
+            RunnablesExecutor executor,
             DataSource beforeApkSigningBlock,
             ByteBuffer apkSignatureSchemeV2Block,
             DataSource centralDir,
@@ -155,7 +157,7 @@ public abstract class V2SchemeVerifier {
             return;
         }
         ApkSigningBlockUtils.verifyIntegrity(
-                beforeApkSigningBlock, centralDir, eocd, contentDigestsToVerify, result);
+                executor, beforeApkSigningBlock, centralDir, eocd, contentDigestsToVerify, result);
         if (!result.containsErrors()) {
             result.verified = true;
         }
diff --git a/src/main/java/com/android/apksig/internal/apk/v3/V3SchemeSigner.java b/src/main/java/com/android/apksig/internal/apk/v3/V3SchemeSigner.java
index fc70a0a..722b304 100644
--- a/src/main/java/com/android/apksig/internal/apk/v3/V3SchemeSigner.java
+++ b/src/main/java/com/android/apksig/internal/apk/v3/V3SchemeSigner.java
@@ -29,7 +29,7 @@ import com.android.apksig.internal.apk.ContentDigestAlgorithm;
 import com.android.apksig.internal.apk.SignatureAlgorithm;
 import com.android.apksig.internal.util.Pair;
 import com.android.apksig.util.DataSource;
-
+import com.android.apksig.util.RunnablesExecutor;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -128,6 +128,7 @@ public abstract class V3SchemeSigner {
     }
 
     public static Pair<byte[], Integer> generateApkSignatureSchemeV3Block(
+            RunnablesExecutor executor,
             DataSource beforeCentralDir,
             DataSource centralDir,
             DataSource eocd,
@@ -136,8 +137,8 @@ public abstract class V3SchemeSigner {
                             SignatureException {
         Pair<List<SignerConfig>,
                 Map<ContentDigestAlgorithm, byte[]>> digestInfo =
-                ApkSigningBlockUtils.computeContentDigests(beforeCentralDir, centralDir, eocd,
-                        signerConfigs);
+                ApkSigningBlockUtils.computeContentDigests(
+                        executor, beforeCentralDir, centralDir, eocd, signerConfigs);
         return generateApkSignatureSchemeV3Block(digestInfo.getFirst(), digestInfo.getSecond());
     }
 
diff --git a/src/main/java/com/android/apksig/internal/apk/v3/V3SchemeVerifier.java b/src/main/java/com/android/apksig/internal/apk/v3/V3SchemeVerifier.java
index 9a2932b..16a6408 100644
--- a/src/main/java/com/android/apksig/internal/apk/v3/V3SchemeVerifier.java
+++ b/src/main/java/com/android/apksig/internal/apk/v3/V3SchemeVerifier.java
@@ -33,7 +33,7 @@ import com.android.apksig.internal.util.ByteBufferUtils;
 import com.android.apksig.internal.util.X509CertificateUtils;
 import com.android.apksig.internal.util.GuaranteedEncodedFormX509Certificate;
 import com.android.apksig.util.DataSource;
-
+import com.android.apksig.util.RunnablesExecutor;
 import java.io.IOException;
 import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
@@ -94,6 +94,7 @@ public abstract class V3SchemeVerifier {
      * @throws IOException if an I/O error occurs when reading the APK
      */
     public static ApkSigningBlockUtils.Result verify(
+            RunnablesExecutor executor,
             DataSource apk,
             ApkUtils.ZipSections zipSections,
             int minSdkVersion,
@@ -118,7 +119,8 @@ public abstract class V3SchemeVerifier {
             minSdkVersion = AndroidSdkVersion.P;
         }
 
-        verify(beforeApkSigningBlock,
+        verify(executor,
+                beforeApkSigningBlock,
                 signatureInfo.signatureBlock,
                 centralDir,
                 eocd,
@@ -131,13 +133,14 @@ public abstract class V3SchemeVerifier {
     /**
      * Verifies the provided APK's v3 signatures and outputs the results into the provided
      * {@code result}. APK is considered verified only if there are no errors reported in the
-     * {@code result}. See {@link #verify(DataSource, ApkUtils.ZipSections, int, int)} for more
-     * information about the contract of this method.
+     * {@code result}. See {@link #verify(RunnablesExecutor, DataSource, ApkUtils.ZipSections, int,
+     * int)} for more information about the contract of this method.
      *
      * @param result result populated by this method with interesting information about the APK,
      *        such as information about signers, and verification errors and warnings.
      */
     private static void verify(
+            RunnablesExecutor executor,
             DataSource beforeApkSigningBlock,
             ByteBuffer apkSignatureSchemeV3Block,
             DataSource centralDir,
@@ -153,7 +156,7 @@ public abstract class V3SchemeVerifier {
             return;
         }
         ApkSigningBlockUtils.verifyIntegrity(
-                beforeApkSigningBlock, centralDir, eocd, contentDigestsToVerify, result);
+                executor, beforeApkSigningBlock, centralDir, eocd, contentDigestsToVerify, result);
 
         // make sure that the v3 signers cover the entire targeted sdk version ranges and that the
         // longest SigningCertificateHistory, if present, corresponds to the newest platform
diff --git a/src/main/java/com/android/apksig/util/RunnablesExecutor.java b/src/main/java/com/android/apksig/util/RunnablesExecutor.java
new file mode 100644
index 0000000..04ec1d8
--- /dev/null
+++ b/src/main/java/com/android/apksig/util/RunnablesExecutor.java
@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) 2019 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.apksig.util;
+
+public interface RunnablesExecutor {
+    RunnablesExecutor SINGLE_THREADED = p -> p.getRunnable().run();
+
+    void execute(RunnablesProvider provider);
+}
diff --git a/src/main/java/com/android/apksig/util/RunnablesProvider.java b/src/main/java/com/android/apksig/util/RunnablesProvider.java
new file mode 100644
index 0000000..5b7bad2
--- /dev/null
+++ b/src/main/java/com/android/apksig/util/RunnablesProvider.java
@@ -0,0 +1,21 @@
+/*
+ * Copyright (C) 2019 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.apksig.util;
+
+public interface RunnablesProvider {
+    Runnable getRunnable();
+}
diff --git a/src/test/java/com/android/apksig/internal/apk/ApkSigningBlockUtilsTest.java b/src/test/java/com/android/apksig/internal/apk/ApkSigningBlockUtilsTest.java
index 77a8dab..7eb7c9b 100644
--- a/src/test/java/com/android/apksig/internal/apk/ApkSigningBlockUtilsTest.java
+++ b/src/test/java/com/android/apksig/internal/apk/ApkSigningBlockUtilsTest.java
@@ -5,15 +5,22 @@ import static org.junit.Assert.assertEquals;
 
 import com.android.apksig.util.DataSource;
 import com.android.apksig.util.DataSources;
+import com.android.apksig.util.RunnablesExecutor;
+import com.android.apksig.util.RunnablesProvider;
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.RandomAccessFile;
 import java.nio.ByteBuffer;
-import java.util.Arrays;
-import java.util.HashMap;
-import java.util.HashSet;
+import java.util.ArrayList;
+import java.util.EnumMap;
+import java.util.EnumSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.Future;
+import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
@@ -24,15 +31,14 @@ import org.junit.runners.JUnit4;
 public class ApkSigningBlockUtilsTest {
     @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder();
 
-    private static int BASE = 255; // Intentionally not power of 2 to test properly
+    private static final int BASE = 255; // Intentionally not power of 2 to test properly
 
-    @Test
-    public void testMultithreadVersionMatchesSinglethreaded() throws Exception {
-        Set<ContentDigestAlgorithm> algos = new HashSet<>(Arrays
-                .asList(ContentDigestAlgorithm.CHUNKED_SHA512));
-        Map<ContentDigestAlgorithm, byte[]> outputContentDigests = new HashMap<>();
-        Map<ContentDigestAlgorithm, byte[]> outputContentDigestsMultithread = new HashMap<>();
+    DataSource[] dataSource;
+
+    final Set<ContentDigestAlgorithm> algos = EnumSet.of(ContentDigestAlgorithm.CHUNKED_SHA512);
 
+    @Before
+    public void setUp() throws Exception {
         byte[] part1 = new byte[80 * 1024 * 1024 + 12345];
         for (int i = 0; i < part1.length; ++i) {
             part1[i] = (byte)(i % BASE);
@@ -53,23 +59,73 @@ public class ApkSigningBlockUtilsTest {
         for (int i = 0; i < part3.length; ++i) {
             part3[i] = (byte)(i % BASE);
         }
-
-        DataSource[] dataSource = {
+        dataSource = new DataSource[] {
                 DataSources.asDataSource(raf),
                 DataSources.asDataSource(ByteBuffer.wrap(part2)),
                 DataSources.asDataSource(ByteBuffer.wrap(part3)),
         };
+    }
+
+    @Test
+    public void testNewVersionMatchesOld() throws Exception {
+        Map<ContentDigestAlgorithm, byte[]> outputContentDigestsOld =
+                new EnumMap<>(ContentDigestAlgorithm.class);
+        Map<ContentDigestAlgorithm, byte[]> outputContentDigestsNew =
+                new EnumMap<>(ContentDigestAlgorithm.class);
+
+        ApkSigningBlockUtils.computeOneMbChunkContentDigests(
+                algos, dataSource, outputContentDigestsOld);
 
         ApkSigningBlockUtils.computeOneMbChunkContentDigests(
+                RunnablesExecutor.SINGLE_THREADED,
+                algos, dataSource, outputContentDigestsNew);
+
+        assertEqualDigests(outputContentDigestsOld, outputContentDigestsNew);
+    }
+
+    @Test
+    public void testMultithreadedVersionMatchesSinglethreaded() throws Exception {
+        Map<ContentDigestAlgorithm, byte[]> outputContentDigests =
+                new EnumMap<>(ContentDigestAlgorithm.class);
+        Map<ContentDigestAlgorithm, byte[]> outputContentDigestsMultithreaded =
+                new EnumMap<>(ContentDigestAlgorithm.class);
+
+        ApkSigningBlockUtils.computeOneMbChunkContentDigests(
+                RunnablesExecutor.SINGLE_THREADED,
                 algos, dataSource, outputContentDigests);
 
-        ApkSigningBlockUtils.computeOneMbChunkContentDigestsMultithread(
-                algos, dataSource, outputContentDigestsMultithread);
+        ApkSigningBlockUtils.computeOneMbChunkContentDigests(
+                (RunnablesProvider provider) -> {
+                    ForkJoinPool forkJoinPool = ForkJoinPool.commonPool();
+                    int jobCount = forkJoinPool.getParallelism();
+                    List<Future<?>> jobs = new ArrayList<>(jobCount);
+
+                    for (int i = 0; i < jobCount; i++) {
+                        jobs.add(forkJoinPool.submit(provider.getRunnable()));
+                    }
+
+                    try {
+                        for (Future<?> future : jobs) {
+                            future.get();
+                        }
+                    } catch (InterruptedException e) {
+                        Thread.currentThread().interrupt();
+                        throw new RuntimeException(e);
+                    } catch (ExecutionException e) {
+                        throw new RuntimeException(e);
+                    }
+                },
+                algos, dataSource, outputContentDigestsMultithreaded);
+
+        assertEqualDigests(outputContentDigestsMultithreaded, outputContentDigests);
+    }
 
-        assertEquals(outputContentDigestsMultithread.keySet(), outputContentDigests.keySet());
-        for (ContentDigestAlgorithm algo : outputContentDigests.keySet()) {
-            byte[] digest1 = outputContentDigestsMultithread.get(algo);
-            byte[] digest2 = outputContentDigests.get(algo);
+    private void assertEqualDigests(
+            Map<ContentDigestAlgorithm, byte[]> d1, Map<ContentDigestAlgorithm, byte[]> d2) {
+        assertEquals(d1.keySet(), d2.keySet());
+        for (ContentDigestAlgorithm algo : d1.keySet()) {
+            byte[] digest1 = d1.get(algo);
+            byte[] digest2 = d2.get(algo);
             assertArrayEquals(digest1, digest2);
         }
     }