1 package com.android.apksig.internal.apk;
2 
3 import static org.junit.Assert.assertArrayEquals;
4 import static org.junit.Assert.assertEquals;
5 
6 import com.android.apksig.util.DataSource;
7 import com.android.apksig.util.DataSources;
8 import com.android.apksig.util.RunnablesExecutor;
9 import com.android.apksig.util.RunnablesProvider;
10 import java.io.File;
11 import java.io.FileOutputStream;
12 import java.io.RandomAccessFile;
13 import java.nio.ByteBuffer;
14 import java.util.ArrayList;
15 import java.util.EnumMap;
16 import java.util.EnumSet;
17 import java.util.List;
18 import java.util.Map;
19 import java.util.Set;
20 import java.util.concurrent.ExecutionException;
21 import java.util.concurrent.ForkJoinPool;
22 import java.util.concurrent.Future;
23 import org.junit.Before;
24 import org.junit.Rule;
25 import org.junit.Test;
26 import org.junit.rules.TemporaryFolder;
27 import org.junit.runner.RunWith;
28 import org.junit.runners.JUnit4;
29 
30 @RunWith(JUnit4.class)
31 public class ApkSigningBlockUtilsTest {
32     @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder();
33 
34     private static final int BASE = 255; // Intentionally not power of 2 to test properly
35 
36     DataSource[] dataSource;
37 
38     final Set<ContentDigestAlgorithm> algos = EnumSet.of(ContentDigestAlgorithm.CHUNKED_SHA512);
39 
40     @Before
setUp()41     public void setUp() throws Exception {
42         byte[] part1 = new byte[80 * 1024 * 1024 + 12345];
43         for (int i = 0; i < part1.length; ++i) {
44             part1[i] = (byte)(i % BASE);
45         }
46 
47         File dataFile = temporaryFolder.newFile("fake.apk");
48 
49         try (FileOutputStream fos = new FileOutputStream(dataFile)) {
50             fos.write(part1);
51         }
52         RandomAccessFile raf = new RandomAccessFile(dataFile, "r");
53 
54         byte[] part2 = new byte[1_500_000];
55         for (int i = 0; i < part2.length; ++i) {
56             part2[i] = (byte)(i % BASE);
57         }
58         byte[] part3 = new byte[30_000];
59         for (int i = 0; i < part3.length; ++i) {
60             part3[i] = (byte)(i % BASE);
61         }
62         dataSource = new DataSource[] {
63                 DataSources.asDataSource(raf),
64                 DataSources.asDataSource(ByteBuffer.wrap(part2)),
65                 DataSources.asDataSource(ByteBuffer.wrap(part3)),
66         };
67     }
68 
69     @Test
testNewVersionMatchesOld()70     public void testNewVersionMatchesOld() throws Exception {
71         Map<ContentDigestAlgorithm, byte[]> outputContentDigestsOld =
72                 new EnumMap<>(ContentDigestAlgorithm.class);
73         Map<ContentDigestAlgorithm, byte[]> outputContentDigestsNew =
74                 new EnumMap<>(ContentDigestAlgorithm.class);
75 
76         ApkSigningBlockUtils.computeOneMbChunkContentDigests(
77                 algos, dataSource, outputContentDigestsOld);
78 
79         ApkSigningBlockUtils.computeOneMbChunkContentDigests(
80                 RunnablesExecutor.SINGLE_THREADED,
81                 algos, dataSource, outputContentDigestsNew);
82 
83         assertEqualDigests(outputContentDigestsOld, outputContentDigestsNew);
84     }
85 
86     @Test
testMultithreadedVersionMatchesSinglethreaded()87     public void testMultithreadedVersionMatchesSinglethreaded() throws Exception {
88         Map<ContentDigestAlgorithm, byte[]> outputContentDigests =
89                 new EnumMap<>(ContentDigestAlgorithm.class);
90         Map<ContentDigestAlgorithm, byte[]> outputContentDigestsMultithreaded =
91                 new EnumMap<>(ContentDigestAlgorithm.class);
92 
93         ApkSigningBlockUtils.computeOneMbChunkContentDigests(
94                 RunnablesExecutor.SINGLE_THREADED,
95                 algos, dataSource, outputContentDigests);
96 
97         ApkSigningBlockUtils.computeOneMbChunkContentDigests(
98                 (RunnablesProvider provider) -> {
99                     ForkJoinPool forkJoinPool = ForkJoinPool.commonPool();
100                     int jobCount = forkJoinPool.getParallelism();
101                     List<Future<?>> jobs = new ArrayList<>(jobCount);
102 
103                     for (int i = 0; i < jobCount; i++) {
104                         jobs.add(forkJoinPool.submit(provider.createRunnable()));
105                     }
106 
107                     try {
108                         for (Future<?> future : jobs) {
109                             future.get();
110                         }
111                     } catch (InterruptedException e) {
112                         Thread.currentThread().interrupt();
113                         throw new RuntimeException(e);
114                     } catch (ExecutionException e) {
115                         throw new RuntimeException(e);
116                     }
117                 },
118                 algos, dataSource, outputContentDigestsMultithreaded);
119 
120         assertEqualDigests(outputContentDigestsMultithreaded, outputContentDigests);
121     }
122 
assertEqualDigests( Map<ContentDigestAlgorithm, byte[]> d1, Map<ContentDigestAlgorithm, byte[]> d2)123     private void assertEqualDigests(
124             Map<ContentDigestAlgorithm, byte[]> d1, Map<ContentDigestAlgorithm, byte[]> d2) {
125         assertEquals(d1.keySet(), d2.keySet());
126         for (ContentDigestAlgorithm algo : d1.keySet()) {
127             byte[] digest1 = d1.get(algo);
128             byte[] digest2 = d2.get(algo);
129             assertArrayEquals(digest1, digest2);
130         }
131     }
132 }
133