signflinger/src/com/android/signflinger/SigningBlockUtils.java

/*
 * Copyright (C) 2020 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.signflinger;

import static java.nio.ByteOrder.LITTLE_ENDIAN;

import com.android.apksig.apk.ApkSigningBlockNotFoundException;
import com.android.apksig.apk.ApkUtils;
import com.android.apksig.apk.ApkUtils.ApkSigningBlock;
import com.android.apksig.internal.apk.ApkSigningBlockUtils;
import com.android.apksig.internal.util.Pair;
import com.android.apksig.util.DataSource;
import com.android.apksig.util.DataSources;
import com.android.apksig.zip.ZipFormatException;
import java.io.File;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/** Generates and appends a new block to APK v2 Signature block. */
public final class SigningBlockUtils {

    private static final int MAGIC_NUM_BYTES = 16;
    private static final int BLOCK_LENGTH_NUM_BYTES = 8;
    static final int SIZE_OF_BLOCK_NUM_BYTES = 8;
    static final int BLOCK_ID_NUM_BYTES = 4;

    static final int VERITY_PADDING_BLOCK_ID = 0x42726577;

    /**
     * Generates a new block with the given block value and block id, and appends it to the signing
     * block.
     *
     * @param signingBlock Block containing v2 signature and (optionally) padding block or null.
     * @param blockValue byte array containing block value of the new block or null.
     * @param blockId block id of the new block.
     * @return APK v2 block with signatures and the new block. If {@code blockValue} is null the
     *     {@code signingBlock} is returned without any modification. If {@code signingBlock} is
     *     null, a new signature block is created containing the new block and, optionally, padding
     *     block.
     */
    public static byte[] addToSigningBlock(byte[] signingBlock, byte[] blockValue, int blockId)
            throws IOException {
        if (blockValue == null || blockValue.length == 0) {
            return signingBlock;
        }
        if (signingBlock == null || signingBlock.length == 0) {
            return createSigningBlock(blockValue, blockId);
        }
        return appendToSigningBlock(signingBlock, blockValue, blockId);
    }

    /**
     * Adds a new block to the signature block and a padding block, if required.
     *
     * @param signingBlock APK v2 signing block containing : length prefix, signers (can include
     *     padding block), length postfix and APK sig v2 block magic.
     * @param blockValue byte array containing block value of the new block.
     * @param blockId block id of the new block.
     * @return APK v2 signing block containing : length prefix, signers including the new block (may
     *     include padding block as well), length postfix and APK sig v2 block magic.
     */
    private static byte[] appendToSigningBlock(byte[] signingBlock, byte[] blockValue, int blockId)
            throws IOException {
        List<Pair<byte[], Integer>> entries =
                new ArrayList<>(
                        extractAllSigners(DataSources.asDataSource(ByteBuffer.wrap(signingBlock))));
        entries.add(Pair.of(blockValue, blockId));
        return ApkSigningBlockUtils.generateApkSigningBlock(entries);
    }

    /**
     * Generate APK sig v2 block containing a block composed of the provided block value and id, and
     * (optionally) padding block.
     */
    private static byte[] createSigningBlock(byte[] blockValue, int blockId) {
        return ApkSigningBlockUtils.generateApkSigningBlock(
                Arrays.asList(Pair.of(blockValue, blockId)));
    }

