Support regular fs-verity Merkle tree

In addition to the legacy "apk-verity" Merkle tree that skips the
Signing Block, now the builder supports the regular Merkle tree.

The class and methods are made public because it will be used in
different package.

Test: atest PtsApkVerityTestCases
Test: With other changes, observe the hash matches what other tool
      generates.
Bug: 112037636
Change-Id: I3d96932cd16f4f330170547d6a7689d2997e4b85

1 files changed, 91 insertions, 25 deletions

diff --git a/core/java/android/util/apk/ApkVerityBuilder.java b/core/java/android/util/apk/ApkVerityBuilder.java
index 553511d73670..edd09f8f73c4 100644
--- a/core/java/android/util/apk/ApkVerityBuilder.java
+++ b/core/java/android/util/apk/ApkVerityBuilder.java
@@ -40,7 +40,7 @@ import java.util.ArrayList;
  *
  * @hide
  */
-abstract class ApkVerityBuilder {
+public abstract class ApkVerityBuilder {
     private ApkVerityBuilder() {}
 
     private static final int CHUNK_SIZE_BYTES = 4096;  // Typical Linux block size
@@ -51,12 +51,18 @@ abstract class ApkVerityBuilder {
     private static final String JCA_DIGEST_ALGORITHM = "SHA-256";
     private static final byte[] DEFAULT_SALT = new byte[8];
 
-    static class ApkVerityResult {
+    /** Result generated by the builder. */
+    public static class ApkVerityResult {
+        /** Raw fs-verity metadata and Merkle tree ready to be deployed on disk. */
         public final ByteBuffer verityData;
+
+        /** Size of the Merkle tree in {@code verityData}. */
         public final int merkleTreeSize;
+
+        /** Root hash of the Merkle tree. */
         public final byte[] rootHash;
 
-        ApkVerityResult(ByteBuffer verityData, int merkleTreeSize, byte[] rootHash) {
+        private ApkVerityResult(ByteBuffer verityData, int merkleTreeSize, byte[] rootHash) {
             this.verityData = verityData;
             this.merkleTreeSize = merkleTreeSize;
             this.rootHash = rootHash;
@@ -65,19 +71,47 @@ abstract class ApkVerityBuilder {
 
     /**
      * Generates the 4k, SHA-256 based Merkle tree for the given APK and stores in the {@link
-     * ByteBuffer} created by the {@link ByteBufferFactory}.  The Merkle tree is suitable to be used
-     * as the on-disk format for apk-verity.
+     * ByteBuffer} created by the {@link ByteBufferFactory}.  The output is suitable to be used as
+     * the on-disk format for fs-verity to use.
      *
      * @return ApkVerityResult containing a buffer with the generated Merkle tree stored at the
      *         front, the tree size, and the calculated root hash.
      */
     @NonNull
-    static ApkVerityResult generateApkVerityTree(@NonNull RandomAccessFile apk,
+    public static ApkVerityResult generateFsVerityTree(@NonNull RandomAccessFile apk,
+            @NonNull ByteBufferFactory bufferFactory)
+            throws IOException, SecurityException, NoSuchAlgorithmException, DigestException {
+        return generateVerityTree(apk, bufferFactory, null /* signatureInfo */,
+                false /* skipSigningBlock */);
+    }
+
+    /**
+     * Generates the 4k, SHA-256 based Merkle tree for the given APK and stores in the {@link
+     * ByteBuffer} created by the {@link ByteBufferFactory}.  The Merkle tree does not cover Signing
+     * Block specificed in {@code signatureInfo}.  The output is suitable to be used as the on-disk
+     * format for fs-verity to use (with elide and patch extensions).
+     *
+     * @return ApkVerityResult containing a buffer with the generated Merkle tree stored at the
+     *         front, the tree size, and the calculated root hash.
+     */
+    @NonNull
+    public static ApkVerityResult generateApkVerityTree(@NonNull RandomAccessFile apk,
             @Nullable SignatureInfo signatureInfo, @NonNull ByteBufferFactory bufferFactory)
             throws IOException, SecurityException, NoSuchAlgorithmException, DigestException {
-        long signingBlockSize =
-                signatureInfo.centralDirOffset - signatureInfo.apkSigningBlockOffset;
-        long dataSize = apk.length() - signingBlockSize;
+        return generateVerityTree(apk, bufferFactory, signatureInfo, true /* skipSigningBlock */);
+    }
+
+    @NonNull
+    private static ApkVerityResult generateVerityTree(@NonNull RandomAccessFile apk,
+            @NonNull ByteBufferFactory bufferFactory, @Nullable SignatureInfo signatureInfo,
+            boolean skipSigningBlock)
+            throws IOException, SecurityException, NoSuchAlgorithmException, DigestException {
+        long dataSize = apk.length();
+        if (skipSigningBlock) {
+            long signingBlockSize =
+                    signatureInfo.centralDirOffset - signatureInfo.apkSigningBlockOffset;
+            dataSize -= signingBlockSize;
+        }
         int[] levelOffset = calculateVerityLevelOffset(dataSize);
         int merkleTreeSize = levelOffset[levelOffset.length - 1];
 
@@ -85,10 +119,11 @@ abstract class ApkVerityBuilder {
                 merkleTreeSize
                 + CHUNK_SIZE_BYTES);  // maximum size of apk-verity metadata
         output.order(ByteOrder.LITTLE_ENDIAN);
-
         ByteBuffer tree = slice(output, 0, merkleTreeSize);
-        byte[] apkRootHash = generateApkVerityTreeInternal(apk, signatureInfo, DEFAULT_SALT,
-                levelOffset, tree);
+        // Only use default salt in legacy case.
+        byte[] salt = skipSigningBlock ? DEFAULT_SALT : null;
+        byte[] apkRootHash = generateVerityTreeInternal(apk, signatureInfo, salt, levelOffset,
+                tree, skipSigningBlock);
         return new ApkVerityResult(output, merkleTreeSize, apkRootHash);
     }
 
@@ -138,7 +173,8 @@ abstract class ApkVerityBuilder {
             throws IOException, SignatureNotFoundException, SecurityException, DigestException,
                    NoSuchAlgorithmException {
         try (RandomAccessFile apk = new RandomAccessFile(apkPath, "r")) {
-            ApkVerityResult result = generateApkVerityTree(apk, signatureInfo, bufferFactory);
+            ApkVerityResult result = generateVerityTree(apk, bufferFactory, signatureInfo,
+                    true /* skipSigningBlock */);
             ByteBuffer footer = slice(result.verityData, result.merkleTreeSize,
                     result.verityData.limit());
             generateApkVerityFooter(apk, signatureInfo, footer);
@@ -170,11 +206,14 @@ abstract class ApkVerityBuilder {
         private final byte[] mDigestBuffer = new byte[DIGEST_SIZE_BYTES];
         private final byte[] mSalt;
 
-        private BufferedDigester(byte[] salt, ByteBuffer output) throws NoSuchAlgorithmException {
+        private BufferedDigester(@Nullable byte[] salt, @NonNull ByteBuffer output)
+                throws NoSuchAlgorithmException {
             mSalt = salt;
             mOutput = output.slice();
             mMd = MessageDigest.getInstance(JCA_DIGEST_ALGORITHM);
-            mMd.update(mSalt);
+            if (mSalt != null) {
+                mMd.update(mSalt);
+            }
             mBytesDigestedSinceReset = 0;
         }
 
@@ -201,7 +240,9 @@ abstract class ApkVerityBuilder {
                     mMd.digest(mDigestBuffer, 0, mDigestBuffer.length);
                     mOutput.put(mDigestBuffer);
                     // After digest, MessageDigest resets automatically, so no need to reset again.
-                    mMd.update(mSalt);
+                    if (mSalt != null) {
+                        mMd.update(mSalt);
+                    }
                     mBytesDigestedSinceReset = 0;
                 }
             }
@@ -242,6 +283,26 @@ abstract class ApkVerityBuilder {
     // thus the syscall overhead is not too big.
     private static final int MMAP_REGION_SIZE_BYTES = 1024 * 1024;
 
+    private static void generateFsVerityDigestAtLeafLevel(RandomAccessFile file, ByteBuffer output)
+            throws IOException, NoSuchAlgorithmException, DigestException {
+        BufferedDigester digester = new BufferedDigester(null /* salt */, output);
+
+        // 1. Digest the whole file by chunks.
+        consumeByChunk(digester,
+                new MemoryMappedFileDataSource(file.getFD(), 0, file.length()),
+                MMAP_REGION_SIZE_BYTES);
+
+        // 2. Pad 0s up to the nearest 4096-byte block before hashing.
+        int lastIncompleteChunkSize = (int) (file.length() % CHUNK_SIZE_BYTES);
+        if (lastIncompleteChunkSize != 0) {
+            digester.consume(ByteBuffer.allocate(CHUNK_SIZE_BYTES - lastIncompleteChunkSize));
+        }
+        digester.assertEmptyBuffer();
+
+        // 3. Fill up the rest of buffer with 0s.
+        digester.fillUpLastOutputChunk();
+    }
+
     private static void generateApkVerityDigestAtLeafLevel(RandomAccessFile apk,
             SignatureInfo signatureInfo, byte[] salt, ByteBuffer output)
             throws IOException, NoSuchAlgorithmException, DigestException {
@@ -288,15 +349,19 @@ abstract class ApkVerityBuilder {
     }
 
     @NonNull
-    private static byte[] generateApkVerityTreeInternal(@NonNull RandomAccessFile apk,
-            @Nullable SignatureInfo signatureInfo, @NonNull byte[] salt,
-            @NonNull int[] levelOffset, @NonNull ByteBuffer output)
+    private static byte[] generateVerityTreeInternal(@NonNull RandomAccessFile apk,
+            @Nullable SignatureInfo signatureInfo, @Nullable byte[] salt,
+            @NonNull int[] levelOffset, @NonNull ByteBuffer output, boolean skipSigningBlock)
             throws IOException, NoSuchAlgorithmException, DigestException {
-        assertSigningBlockAlignedAndHasFullPages(signatureInfo);
-
         // 1. Digest the apk to generate the leaf level hashes.
-        generateApkVerityDigestAtLeafLevel(apk, signatureInfo, salt, slice(output,
-                    levelOffset[levelOffset.length - 2], levelOffset[levelOffset.length - 1]));
+        if (skipSigningBlock) {
+            assertSigningBlockAlignedAndHasFullPages(signatureInfo);
+            generateApkVerityDigestAtLeafLevel(apk, signatureInfo, salt, slice(output,
+                        levelOffset[levelOffset.length - 2], levelOffset[levelOffset.length - 1]));
+        } else {
+            generateFsVerityDigestAtLeafLevel(apk, slice(output,
+                        levelOffset[levelOffset.length - 2], levelOffset[levelOffset.length - 1]));
+        }
 
         // 2. Digest the lower level hashes bottom up.
         for (int level = levelOffset.length - 3; level >= 0; level--) {
@@ -434,10 +499,11 @@ abstract class ApkVerityBuilder {
         return levelOffset;
     }
 
-    private static void assertSigningBlockAlignedAndHasFullPages(SignatureInfo signatureInfo) {
+    private static void assertSigningBlockAlignedAndHasFullPages(
+            @NonNull SignatureInfo signatureInfo) {
         if (signatureInfo.apkSigningBlockOffset % CHUNK_SIZE_BYTES != 0) {
             throw new IllegalArgumentException(
-                    "APK Signing Block does not start at the page  boundary: "
+                    "APK Signing Block does not start at the page boundary: "
                     + signatureInfo.apkSigningBlockOffset);
         }