    /**
     * Extracts all signing block entries except padding block.
     *
     * @param signingBlock APK v2 signing block containing: length prefix, signers (can include
     *     padding block), length postfix and APK sig v2 block magic.
     * @return list of block entry value and block entry id pairs.
     */
    private static List<Pair<byte[], Integer>> extractAllSigners(DataSource signingBlock)
            throws IOException {
        long wholeBlockSize = signingBlock.size();
        // Take the segment of the existing signing block without the length prefix (8 bytes)
        // at the beginning and the length and magic (24 bytes) at the end, so it is just the
        // sequence
        // of length prefix id value pairs.
        DataSource lengthPrefixedIdValuePairsSource =
                signingBlock.slice(
                        SIZE_OF_BLOCK_NUM_BYTES,
                        wholeBlockSize - 2 * SIZE_OF_BLOCK_NUM_BYTES - MAGIC_NUM_BYTES);
        final int lengthAndIdByteCount = BLOCK_LENGTH_NUM_BYTES + BLOCK_ID_NUM_BYTES;
        ByteBuffer lengthAndId = ByteBuffer.allocate(lengthAndIdByteCount).order(LITTLE_ENDIAN);
        List<Pair<byte[], Integer>> idValuePairs = new ArrayList<>();

        for (int index = 0;
                index <= lengthPrefixedIdValuePairsSource.size() - lengthAndIdByteCount; ) {
            lengthPrefixedIdValuePairsSource.copyTo(index, lengthAndIdByteCount, lengthAndId);
            lengthAndId.flip();
            int blockLength = Math.toIntExact(lengthAndId.getLong());
            int id = lengthAndId.getInt();
            lengthAndId.clear();

            if (id != VERITY_PADDING_BLOCK_ID) {
                int blockValueSize = blockLength - BLOCK_ID_NUM_BYTES;
                ByteBuffer blockValue = ByteBuffer.allocate(blockValueSize);
                lengthPrefixedIdValuePairsSource.copyTo(
                        index + BLOCK_LENGTH_NUM_BYTES + BLOCK_ID_NUM_BYTES,
                        blockValueSize,
                        blockValue);
                idValuePairs.add(Pair.of(blockValue.array(), id));
            }

            index += blockLength + BLOCK_LENGTH_NUM_BYTES;
        }
        return idValuePairs;
    }

    /**
     * Extract a block with the given id from the APK. If there is more than one block with the same
     * ID, the first block will be returned. If there are no block with the give id, {@code null}
     * will be returned.
     *
     * @param apk APK file
     * @param blockId id of the block to be extracted.
     */
    public static ByteBuffer extractBlock(File apk, int blockId)
            throws IOException, ZipFormatException, ApkSigningBlockNotFoundException {
        try (RandomAccessFile file = new RandomAccessFile(apk, "r")) {
            DataSource apkDataSource = DataSources.asDataSource(file);
            ApkSigningBlock signingBlockInfo =
                    ApkUtils.findApkSigningBlock(
                            apkDataSource, ApkUtils.findZipSections(apkDataSource));

            DataSource wholeV2Block = signingBlockInfo.getContents();
            final int lengthAndIdByteCount = BLOCK_LENGTH_NUM_BYTES + BLOCK_ID_NUM_BYTES;
            DataSource signingBlock =
                    wholeV2Block.slice(
                            SIZE_OF_BLOCK_NUM_BYTES,
                            wholeV2Block.size() - SIZE_OF_BLOCK_NUM_BYTES - MAGIC_NUM_BYTES);
            ByteBuffer lengthAndId =
                    ByteBuffer.allocate(lengthAndIdByteCount).order(ByteOrder.LITTLE_ENDIAN);
            for (int index = 0; index <= signingBlock.size() - lengthAndIdByteCount; ) {
                signingBlock.copyTo(index, lengthAndIdByteCount, lengthAndId);
                lengthAndId.flip();
                int blockLength = (int) lengthAndId.getLong();
                int id = lengthAndId.getInt();
                lengthAndId.flip();
                if (id == blockId) {
                    ByteBuffer block = ByteBuffer.allocate(blockLength - BLOCK_ID_NUM_BYTES);
                    signingBlock.copyTo(
                            index + lengthAndIdByteCount, blockLength - BLOCK_ID_NUM_BYTES, block);
                    block.flip();
                    return block;
                }
                index += blockLength + BLOCK_LENGTH_NUM_BYTES;
            }
            return null;
        }
    }

    private SigningBlockUtils() {}
}